Duplicate from deepseek-ai/DeepSeek-V4.1-Flash

Co-authored-by: Xingkai Yu <GeeeekExplorer@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,38 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+inference/examples/images/carrots.jpeg filter=lfs diff=lfs merge=lfs -text
+DeepSeek_V41_Tech_Report.pdf filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

DeepSeek_V41_Tech_Report.pdf

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

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 DeepSeek
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,197 @@
+---
+license: mit
+library_name: transformers
+pipeline_tag: image-text-to-text
+---
+
+# DeepSeek-V4.1-Flash: Pushing the Limits of KV Cache Compression
+
+<!-- markdownlint-disable first-line-h1 -->
+<!-- markdownlint-disable html -->
+<!-- markdownlint-disable no-duplicate-header -->
+
+<div align="center">
+  <img src="https://github.com/deepseek-ai/DeepSeek-V2/blob/main/figures/logo.svg?raw=true" width="60%" alt="DeepSeek-V4.1" />
+</div>
+<hr>
+<div align="center" style="line-height: 1;">
+  <a href="https://www.deepseek.com/" target="_blank" style="margin: 2px;">
+    <img alt="Homepage" src="https://github.com/deepseek-ai/DeepSeek-V2/blob/main/figures/badge.svg?raw=true" style="display: inline-block; vertical-align: middle;"/>
+  </a>
+  <a href="https://chat.deepseek.com/" target="_blank" style="margin: 2px;">
+    <img alt="Chat" src="https://img.shields.io/badge/🤖%20Chat-DeepSeek%20V4.1-536af5?color=536af5&logoColor=white" style="display: inline-block; vertical-align: middle;"/>
+  </a>
+</div>
+<div align="center" style="line-height: 1;">
+  <a href="https://huggingface.co/deepseek-ai" target="_blank" style="margin: 2px;">
+    <img alt="Hugging Face" src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-DeepSeek%20AI-ffc107?color=ffc107&logoColor=white" style="display: inline-block; vertical-align: middle;"/>
+  </a>
+  <a href="https://twitter.com/deepseek_ai" target="_blank" style="margin: 2px;">
+    <img alt="Twitter Follow" src="https://img.shields.io/badge/Twitter-deepseek_ai-white?logo=x&logoColor=white" style="display: inline-block; vertical-align: middle;"/>
+  </a>
+</div>
+<div align="center" style="line-height: 1;">
+  <a href="LICENSE" style="margin: 2px;">
+    <img alt="License" src="https://img.shields.io/badge/License-MIT-f5de53?&color=f5de53" style="display: inline-block; vertical-align: middle;"/>
+  </a>
+</div>
+
+<p align="center">
+  <a href="https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/blob/main/DeepSeek_V41_Tech_Report.pdf"><b>Technical Report</b> 👁️</a>
+</p>
+
+## Introduction
+
+We introduce **DeepSeek-V4.1-Flash**, a multimodal Mixture-of-Experts (MoE) model with 552B backbone parameters and support for contexts of up to one million tokens. The model natively processes images and text, and generates text autoregressively.
+
+**Architecture.** DeepSeek-V4.1-Flash adopts a **Causal Encoder-Decoder (CED)** architecture: a 40-layer Transformer organized as a 20-layer causal encoder followed by a 20-layer decoder. With CED, the decoder's global KV cache is projected from the final encoder hidden states rather than derived from each decoder layer's own hidden states. This allows the model to activate only **8B parameters per token during prefill** and **16B during decode**, substantially improving cost efficiency for input-heavy agentic workloads. **SWA Bounded Replay** reconstructs missing SWA KV states by replaying only the most recent *n*_win tokens, avoiding the need to persist SWA KV to SSD and reducing the persistent KV cache footprint to roughly **1/8** of that of DeepSeek-V4-Flash.
+
+**Compressed Sparse Attention 2 (CSA2).** DeepSeek-V4.1-Flash uses CSA2, which assigns each attention layer one of three static modes — **Full**, **Reindex**, or **Reuse** — to share main KV and indexer K across layers and reuse Top-K sparse-attention indices. In the decoder, a **Hierarchical Sparse Indexer** further restricts later indexing layers to a candidate pool constructed by the first Full Mode layer, bounding deeper indexer cost independently of context length. Combined with **FP4 main KV caching** (E2M1 format, one E4M3 scale per 16 channels), these designs reduce the global KV cache footprint to **890 bytes per token** — roughly **1/4** of DeepSeek-V4-Flash.
+
+**Additional architectural components** include Single-Pass mHC (revised residual-stream mixing with an efficient Mega-mHC kernel), Engram conditional memory (196B parameters, sparsely accessed via token-based lookup), and DSpark speculative decoding (semi-autoregressive draft generation with confidence-scheduled verification). The model uses 1 shared expert and 384 routed experts per MoE layer, activating 6 routed experts per token.
+
+**Multimodal architecture.** A vision encoder (DeepSeek-ViT, trained from scratch with 2D-RoPE and 3×3 pixel-unshuffle downsampling) and a two-layer MLP projector convert images into visual embeddings, processed jointly with text embeddings from the start of language-model pre-training.
+
+**Pre-training.** DeepSeek-V4.1-Flash is trained from scratch on a multimodal corpus comprising **45T tokens**, with sparse attention trained at a sequence length of 64K and context extended to 1M tokens at 34T tokens.
+
+**Post-training.** The post-training recipe follows the standard SFT → RL → on-policy distillation (OPD) paradigm without algorithmic modifications. All substantive changes lie instead in the data pipeline: large-scale automated synthesis of agent tasks and environments with progressive scaling of data, tasks, and rollouts. The model supports a **continuously controllable reasoning effort** setting (integer 1–100) that trades inference cost for accuracy.
+
+<div align="center">
+  <img src="assets/dsv41_agentic_performance.png" width="48%" alt="DeepSeek-V4.1-Flash agentic benchmark performance" style="display: inline-block; margin: 0 1%;" />
+  <img src="assets/dsv41_kv_cache.png" width="48%" alt="Global KV cache size per token across DeepSeek generations" style="display: inline-block; margin: 0 1%;" />
+</div>
+
+*Figure 1. (a) Performance of DeepSeek-V4.1-Flash and counterparts on agentic benchmarks. (b) Global KV cache size per token (bytes) across generations of DeepSeek models. DeepSeek-V4.1-Flash achieves approximately 4-fold and 437-fold reductions relative to DeepSeek-V4-Flash and DeepSeek-V1, respectively.*
+
+## Evaluation Results
+
+### Base Model
+
+All base models are evaluated in our internal framework under the same evaluation settings. Scores within 0.3 of each other are considered equivalent.
+
+<div align="center">
+
+| Benchmark (Metric) | # Shots | DeepSeek-V4-Flash-Base | DeepSeek-V4-Pro-Base | DeepSeek-V4.1-Flash-Base |
+| :--- | :---: | :---: | :---: | :---: |
+| Architecture | — | MoE | MoE | MoE |
+| # Backbone Params | — | 284B | 1.6T | 552B |
+| # Activated Params | — | 13B | 49B | 8B / 16B |
+| **World Knowledge** | | | | |
+| AGIEval (EM) | 3–5-shot | 83.9 | **84.4** | 83.4 |
+| MMLU-Pro (EM) | 5-shot | 68.3 | 73.5 | **74.1** |
+| C-Eval (EM) | 5-shot | 92.1 | **93.1** | 92.1 |
+| MultiLoKo (LLM-Judge) | 5-shot | 42.6 | **50.9** | 45.5 |
+| SimpleQA-Verified (EM) | 25-shot | 30.1 | **55.2** | 42.3 |
+| SuperGPQA (EM) | 5-shot | 46.5 | **53.9** | 53.1 |
+| **Language & Reasoning** | | | | |
+| BBH (EM) | 3-shot | 86.9 | **87.5** | 86.1 |
+| BBEH (EM) | 1-shot | 25.4 | **29.8** | 27.2 |
+| DROP (F1) | 1-shot | **88.6** | **88.7** | 87.9 |
+| HellaSwag (EM) | 0-shot | 85.7 | **88.0** | 87.2 |
+| **Code & Math** | | | | |
+| BigCodeBench (Pass@1) | 3-shot | 56.8 | 59.2 | **60.6** |
+| HumanEval (Pass@1) | 0-shot | 69.5 | 76.8 | **79.4** |
+| GSM8K (EM) | 8-shot | 90.8 | 92.6 | **93.0** |
+| MATH (EM) | 4-shot | 57.4 | **64.5** | 61.1 |
+| MGSM (EM) | 8-shot | **85.7** | 84.4 | 80.2 |
+| **Long Context** | | | | |
+| LongBench-V2 (EM) | 1-shot | 44.7 | **51.5** | 45.2 |
+| **Multimodal** | | | | |
+| MMMU-Pro (EM) | 4-shot | — | — | 56.5 |
+| CVBench (EM) | 4-shot | — | — | 77.9 |
+| DocVQA (LLM-Judge) | 4-shot | — | — | 95.6 |
+| RefCOCO-avg (Acc@0.5) | 0-shot | — | — | 86.0 |
+
+</div>
+
+### Instruct Model
+
+DeepSeek-V4.1-Flash supports a continuously controllable reasoning effort from 1 to 100. All instruct results below use the maximum effort setting (`reasoning_effort=100`). Evaluations use `temperature=1.0, top_p=0.95`.
+
+For code agent benchmarks (Terminal-Bench 2.1/3.0/4.0, DeepSWE v1.1, NL2Repo-Bench, ProgramBench), the model is evaluated with the Minimal mode of DeepSeek Harness and a 1M-token context window. To align with official setup requirements, the mini-SWE harness is used for DeepSWE v1.1, and the Claude Code harness for SEC-Bench Pro. Visual agent benchmarks (Chartography, BabyVision, ZeroBench) use the Claude Code harness with a 512k-token context window. Agent's Last Exam and AutomationBench use their official scaffolds. All agentic evaluations use `temperature=1.0, top_p=0.95`.
+
+#### Comparison with frontier models (Max reasoning effort)
+
+<div align="center">
+
+| Benchmark (Metric) | Opus-5.0 | GPT-5.6 Sol | K3 | GLM-5.3 | DS-V4-Pro | DS-V4-Flash | DS-V4.1-Flash |
+| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
+| **Reasoning** | | | | | | | |
+| GPQA Diamond (Pass@1) | 93.4 | **94.1** | 92.9 | 88.1 | 92.4 | 89.9 | 90.9 |
+| HLE (Pass@1) | **56.3** | 44.5 | 43.5 | 42.0† | 42.7† | 37.8† | 36.8 (39.1†) |
+| Codeforces (Rating) | — | — | — | — | 3348 | 3289 | **3471** |
+| MathArena Apex (Pass@1) | — | — | **65.6** | — | 65.3 | 58.6 | **65.6** |
+| **Agentic** | | | | | | | |
+| Terminal-Bench 2.1 (Pass@1) | 89.1 | 88.8 | 88.3 | 88.2 | 87.9 | 82.7 | **90.6** |
+| Terminal-Bench 3.0 (Pass@1) | **43.3** | 34.4 | 17.7 | 28.3 | 11.8 | 7.6 | 30.0 |
+| Terminal-Bench 4.0 (Pass@1) | **51.8** | 39.9 | 12.6 | 37.9 | 12.4 | 7.0 | 31.2 |
+| DeepSWE v1.1 (Resolved) | 74.0 | 73.0 | 67.5 | 66.9 | 62.7 | 54.4 | **74.2** |
+| ProgramBench (Almost@1) | **37.0** | 23.0 | 17.5 | 19.0 | 15.5 | — | 20.3 |
+| NL2Repo-Bench (Score) | **75.3** | 56.8 | 58.0 | 58.0 | 61.5 | 54.2 | 64.0 |
+| CyberGym (Pass@1) | — | 84.5 | 80.0 | 84.5 | 83.3 | 76.7 | **88.1** |
+| SEC-Bench Pro (Pass@1) | — | **74.3** | — | — | 56.4 | 30.9 | 62.8 |
+| ExploitGym (Pass@1) | 22.1 | **33.7** | — | 15.0 | 5.4 | 1.8 | 15.3 |
+| HLE w/ tools (Pass@1) | 63.6 | — | 59.8 | 62.5 | 60.0 | 51.5 | **63.9** |
+| AutomationBench (Pass@1) | 50.3 | 45.8 | 46.7 | 48.8 | 43.2 | 37.7 | **54.8** |
+| Agent's Last Exam (Pass@1) | 28.6 | 26.7 | 27.6 | 28.5 | 25.7 | 25.2 | **31.8** |
+| Chartography w/ tools (Pass@1) | **84.0** | 79.9 | 68.1 | — | — | — | 78.9 |
+| BabyVision w/ tools (Pass@1) | **94.1** | 88.9 | 85.7 | — | — | — | 89.6 |
+| ZeroBench-main w/ tools (Pass@5) | 52.0 | **53.0** | 41.0 | — | — | — | 49.0 |
+
+</div>
+
+*† Text-only subset of HLE.*
+
+#### Performance across agent scaffolds (DeepSWE v1.1 and Terminal-Bench 2.1, Max reasoning effort)
+
+All scaffolds use N=8 samples per task on DeepSWE v1.1 and N=3 on Terminal-Bench 2.1, with Linux containers, `temperature=1.0`, `top_p=0.95`, a 1M-token context limit, and max_steps=500 per agent. Terminal-Bench 2.1 is evaluated without network access.
+
+<div align="center">
+
+| Benchmark (Metric) | Claude Code | Codex | OpenCode | Pi | mini-SWE | DSH Minimal | DSH Standard | DSH PTC |
+| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
+| DeepSWE v1.1 (Resolved) | 69.8 | 65.6 | 65.5 | 66.2 | 74.2 | 72.6 | 70.5 | 67.6 |
+| Terminal-Bench 2.1 (Pass@1) | 88.0 | 84.1 | 85.0 | 86.1 | 90.3 | 90.6 | 85.8 | 85.8 |
+
+</div>
+
+## Prompt Encoding
+
+This release does not include a Jinja-format chat template. The [`encoding`](encoding/README.md) folder contains a self-contained Python reference implementation (`encoding.py`) with test cases for multi-turn conversations, tool calling, thinking mode, numeric reasoning effort, mid-conversation system messages, and interleaved image content.
+
+For production use, we additionally release [deepseek-recipe](https://github.com/deepseek-ai/deepseek-recipe), a set of Rust libraries with Python bindings that provides the same prompt format as a maintained, protocol-aware toolkit. It converts Messages, Chat Completions, and Responses API requests into the Conversation format, encodes them into DeepSeek V4 and V4.1 prompts or token IDs, and parses model output back into complete or streamed responses — covering thinking, tool calls, images, and generation settings. Model inference, tool execution, and HTTP transport are left to the caller.
+
+## Minimal Inference
+
+Please refer to the [`inference`](inference/README.md) folder for instructions on weight conversion and running inference locally.
+
+**Recommended sampling parameters:**
+
+| Parameter | Value |
+| :--- | :--- |
+| `temperature` | 1.0 |
+| `top_p` | 0.95 or 1.0 |
+| `context_window` | 1M tokens |
+| `max_tokens` | ≥ 256K |
+
+## Reproducing DeepSWE Benchmark Results
+
+The [`evaluation`](evaluation/README.md) folder contains step-by-step instructions for reproducing the DeepSWE v1.1 benchmark results, covering both the `dsh-minimal` agent and the official `mini-swe-agent`. The patch required to integrate `dsh-minimal` with [Pier](https://github.com/datacurve-ai/pier) is also included there.
+
+## License
+
+This repository and the model weights are licensed under the [MIT License](LICENSE).
+
+## Citation
+
+```bibtex
+@misc{deepseekai2026deepseekv41flash,
+      title={DeepSeek-V4.1-Flash: Pushing the Limits of KV Cache Compression},
+      author={DeepSeek-AI},
+      year={2026},
+}
+```
+
+## Contact
+
+If you have any questions, please raise an issue or contact us at [service@deepseek.com](mailto:service@deepseek.com).

config.json

@@ -0,0 +1,170 @@
+{
+  "architectures": [
+    "DeepseekV41ForCausalLM"
+  ],
+  "model_type": "deepseek_v41",
+  "dtype": "bfloat16",
+  "transformers_version": "5.6.0",
+  "bos_token_id": 0,
+  "eos_token_id": 1,
+  "pad_token_id": 2,
+  "image_token_id": 129264,
+  "quantization_config": {
+    "quant_method": "fp8",
+    "activation_scheme": "dynamic",
+    "weight_block_size": [
+      32,
+      32
+    ],
+    "scale_fmt": "ue8m0",
+    "expert_dtype": "fp4"
+  },
+  "text_config": {
+    "model_type": "deepseek_v41_text",
+    "vocab_size": 129280,
+    "hidden_size": 5120,
+    "moe_intermediate_size": 2304,
+    "num_hidden_layers": 40,
+    "num_attention_heads": 64,
+    "num_key_value_heads": 1,
+    "head_dim": 512,
+    "qk_rope_head_dim": 64,
+    "q_lora_rank": 1280,
+    "o_lora_rank": 1024,
+    "o_groups": 8,
+    "hidden_act": "silu",
+    "swiglu_limit": 10.0,
+    "rms_norm_eps": 1e-20,
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "initializer_range": 0.02,
+    "use_cache": true,
+    "tie_word_embeddings": false,
+    "max_position_embeddings": 1048576,
+    "rope_theta": 10000,
+    "rope_scaling": {
+      "rope_type": "yarn",
+      "factor": 16,
+      "beta_fast": 32,
+      "beta_slow": 1,
+      "original_max_position_embeddings": 65536
+    },
+    "n_routed_experts": 384,
+    "n_shared_experts": 1,
+    "num_experts_per_tok": 6,
+    "scoring_func": "sqrtsoftplus",
+    "topk_method": "noaux_tc",
+    "norm_topk_prob": true,
+    "routed_scaling_factor": 1.5,
+    "sliding_window": 128,
+    "compress_ratios": [
+      0,
+      0,
+      2,
+      2,
+      2,
+      2,
+      2,
+      2,
+      2,
+      2,
+      2,
+      2,
+      2,
+      2,
+      2,
+      2,
+      2,
+      2,
+      2,
+      2,
+      1,
+      1,
+      1,
+      1,
+      1,
+      1,
+      1,
+      1,
+      1,
+      1,
+      1,
+      1,
+      1,
+      1,
+      1,
+      1,
+      1,
+      1,
+      1,
+      1,
+      0,
+      0,
+      0
+    ],
+    "compress_rope_theta": 160000,
+    "kv_source_layer_ids": [
+      2,
+      8,
+      14,
+      20
+    ],
+    "index_source_layer_ids": [
+      2,
+      8,
+      14,
+      20,
+      24,
+      28,
+      32,
+      36
+    ],
+    "index_n_heads": 32,
+    "index_head_dim": 128,
+    "index_topk": 512,
+    "candidate_source_layer_id": 20,
+    "candidate_topk_blocks": 2048,
+    "candidate_block_size": 8,
+    "hc_mult": 4,
+    "hc_sinkhorn_iters": 20,
+    "hc_eps": 1e-06,
+    "engram_layer_ids": [
+      1,
+      14
+    ],
+    "engram_num_embeddings": [
+      384006168,
+      384016682
+    ],
+    "engram_max_ngram_size": 4,
+    "engram_vocab_size": 16000000,
+    "engram_n_heads": 8,
+    "engram_head_dim": 256,
+    "engram_pad_token_id": 2,
+    "engram_compressed_vocab_size": 99092,
+    "num_nextn_predict_layers": 3,
+    "dspark_block_size": 5,
+    "dspark_noise_token_id": 128799,
+    "dspark_target_layer_ids": [
+      37,
+      38,
+      39
+    ],
+    "dspark_markov_rank": 256,
+    "dspark_n_routed_experts": 128,
+    "dspark_num_experts_per_tok": 3
+  },
+  "vision_config": {
+    "model_type": "deepseek_v41_vision",
+    "num_hidden_layers": 32,
+    "hidden_size": 1024,
+    "num_attention_heads": 16,
+    "intermediate_size": 2816,
+    "patch_size": 14,
+    "rope_theta": 10000,
+    "downsample_ratio": 3,
+    "max_image_tokens": 1024,
+    "min_pixels": 295936,
+    "max_wh_ratio": null
+  }
+}

encoding/README.md

@@ -0,0 +1,305 @@
+# DeepSeek-V4.1 text and vision encoding
+
+`encoding.py` is the standalone prompt-format reference for DeepSeek-V4.1. It
+supports multi-turn conversations, tool calls, thinking modes, numeric reasoning
+effort, mid-conversation system messages, and interleaved image content blocks,
+without importing the inference implementation.
+
+## V4.1 changes relative to V4
+
+Three prompt-format changes distinguish V4.1 from V4:
+
+1. **DSML tag names use a leading space.** Tool calls are wrapped in
+   `<｜DSML｜ calls>` blocks with `<｜DSML｜ invoke>` / `<｜DSML｜ parameter>` tags
+   (note the space before `calls`, `invoke`, and `parameter`). The V4 format used
+   `<｜DSML｜tool_calls>` without a space.
+
+2. **Reasoning effort is a numeric budget (1–100).** The effort prefix is
+   rendered as `Reasoning Effort: {budget} (range 1-100, ...)` rather than the
+   verbose natural-language descriptions used in V4. String aliases map as
+   follows: `"low"` → 50, `"high"` → 75, `"max"` → 100. The
+   default is `"high"` (75). The effort prefix is only rendered in
+   `thinking_mode="thinking"` and only at the beginning of the conversation
+   (index 0).
+
+3. **Mid-conversation system messages** are supported via the `<｜System｜>` token.
+   A mid-conversation system message behaves like a user message for the purpose
+   of appending the assistant generation header.
+
+## Quick start
+
+```python
+from encoding import encode_messages, parse_message_from_completion_text
+
+# Text-only conversation
+messages = [
+    {"role": "system", "content": "You are a helpful assistant."},
+    {"role": "user", "content": "What is 2+2?"},
+]
+prompt, media = encode_messages(
+    messages,
+    thinking_mode="thinking",
+    reasoning_effort=75,           # integer 1–100, or "low"/"high"/"max"
+    return_multi_modal_data=True,
+)
+# prompt:
+# '<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the
+#  value, the more thorough the reasoning)\n\nYou are a helpful assistant.
+#  <｜User｜>What is 2+2?<｜Assistant｜><think>'
+
+# Parse model output back to a 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:** `parse_message_from_completion_text` 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.
+
+## OpenAI-style messages
+
+```python
+from encoding import encode_messages
+
+messages = [{
+    "role": "user",
+    "content": [
+        {"type": "text", "text": "第一张图"},
+        {
+            "type": "image_url",
+            "image_url": {"url": "examples/images/image_1.jpeg"},
+        },
+        {"type": "text", "text": "有什么内容？"},
+    ],
+}]
+
+prompt, media = encode_messages(
+    messages,
+    thinking_mode="chat",
+    return_multi_modal_data=True,
+)
+# prompt:
+# '<｜begin▁of▁sentence｜><｜User｜>第一张图\n\n<｜deepseek_image｜>\n\n有什么内容？<｜Assistant｜></think>'
+# media["images"] contains the image records in prompt order
+```
+
+Images are represented in the prompt by `<｜deepseek_image｜>`. `media["images"]`
+contains the corresponding image records in exactly the same order they appear in
+the prompt. Pixel loading and expansion into model image tokens are handled by
+`inference/image_processor.py`.
+
+## Compact TXT notation
+
+`parse_tagged_text()` converts a compact prompt such as
+
+```text
+第一张图<image>examples/images/image_1.jpeg</image>有什么内容？
+```
+
+into the same standard content blocks. It is an input convenience layer, not a
+second encoding implementation.
+
+## 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 |
+| `<｜System｜>` | Mid-conversation system message prefix |
+| `<｜latest_reminder｜>` | Latest reminder (date, locale, etc.) |
+| `<think>` / `</think>` | Reasoning block delimiters |
+| `｜DSML｜` | DSML markup token |
+| `<｜deepseek_image｜>` | Image placeholder in the prompt string |
+
+### Roles
+
+The encoding supports the following message roles: `system`, `user`, `assistant`,
+`tool`, and `latest_reminder`.
+
+A `tool` message is not rendered directly: `merge_tool_messages()` converts it
+into a `<tool_result>` block inside the preceding user message. When multiple
+tool results are present, they are sorted by the order of the corresponding
+`tool_calls` in the preceding assistant message.
+
+### 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.
+
+### Thinking mode
+
+In **thinking mode** (`thinking_mode="thinking"`), the model produces explicit
+reasoning inside `<think>...</think>` blocks before responding.
+
+```
+<｜begin▁of▁sentence｜><｜System｜>{reasoning_effort_prefix}{system_prompt}
+<｜User｜>{message}<｜Assistant｜><think>{reasoning}</think>{response}<｜end▁of▁sentence｜>
+```
+
+The reasoning effort prefix is injected once, before the system message, as a
+`<｜System｜>` block:
+
+```
+<｜System｜>Reasoning Effort: {budget} (range 1-100, the higher the value, the more thorough the reasoning)
+```
+
+The `drop_thinking` parameter (default `True`) controls whether reasoning from
+earlier turns is preserved:
+
+- **Without tools**: reasoning content from assistant turns **before** the last
+  user message is stripped. Only the final assistant turn retains its
+  `<think>...</think>` block.
+- **With tools**: `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` message via the `tools` field
+(OpenAI-compatible format). When tools are present, the following schema block is
+injected into the system 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｜ calls>" block like the following:
+
+<｜DSML｜ 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｜ 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｜ 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｜ 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>...
+```
+
+### Tool namespaces
+
+Tool definitions may include a `namespace` alongside `function`, either as a
+string or as an object with `name` and an optional `description`:
+
+```python
+tool = {
+    "type": "function",
+    "namespace": {"name": "search", "description": "Search tools."},
+    "function": {
+        "name": "lookup",
+        "description": "Look up a value",
+        "parameters": {"type": "object", "properties": {"query": {"type": "string"}}},
+    },
+}
+tool_call = {
+    "type": "function",
+    "namespace": "search",
+    "function": {"name": "lookup", "arguments": '{"query": "value"}'},
+}
+```
+
+The tool schema and DSML invocation both use `search::lookup`. The namespace
+description is prepended to the tool description, separated by a newline.
+The parser returns `function.name="lookup"` and `namespace="search"` on the
+tool call, so its output can be passed back to `encode_messages()` directly.
+
+Input also accepts `namespace` inside `function`, or a qualified function name
+such as `search::lookup`. A qualified name must agree with any explicit
+namespace; `::` separates exactly one namespace from the tool name. Tools
+without a namespace retain their original names and output format.
+
+### Reasoning effort
+
+Pass `reasoning_effort` as an integer in `[1, 100]` or as one of `"low"` (50),
+`"high"` (75), or `"max"` (100). The default is `"high"` (75).
+The setting only affects `thinking_mode="thinking"` and is only rendered at the
+start of the conversation (index 0). Intermediate values may be used to elicit
+interpolated reasoning behavior.
+
+### 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｜>` |
+| `<｜read_url｜>` | Determines whether each URL in the user prompt should be fetched and read. | `...<｜User｜>{prompt}<｜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.
+
+## Tests
+
+From this directory:
+
+```bash
+python -m pytest -q test_encoding.py
+```
+
+Test cases are stored as paired JSON input / TXT expected-output files under
+`tests/`. The tests cover multi-turn conversations, tool calling, thinking mode,
+numeric reasoning effort, mid-conversation system messages, and multimodal image
+ordering. They include a check that the TXT and JSON examples encode to the same
+prompt and preserve the same image ordering.

encoding/encoding.py

@@ -0,0 +1,979 @@
+"""
+DeepSeek-V4.1 Text and Vision Encoding
+
+A fully self-contained implementation for encoding/decoding DeepSeek-V4.1 chat
+messages with tool calling, thinking mode, quick instruction tasks, and image
+content blocks. No dependency on encoding_dsv4.
+
+V4.1 changes relative to V4:
+
+1. DSML tag names: tool calls are wrapped in "<｜DSML｜ calls>" blocks with
+   "<｜DSML｜ invoke>" / "<｜DSML｜ parameter>" tags (leading-space tag names).
+2. Numeric reasoning effort: "Reasoning Effort: {budget} (range 1-100, ...)".
+   Accepts an int in [1, 100] or one of "low"/"high"/"max"
+   (mapped to 50/75/100). Defaults to "high". Only rendered in thinking mode.
+3. Mid-conversation system messages are supported via the "<｜System｜>" token.
+   A mid-conversation system message behaves like a user message for the purpose
+   of appending the assistant generation header.
+"""
+
+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｜>"
+
+IMAGE_PLACEHOLDER = "<｜deepseek_image｜>"
+IMAGE_TAG_PATTERN = re.compile(r"<image>(.*?)</image>", re.DOTALL)
+
+# 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_output_template: str = (
+    "<tool_result>{content}</tool_result>"
+)
+
+# ============================================================
+# 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 with namespace-qualified names."""
+    functions = []
+    for tool in tools:
+        function = dict(tool["function"])
+        if tool.get("namespace") is not None:
+            function["namespace"] = tool["namespace"]
+        function["name"] = _tool_name_for_encoding(function)
+        namespace = function.pop("namespace", None)
+        if isinstance(namespace, dict) and namespace.get("description"):
+            function["description"] = (
+                namespace["description"] + "\n" + (function.get("description") or "")
+            )
+        functions.append(function)
+    return functions
+
+
+def _split_tool_name(name: str, namespace: Optional[str] = None) -> Tuple[Optional[str], str]:
+    """Split a qualified name and validate any explicit namespace."""
+    prefix, separator, bare_name = name.partition("::")
+    if separator:
+        assert namespace in (None, prefix), (
+            f"Conflicting tool namespaces: {namespace} != {prefix}"
+        )
+        namespace, name = prefix, bare_name
+    assert "::" not in name, f"Tool name must not contain '::': {name}"
+    assert namespace is None or "::" not in namespace, (
+        f"Tool namespace must not contain '::': {namespace}"
+    )
+    return namespace, name
+
+
+def _tool_name_for_encoding(tool: Dict[str, Any]) -> str:
+    namespace = tool.get("namespace")
+    if isinstance(namespace, dict):
+        namespace = namespace["name"]
+    namespace, name = _split_tool_name(tool["name"], namespace)
+    return name if namespace is None else f"{namespace}::{name}"
+
+
+def tool_calls_from_openai_format(tool_calls):
+    """Convert OpenAI-format tool calls to internal format."""
+    calls = []
+    for tool_call in tool_calls:
+        function = tool_call["function"]
+        namespace, name = _split_tool_name(
+            function["name"], tool_call.get("namespace") or function.get("namespace")
+        )
+        call = {"name": name, "arguments": function["arguments"]}
+        if namespace is not None:
+            call["namespace"] = namespace
+        calls.append(call)
+    return calls
+
+
+def tool_calls_to_openai_format(tool_calls):
+    """Convert internal tool calls to OpenAI format."""
+    calls = []
+    for tool_call in tool_calls:
+        call = {
+            "type": "function",
+            "function": {
+                "name": tool_call["name"],
+                "arguments": tool_call["arguments"],
+            }
+        }
+        if tool_call.get("namespace") is not None:
+            call["namespace"] = tool_call["namespace"]
+        calls.append(call)
+    return calls
+
+
+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", "arguments" (JSON string), and optional "namespace".
+    """
+    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()]) + "}"
+    namespace, name = _split_tool_name(tool_name)
+    tool_call = dict(name=name, arguments=tool_args_json)
+    if namespace is not None:
+        tool_call["namespace"] = namespace
+    return tool_call
+
+
+# ============================================================
+# 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.1 does not have a standalone "tool" role; instead, tool results
+    are encoded as <tool_result> blocks within 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":
+            content_blocks = msg.get("content_blocks")
+            if content_blocks is None:
+                content_blocks = [{"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"].extend(content_blocks)
+            else:
+                # Preserve structured content and all message-level metadata.
+                new_msg = msg
+                new_msg["content_blocks"] = content_blocks
+                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.
+    """
+    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
+
+
+# ============================================================
+# Vision Message Preprocessing
+# ============================================================
+
+def parse_tagged_text(text: str) -> Union[str, List[Dict[str, Any]]]:
+    """Convert ``<image>path</image>`` text into standard content blocks."""
+    matches = list(IMAGE_TAG_PATTERN.finditer(text))
+    remaining = IMAGE_TAG_PATTERN.sub("", text)
+    if "<image>" in remaining or "</image>" in remaining:
+        raise ValueError("Malformed <image>path</image> tag")
+    if not matches:
+        return text
+
+    blocks: List[Dict[str, Any]] = []
+    cursor = 0
+    for match in matches:
+        if match.start() > cursor:
+            blocks.append({"type": "text", "text": text[cursor:match.start()]})
+        path = match.group(1)
+        if not path:
+            raise ValueError("Image path must not be empty")
+        blocks.append({
+            "type": "image_url",
+            "image_url": {"url": path},
+        })
+        cursor = match.end()
+    if cursor < len(text):
+        blocks.append({"type": "text", "text": text[cursor:]})
+    return blocks
+
+
+def _is_image_block(block: Dict[str, Any]) -> bool:
+    """Return whether a content block is an OpenAI/Anthropic/internal image."""
+    return isinstance(block, dict) and block.get("type") in ("image", "image_url")
+
+
+def _extract_image(block: Dict[str, Any]) -> Dict[str, Any]:
+    """Normalize a supported image block into an internal image record."""
+    record: Dict[str, Any] = {"type": "image"}
+    if block.get("type") == "image_url":
+        image_url = block.get("image_url")
+        if isinstance(image_url, str):
+            record["url"] = image_url
+        else:
+            record["url"] = (image_url or {}).get("url", "")
+    else:
+        for key in ("source", "url", "data"):
+            if key in block:
+                record[key] = block[key]
+    if not any(record.get(key) for key in ("source", "url", "data")):
+        raise ValueError("Image block does not contain a valid source")
+    return record
+
+
+def _process_image_blocks(
+    blocks: List[Any], image_placeholder: str = IMAGE_PLACEHOLDER
+) -> Tuple[List[Any], List[Dict[str, Any]]]:
+    """Replace image blocks and collect their records in one ordered traversal."""
+    new_blocks: List[Any] = []
+    images: List[Dict[str, Any]] = []
+    for block in blocks:
+        if not isinstance(block, dict):
+            new_blocks.append(block)
+            continue
+        if _is_image_block(block):
+            new_blocks.append({"type": "text", "text": image_placeholder})
+            images.append(_extract_image(block))
+        elif block.get("type") == "tool_result" and isinstance(block.get("content"), list):
+            block = copy.copy(block)
+            block["content"], nested_images = _process_image_blocks(
+                block["content"], image_placeholder)
+            new_blocks.append(block)
+            images.extend(nested_images)
+        elif block.get("type") == "text":
+            text = block.get("text") or ""
+            if IMAGE_PLACEHOLDER in text:
+                raise ValueError(
+                    f"Text block contains image placeholder '{IMAGE_PLACEHOLDER}': "
+                    f"'{text[:100]}'. Images should be separate content blocks."
+                )
+            new_blocks.append(block)
+        else:
+            new_blocks.append(block)
+    return new_blocks, images
+
+
+def _validate_no_image_sp_tokens(msg: Dict[str, Any]) -> None:
+    """Reject user-supplied image placeholder tokens in textual fields."""
+    content = msg.get("content")
+    if isinstance(content, str) and IMAGE_PLACEHOLDER in content:
+        raise ValueError(
+            f"Message content contains image special token '{IMAGE_PLACEHOLDER}'. "
+            "Images should be provided as image content blocks."
+        )
+    reasoning_content = msg.get("reasoning_content")
+    if isinstance(reasoning_content, str) and IMAGE_PLACEHOLDER in reasoning_content:
+        raise ValueError(
+            f"reasoning_content contains image special token '{IMAGE_PLACEHOLDER}'"
+        )
+
+
+def process_image_messages(
+    messages: List[Dict[str, Any]],
+) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
+    """Normalize image blocks and return their records in prompt order."""
+    processed: List[Dict[str, Any]] = []
+    images: List[Dict[str, Any]] = []
+    for msg in messages:
+        msg = copy.deepcopy(msg)
+        _validate_no_image_sp_tokens(msg)
+
+        if isinstance(msg.get("content"), list) and "content_blocks" not in msg:
+            msg["content_blocks"] = msg.pop("content")
+
+        if msg.get("content_blocks"):
+            msg["content_blocks"], message_images = _process_image_blocks(
+                msg["content_blocks"])
+            images.extend(message_images)
+            if not isinstance(msg.get("content"), str):
+                texts = [
+                    block.get("text", "")
+                    for block in msg["content_blocks"]
+                    if isinstance(block, dict) and block.get("type") == "text"
+                ]
+                msg["content"] = "\n\n".join(texts)
+
+        processed.append(msg)
+    return processed, images
+
+
+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
+
+# ============================================================
+# V4.1 Special Tokens and DSML Tag Names
+# ============================================================
+
+SYSTEM_SP_TOKEN = "<｜System｜>"
+
+tool_calls_block_name: str = " calls"
+tool_call_tag_name: str = " invoke"
+tool_parameter_tag_name: str = " parameter"
+
+tool_call_template: str = (
+    "<{dsml_token}{tool_call_tag_name} name=\"{name}\">\n{arguments}\n</{dsml_token}{tool_call_tag_name}>"
+)
+tool_calls_template = (
+    "<{dsml_token}{tc_block_name}>\n{tool_calls}\n</{dsml_token}{tc_block_name}>"
+)
+
+# ============================================================
+# Reasoning Effort (numeric budget)
+# ============================================================
+
+REASONING_EFFORT_TEMPLATE = (
+    "Reasoning Effort: {budget} "
+    "(range 1-100, the higher the value, the more thorough the reasoning)\n\n"
+)
+
+REASONING_EFFORT_MAPPINGS: Dict[str, int] = {
+    "low": 50,
+    "high": 75,
+    "max": 100,
+}
+DEFAULT_REASONING_EFFORT = "high"
+
+
+def render_reasoning_effort(
+    index: int,
+    thinking_mode: str,
+    effort: Union[str, int, None],
+) -> str:
+    """Render the V4.1 numeric reasoning effort prefix (thinking mode, index 0 only)."""
+    if effort is None:
+        effort = DEFAULT_REASONING_EFFORT
+    assert (
+        type(effort) is int and 1 <= effort <= 100
+    ) or effort in REASONING_EFFORT_MAPPINGS, (
+        "Invalid reasoning effort for deepseek_v41: "
+        f"{effort}, should be int within [1,100] or {list(REASONING_EFFORT_MAPPINGS)}"
+    )
+    if type(effort) is str:
+        effort = REASONING_EFFORT_MAPPINGS[effort]
+    if index == 0 and thinking_mode == "thinking":
+        return REASONING_EFFORT_TEMPLATE.format(budget=effort)
+    return ""
+
+
+# ============================================================
+# Tools rendering
+# ============================================================
+
+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}{tc_block_name}>" block like the following:
+
+<{dsml_token}{tc_block_name}>
+<{dsml_token}{tool_call_tag_name} name="$TOOL_NAME">
+<{dsml_token}{tool_parameter_tag_name} name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</{dsml_token}{tool_parameter_tag_name}>
+...
+</{dsml_token}{tool_call_tag_name}>
+<{dsml_token}{tool_call_tag_name} name="$TOOL_NAME2">
+...
+</{dsml_token}{tool_call_tag_name}>
+</{dsml_token}{tc_block_name}>
+
+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.
+"""
+
+
+def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str:
+    """Render tool schemas into the V4.1 system prompt format."""
+    tools_json = [to_json(t) for t in tools]
+
+    return TOOLS_TEMPLATE.format(
+        tool_schemas="\n".join(tools_json),
+        dsml_token=dsml_token,
+        tc_block_name=tool_calls_block_name,
+        tool_call_tag_name=tool_call_tag_name,
+        tool_parameter_tag_name=tool_parameter_tag_name,
+        thinking_start_token=thinking_start_token,
+        thinking_end_token=thinking_end_token,
+    )
+
+
+def encode_arguments_to_dsml(tool_call: Dict[str, Any]) -> str:
+    """Encode tool call arguments into V4.1 DSML parameter format."""
+    p_dsml_template = (
+        '<{dsml_token}{tool_parameter_tag_name} name="{key}" string="{is_str}">'
+        '{value}</{dsml_token}{tool_parameter_tag_name}>'
+    )
+    P_dsml_strs = []
+
+    arguments = tool_call["arguments"]
+    if not isinstance(arguments, dict):
+        # Tolerate JSON strings, including double-encoded ones.
+        for _ in range(2):
+            if isinstance(arguments, str):
+                try:
+                    arguments = json.loads(arguments)
+                except Exception:
+                    break
+            else:
+                break
+        if not isinstance(arguments, dict):
+            arguments = {"arguments": tool_call["arguments"]}
+
+    for k, v in arguments.items():
+        P_dsml_strs.append(p_dsml_template.format(
+            dsml_token=dsml_token,
+            tool_parameter_tag_name=tool_parameter_tag_name,
+            key=k,
+            is_str="true" if isinstance(v, str) else "false",
+            value=v if isinstance(v, str) else to_json(v),
+        ))
+
+    return "\n".join(P_dsml_strs)
+
+
+# ============================================================
+# Message Rendering
+# ============================================================
+
+def find_last_user_index(messages: List[Dict[str, Any]]) -> int:
+    """
+    Find the index of the last user message.
+
+    V4.1 supports mid-conversation system messages, which count as user
+    messages for the purposes of the assistant generation header.
+    """
+    last_user_index = -1
+    for idx in range(len(messages) - 1, -1, -1):
+        role = messages[idx].get("role")
+        if role == "user" or (role == "system" and idx > 0):
+            last_user_index = idx
+            break
+    return last_user_index
+
+
+def render_message(
+    index: int,
+    messages: List[Dict[str, Any]],
+    thinking_mode: str,
+    drop_thinking: bool = True,
+    reasoning_effort: Union[str, int, None] = None,
+) -> str:
+    """
+    Render a single message at the given index into its V4.1 encoded string form.
+    """
+    assert 0 <= index < len(messages)
+    assert thinking_mode in ["chat", "thinking"], f"Invalid thinking_mode `{thinking_mode}`"
+
+    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 (thinking mode, index 0 only)
+    reasoning_effort_prompt = render_reasoning_effort(index, thinking_mode, reasoning_effort)
+    # System token leads the conversation when there is a reasoning effort prompt
+    # or the first message is a system message.
+    prompt = SYSTEM_SP_TOKEN if index == 0 and (reasoning_effort_prompt or role == "system") else ""
+    prompt += reasoning_effort_prompt
+
+    if role == "system":
+        if index > 0:
+            # Mid-conversation system message
+            prompt += SYSTEM_SP_TOKEN
+        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 == "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_v41 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,
+                    tool_call_tag_name=tool_call_tag_name,
+                    name=_tool_name_for_encoding(tc),
+                    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 role == "user" or (role == "system" and index > 0):
+        # Normal generation: append Assistant + thinking token
+        # (mid-conversation system messages also trigger the assistant header)
+        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
+
+
+# ============================================================
+# Main Encoding Function
+# ============================================================
+
+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.
+    Same as V4, but uses the V4.1 last-user definition (mid systems count).
+    """
+    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)
+
+    return result
+
+
+def _encode_messages_text(
+    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: Union[str, int, None] = None,
+) -> str:
+    """Encode preprocessed (text-only) messages into the V4.1 prompt format."""
+    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)
+        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 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: Union[str, int, None] = None,
+    return_multi_modal_data: bool = False,
+) -> Any:
+    """Encode text or multimodal messages into the DeepSeek-V4.1 prompt format.
+
+    Text-only calls return the prompt string. When return_multi_modal_data is
+    true, the result is ``(prompt, media_data)``.
+    """
+    context = context or []
+    processed_context, _ = process_image_messages(context) if context else ([], [])
+    processed_messages, images = process_image_messages(messages)
+    prompt = _encode_messages_text(
+        processed_messages,
+        thinking_mode=thinking_mode,
+        context=processed_context if processed_context else None,
+        drop_thinking=drop_thinking,
+        add_default_bos_token=add_default_bos_token,
+        reasoning_effort=reasoning_effort,
+    )
+    if return_multi_modal_data:
+        return prompt, {"images": images}
+    return prompt
+
+
+def load_cases(input_file: str) -> List[Dict[str, Any]]:
+    """Load one or more OpenAI-format conversation cases from JSON."""
+    with open(input_file) as file:
+        data = json.load(file)
+    if isinstance(data, dict):
+        data = [data]
+    elif data and isinstance(data[0], dict) and "role" in data[0]:
+        data = [{"messages": data}]
+
+    cases = []
+    for case in data:
+        messages = copy.deepcopy(case["messages"])
+        if "tools" in case:
+            if not messages:
+                raise ValueError("A case with tools must contain at least one message")
+            messages[0]["tools"] = case["tools"]
+        cases.append({
+            "messages": messages,
+            "context": case.get("context"),
+            "thinking_mode": case.get("thinking_mode"),
+            "reasoning_effort": case.get("reasoning_effort"),
+        })
+    return cases
+
+
+def encode_case(
+    case: Dict[str, Any], thinking_mode: str
+) -> Tuple[str, List[Dict[str, Any]]]:
+    """Encode one JSON case and return its current-turn image records."""
+    prompt, media_data = encode_messages(
+        case["messages"],
+        thinking_mode=case.get("thinking_mode") or thinking_mode,
+        context=case.get("context"),
+        reasoning_effort=case.get("reasoning_effort"),
+        return_multi_modal_data=True,
+    )
+    return prompt, media_data["images"]
+
+
+# ============================================================
+# Parsing (Decoding model output)
+# ============================================================
+
+def parse_tool_calls(index: int, text: str) -> Tuple[int, Optional[str], List[Dict[str, str]]]:
+    """
+    Parse V4.1 DSML tool calls from text starting at the given index.
+
+    Returns:
+        Tuple of (new_index, last_stop_token, list_of_tool_call_dicts).
+    """
+    tool_calls: List[Dict[str, Any]] = []
+    stop_token = None
+    tool_calls_end_token = f"</{dsml_token}{tool_calls_block_name}>"
+    tool_call_start_token = f"<{dsml_token}{tool_call_tag_name}"
+    tool_call_end_token = f"</{dsml_token}{tool_call_tag_name}"
+    tool_parameter_start_token = f"<{dsml_token}{tool_parameter_tag_name}"
+    tool_parameter_end_token = f"/{dsml_token}{tool_parameter_tag_name}"
+
+    while index < len(text):
+        index, _, stop_token = _read_until_stop(index, text, [tool_call_start_token, 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, [tool_parameter_start_token, tool_call_end_token])
+
+        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 == tool_parameter_start_token:
+            index, param_content, stop_token = _read_until_stop(index, text, [tool_parameter_end_token])
+
+            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, [tool_parameter_start_token, tool_call_end_token])
+            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 (V4.1 format).
+
+    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.py

@@ -0,0 +1,558 @@
+"""
+Tests for encoding.py (DeepSeek-V4.1 encoding).
+
+Adapted from dsv41-master/deepseek_harmony/tests/test_deepseek_v41.py for the
+self-contained dict-based API in this repo.
+"""
+
+import copy
+import json
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+import encoding as enc
+from encoding import (
+    IMAGE_PLACEHOLDER,
+    SYSTEM_SP_TOKEN,
+    encode_messages,
+    parse_message_from_completion_text,
+    render_message,
+    merge_tool_messages,
+)
+
+
+REASONING_EFFORT_TEMPLATE = (
+    SYSTEM_SP_TOKEN + "Reasoning Effort: {budget} "
+    "(range 1-100, the higher the value, the more thorough the reasoning)\n\n"
+)
+
+V41_TOOL_CALL_OUTPUT = (
+    '  reason  </think>summary\n\n'
+    '<｜DSML｜ calls>\n'
+    '<｜DSML｜ invoke name="lookup">\n'
+    '<｜DSML｜ parameter name="query" string="true">value'
+    '</｜DSML｜ parameter>\n'
+    '<｜DSML｜ parameter name="limit" string="false">2'
+    '</｜DSML｜ parameter>\n'
+    '</｜DSML｜ invoke>\n'
+    '</｜DSML｜ calls><｜end▁of▁sentence｜>'
+)
+
+
+def make_tool() -> dict:
+    return {
+        "type": "function",
+        "function": {
+            "name": "lookup",
+            "description": "Look up a value",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "query": {"type": "string"},
+                    "limit": {"type": "integer"},
+                },
+            },
+        },
+    }
+
+
+def make_tool_call_messages() -> list:
+    return [
+        {"role": "user", "content": "question"},
+        {
+            "role": "assistant",
+            "reasoning_content": "  reason  ",
+            "content": "summary",
+            "tool_calls": [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "lookup",
+                        "arguments": '{"query":"value","limit":2}',
+                    },
+                }
+            ],
+        },
+    ]
+
+
+# ============================================================
+# Vision
+# ============================================================
+
+def test_v41_renders_images() -> None:
+    prompt, media = encode_messages(
+        [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "inspect"},
+                    {"type": "image_url", "image_url": {"url": "/unused/image.png"}},
+                ],
+            }
+        ],
+        thinking_mode="chat",
+        return_multi_modal_data=True,
+    )
+
+    assert prompt == (
+        '<｜begin▁of▁sentence｜><｜User｜>inspect\n\n'
+        f'{IMAGE_PLACEHOLDER}<｜Assistant｜></think>'
+    )
+    assert media == {"images": [{"type": "image", "url": "/unused/image.png"}]}
+
+
+def test_v41_rejects_image_placeholder_in_text() -> None:
+    with pytest.raises(ValueError):
+        encode_messages(
+            [{"role": "user", "content": f"hi {IMAGE_PLACEHOLDER}"}],
+            thinking_mode="chat",
+        )
+
+
+# ============================================================
+# Reasoning Effort
+# ============================================================
+
+@pytest.mark.parametrize(
+    ("effort", "budget"),
+    [
+        (None, 75),
+        ("low", 50),
+        ("high", 75),
+        ("max", 100),
+        (1, 1),
+        (42, 42),
+        (100, 100),
+    ],
+)
+def test_v41_maps_reasoning_effort_to_1_100_budget(
+    effort: Any,
+    budget: int,
+) -> None:
+    prompt = encode_messages(
+        [{"role": "user", "content": "question"}],
+        thinking_mode="thinking",
+        reasoning_effort=effort,
+    )
+
+    assert prompt == (
+        '<｜begin▁of▁sentence｜>'
+        f'{REASONING_EFFORT_TEMPLATE.format(budget=budget)}'
+        '<｜User｜>question<｜Assistant｜><think>'
+    )
+
+
+def test_v41_only_adds_reasoning_effort_to_first_thinking_message() -> None:
+    messages = [
+        {"role": "system", "content": "system"},
+        {"role": "user", "content": "question"},
+    ]
+
+    later_message = render_message(
+        1, messages, thinking_mode="thinking", reasoning_effort=100
+    )
+    chat_message = render_message(
+        0, messages, thinking_mode="chat", reasoning_effort=100
+    )
+
+    assert "Reasoning Effort:" not in later_message
+    assert "Reasoning Effort:" not in chat_message
+
+
+def test_v41_chat_mode_has_no_reasoning_effort_or_system_token() -> None:
+    prompt = encode_messages(
+        [{"role": "user", "content": "hello"}],
+        thinking_mode="chat",
+        reasoning_effort="max",
+    )
+    assert prompt == '<｜begin▁of▁sentence｜><｜User｜>hello<｜Assistant｜></think>'
+
+
+@pytest.mark.parametrize("effort", [-1, 0, 101, "medium"])
+def test_v41_rejects_out_of_range_or_unknown_reasoning_effort(
+    effort: Any,
+) -> None:
+    with pytest.raises(AssertionError, match=r"int within \[1,100\]"):
+        encode_messages(
+            [{"role": "user", "content": "question"}],
+            thinking_mode="thinking",
+            reasoning_effort=effort,
+        )
+
+
+@pytest.mark.parametrize("effort", [True, False, 1.5])
+def test_v41_rejects_non_string_non_integer_effort_types(effort: Any) -> None:
+    # bool is not `type(...) is int`; float is invalid too
+    with pytest.raises(AssertionError):
+        encode_messages(
+            [{"role": "user", "content": "question"}],
+            thinking_mode="thinking",
+            reasoning_effort=effort,
+        )
+
+
+# ============================================================
+# System token
+# ============================================================
+
+def test_v41_leading_system_message_uses_system_token() -> None:
+    prompt = encode_messages(
+        [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "hello"},
+        ],
+        thinking_mode="chat",
+    )
+    assert prompt == (
+        '<｜begin▁of▁sentence｜><｜System｜>You are a helpful assistant.'
+        '<｜User｜>hello<｜Assistant｜></think>'
+    )
+
+
+def test_v41_mid_conversation_system_message() -> None:
+    prompt = encode_messages(
+        [
+            {"role": "system", "content": "sys"},
+            {"role": "user", "content": "q1"},
+            {"role": "assistant", "content": "a1", "reasoning_content": "r1"},
+            {"role": "system", "content": "mid sys"},
+        ],
+        thinking_mode="thinking",
+        reasoning_effort=88,
+    )
+    # Mid-conversation system gets its own <｜System｜> token and triggers
+    # the assistant generation header afterwards.
+    assert prompt == (
+        '<｜begin▁of▁sentence｜>'
+        f'{REASONING_EFFORT_TEMPLATE.format(budget=88)}'
+        'sys<｜User｜>q1<｜Assistant｜></think>a1<｜end▁of▁sentence｜>'
+        '<｜System｜>mid sys<｜Assistant｜><think>'
+    )
+
+
+# ============================================================
+# DSML tool tags
+# ============================================================
+
+def test_v41_tool_instructions_use_spaced_dsml_tags_in_chat_mode() -> None:
+    prompt = encode_messages(
+        [
+            {"role": "system", "content": "system", "tools": [make_tool()]},
+            {"role": "user", "content": "question"},
+        ],
+        thinking_mode="chat",
+    )
+
+    assert (
+        '<｜DSML｜ calls>\n'
+        '<｜DSML｜ invoke name="$TOOL_NAME">\n'
+        '<｜DSML｜ parameter name="$PARAMETER_NAME" '
+        'string="true|false">$PARAMETER_VALUE</｜DSML｜ parameter>\n'
+        '...\n'
+        '</｜DSML｜ invoke>'
+    ) in prompt
+    assert '<｜DSML｜tool_calls>' not in prompt
+    assert '<｜DSML｜invoke' not in prompt
+    assert '<｜DSML｜parameter' not in prompt
+
+
+def test_v41_renders_spaced_dsml_with_v4_assistant_semantics() -> None:
+    messages = make_tool_call_messages()
+
+    prompt = render_message(1, messages, thinking_mode="thinking")
+
+    assert prompt == V41_TOOL_CALL_OUTPUT
+
+
+def test_v41_parses_spaced_dsml_roundtrip() -> None:
+    messages = make_tool_call_messages()
+
+    parsed = parse_message_from_completion_text(
+        V41_TOOL_CALL_OUTPUT, thinking_mode="thinking"
+    )
+
+    assert parsed["role"] == "assistant"
+    assert parsed["reasoning_content"] == "  reason  "
+    assert parsed["content"] == "summary"
+    assert parsed["tool_calls"]
+    assert parsed["tool_calls"][0]["function"]["name"] == "lookup"
+    assert json.loads(parsed["tool_calls"][0]["function"]["arguments"]) == {
+        "query": "value",
+        "limit": 2,
+    }
+
+    # Re-encoding the parsed message reproduces the original completion text
+    assert encode_messages(
+        [parsed],
+        thinking_mode="thinking",
+        context=messages[:1],
+    ) == V41_TOOL_CALL_OUTPUT
+
+
+def test_v41_parse_rejects_unspaced_v4_dsml() -> None:
+    v4_output = V41_TOOL_CALL_OUTPUT.replace("｜DSML｜ calls", "｜DSML｜tool_calls") \
+        .replace("｜DSML｜ invoke", "｜DSML｜invoke") \
+        .replace("｜DSML｜ parameter", "｜DSML｜parameter")
+    with pytest.raises(AssertionError):
+        parse_message_from_completion_text(v4_output, thinking_mode="thinking")
+
+
+# ============================================================
+# Tool namespaces
+# ============================================================
+
+@pytest.mark.parametrize("location", ["tool", "function"])
+@pytest.mark.parametrize("namespace", ["search", {"name": "search", "description": "Search tools."}])
+def test_v41_renders_namespaced_tool_schemas(location: str, namespace: Any) -> None:
+    tool = make_tool()
+    target = tool if location == "tool" else tool["function"]
+    target["namespace"] = namespace
+    original = copy.deepcopy(tool)
+
+    prompt = encode_messages(
+        [{"role": "system", "content": "system", "tools": [tool]}],
+        thinking_mode="chat",
+    )
+
+    schema = dict(make_tool()["function"], name="search::lookup")
+    if isinstance(namespace, dict):
+        schema["description"] = "Search tools.\nLook up a value"
+    assert json.dumps(schema) in prompt
+    assert '"namespace":' not in prompt
+    assert tool == original
+
+
+@pytest.mark.parametrize("thinking_mode", ["chat", "thinking"])
+@pytest.mark.parametrize("location", ["tool", "function", "qualified_name"])
+def test_v41_namespaced_tool_calls_roundtrip(thinking_mode: str, location: str) -> None:
+    messages = make_tool_call_messages()
+    call = messages[1]["tool_calls"][0]
+    if location == "qualified_name":
+        call["function"]["name"] = "search::lookup"
+    else:
+        target = call if location == "tool" else call["function"]
+        target["namespace"] = "search"
+    original = copy.deepcopy(messages)
+
+    expected = V41_TOOL_CALL_OUTPUT.replace('name="lookup"', 'name="search::lookup"')
+    if thinking_mode == "chat":
+        expected = expected.split("</think>", 1)[1]
+    assert render_message(1, messages, thinking_mode=thinking_mode) == expected
+
+    parsed = parse_message_from_completion_text(expected, thinking_mode=thinking_mode)
+    assert parsed["tool_calls"] == [{
+        "type": "function",
+        "namespace": "search",
+        "function": {
+            "name": "lookup",
+            "arguments": '{"query": "value", "limit": 2}',
+        },
+    }]
+    assert encode_messages(
+        [parsed], thinking_mode=thinking_mode, context=messages[:1]
+    ) == expected
+    assert messages == original
+
+
+def test_v41_keeps_same_named_tools_in_separate_namespaces() -> None:
+    tools, calls = [], []
+    for namespace in (None, "search", "files"):
+        tool = make_tool()
+        call = {
+            "type": "function",
+            "function": {"name": "lookup", "arguments": '{"query":"value"}'},
+        }
+        if namespace is not None:
+            tool["namespace"] = {"name": namespace}
+            call["namespace"] = namespace
+        tools.append(tool)
+        calls.append(call)
+
+    messages = [
+        {"role": "system", "content": "system", "tools": tools},
+        {"role": "user", "content": "question"},
+        {"role": "assistant", "content": "summary", "tool_calls": calls},
+    ]
+    prompt = encode_messages(messages, thinking_mode="chat")
+    for name in ("lookup", "search::lookup", "files::lookup"):
+        assert f'"name": "{name}"' in prompt
+        assert f'<｜DSML｜ invoke name="{name}">' in prompt
+
+    completion = render_message(2, messages, thinking_mode="chat")
+    parsed = parse_message_from_completion_text(completion, thinking_mode="chat")
+    assert "namespace" not in parsed["tool_calls"][0]
+    assert [call.get("namespace") for call in parsed["tool_calls"]] == [None, "search", "files"]
+    assert all(call["function"]["name"] == "lookup" for call in parsed["tool_calls"])
+
+
+def test_v41_does_not_duplicate_a_qualified_namespace() -> None:
+    tool = make_tool()
+    tool["function"]["name"] = "search::lookup"
+    tool["namespace"] = {"name": "search", "description": "Search tools."}
+    schema = enc.tools_from_openai_format([tool])[0]
+    assert schema["name"] == "search::lookup"
+    assert schema["description"] == "Search tools.\nLook up a value"
+
+    messages = make_tool_call_messages()
+    call = messages[1]["tool_calls"][0]
+    call["function"]["name"] = "search::lookup"
+    call["namespace"] = "search"
+    assert render_message(1, messages, thinking_mode="thinking") == (
+        V41_TOOL_CALL_OUTPUT.replace('name="lookup"', 'name="search::lookup"')
+    )
+
+
+@pytest.mark.parametrize(
+    ("name", "namespace", "error"),
+    [
+        ("search::lookup", "files", "Conflicting tool namespaces"),
+        ("search::nested::lookup", None, "Tool name must not contain"),
+        ("lookup", "search::nested", "Tool namespace must not contain"),
+    ],
+)
+def test_v41_rejects_ambiguous_tool_namespaces(name: str, namespace: Any, error: str) -> None:
+    tool = make_tool()
+    tool["function"]["name"] = name
+    tool["namespace"] = namespace
+    with pytest.raises(AssertionError, match=error):
+        enc.tools_from_openai_format([tool])
+
+    messages = make_tool_call_messages()
+    call = messages[1]["tool_calls"][0]
+    call["function"]["name"] = name
+    call["namespace"] = namespace
+    with pytest.raises(AssertionError, match=error):
+        render_message(1, messages, thinking_mode="thinking")
+
+
+# ============================================================
+# Multi-turn flow
+# ============================================================
+
+def test_v41_drop_thinking_without_tools() -> None:
+    prompt = encode_messages(
+        [
+            {"role": "user", "content": "q1"},
+            {"role": "assistant", "content": "a1", "reasoning_content": "r1"},
+            {"role": "user", "content": "q2"},
+        ],
+        thinking_mode="thinking",
+        drop_thinking=True,
+    )
+    # Earlier turn reasoning dropped, </think> form; new turn opens <think>
+    assert '<｜User｜>q1<｜Assistant｜></think>a1<｜end▁of▁sentence｜>' in prompt
+    assert 'r1' not in prompt
+    assert prompt.endswith('<｜User｜>q2<｜Assistant｜><think>')
+
+
+# ============================================================
+# Preprocessing
+# ============================================================
+
+def test_merge_tool_messages_creates_tool_result_blocks() -> None:
+    merged = merge_tool_messages([
+        {"role": "assistant", "content": "", "tool_calls": []},
+        {"role": "tool", "tool_call_id": "a", "content": "r1"},
+        {"role": "tool", "tool_call_id": "b", "content": "r2"},
+    ])
+    assert len(merged) == 2
+    assert merged[1]["role"] == "user"
+    assert [b["type"] for b in merged[1]["content_blocks"]] == ["tool_result", "tool_result"]
+
+
+def test_v41_task_sp_token() -> None:
+    prompt = encode_messages(
+        [{"role": "user", "content": "classify me", "task": "query"}],
+        thinking_mode="chat",
+    )
+    assert prompt.endswith("classify me<｜query｜>")
+    assert "<｜Assistant｜>" not in prompt
+
+
+# ============================================================
+# Golden fixtures from encoding/tests
+# ============================================================
+
+ENCODING_DIR = Path(__file__).resolve().parent
+ENCODING_FIXTURES_DIR = ENCODING_DIR / "tests"
+INFERENCE_EXAMPLES_DIR = ENCODING_DIR.parent / "inference" / "examples"
+
+FIXTURE_CASE_IDS = sorted(
+    int(p.stem.split("_")[-1])
+    for p in ENCODING_FIXTURES_DIR.glob("test_input_*.json")
+)
+
+
+@pytest.mark.parametrize("case_id", FIXTURE_CASE_IDS)
+def test_examples_encoding_golden_outputs(case_id: int) -> None:
+    """Each tests/encoding input must encode to its checked-in golden output."""
+    input_file = ENCODING_FIXTURES_DIR / f"test_input_{case_id}.json"
+    output_file = ENCODING_FIXTURES_DIR / f"test_output_{case_id}.txt"
+    assert output_file.exists(), f"missing golden output: {output_file.name} (run tests/encoding/regen_outputs.py)"
+
+    case = enc.load_cases(str(input_file))[0]
+    prompt, _ = enc.encode_case(case, thinking_mode="chat")
+
+    assert prompt == output_file.read_text(), (
+        f"{output_file.name} is stale; regenerate with tests/encoding/regen_outputs.py"
+    )
+
+
+def test_examples_v41_output_uses_v41_format_markers() -> None:
+    """Sanity-check the V4.1 goldens actually exercise V4.1-specific format."""
+    # case 1: tool calls with spaced DSML tags
+    out1 = (ENCODING_FIXTURES_DIR / "test_output_1.txt").read_text()
+    assert '<｜DSML｜ calls>' in out1 and '<｜DSML｜ invoke name="get_weather">' in out1
+    assert '<｜DSML｜tool_calls>' not in out1
+
+    # case 5: numeric reasoning effort behind the system token
+    out5 = (ENCODING_FIXTURES_DIR / "test_output_5.txt").read_text()
+    assert out5.startswith(
+        '<｜begin▁of▁sentence｜>' + REASONING_EFFORT_TEMPLATE.format(budget=100)
+    )
+    assert out5.count(IMAGE_PLACEHOLDER) == 2
+
+
+def test_examples_vl_txt_and_json_encode_identically() -> None:
+    """The TXT (last block of example.txt) and JSON vision examples must encode identically."""
+    txt = (INFERENCE_EXAMPLES_DIR / "example.txt").read_text().rstrip("\n").split("\n\n")[-1]
+    messages = [{"role": "user", "content": enc.parse_tagged_text(txt)}]
+    p1, m1 = encode_messages(messages, thinking_mode="chat", return_multi_modal_data=True)
+
+    case = enc.load_cases(str(INFERENCE_EXAMPLES_DIR / "example_harmony.json"))[0]
+    p2, m2 = enc.encode_case(case, thinking_mode="chat")
+
+    assert p1 == p2
+    assert m1["images"] == m2
+    assert len(m2) == 2
+
+
+def test_examples_harmony_cases_encode() -> None:
+    """All example_harmony.json cases encode without error."""
+    cases = enc.load_cases(str(INFERENCE_EXAMPLES_DIR / "example_harmony.json"))
+    assert len(cases) == 4
+
+    # case 1 (vision) is covered by test_examples_vl_txt_and_json_encode_identically
+
+    # cases are pure OpenAI format: mode/effort are passed at call time
+    prompt = encode_messages(
+        cases[1]["messages"], thinking_mode="thinking", reasoning_effort=75
+    )
+    assert REASONING_EFFORT_TEMPLATE.format(budget=75) in prompt
+
+    # case 3: tools with spaced DSML tags
+    prompt, _ = enc.encode_case(cases[2], thinking_mode="chat")
+    assert '<｜DSML｜ calls>' in prompt
+
+    # case 4: mid-conversation system message triggers assistant header
+    prompt, _ = enc.encode_case(cases[3], thinking_mode="chat")
+    assert '<｜System｜>Mid-conversation instruction update' in prompt
+    assert prompt.endswith('<｜Assistant｜></think>')
+
+
+if __name__ == "__main__":
+    import sys
+    sys.exit(pytest.main([__file__, "-v"]))

encoding/tests/test_input_1.json

@@ -0,0 +1,82 @@
+{
+    "thinking_mode": "thinking",
+    "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 like in Beijing?"
+        },
+        {
+            "role": "assistant",
+            "reasoning_content": "The user wants the weather in Beijing. I should call get_weather.",
+            "content": "",
+            "tool_calls": [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "get_weather",
+                        "arguments": "{\"location\": \"Beijing\", \"unit\": \"celsius\"}"
+                    }
+                }
+            ]
+        },
+        {
+            "role": "tool",
+            "tool_call_id": "call_0",
+            "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\u00b0C 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,93 @@
+[
+  {
+    "role": "system",
+    "content": "该助手为DeepSeek，由深度求索公司创造。"
+  },
+  {
+    "role": "latest_reminder",
+    "content": "2026-02-21,星期六,广州,App,中文"
+  },
+  {
+    "role": "system",
+    "content": "CITATION 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
+          }
+        }
+      },
+      {
+        "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
+                    },
+                    "loc": {"type": "integer", "description": "Start line", "default": -1},
+                    "num_lines": {"type": "integer", "description": "", "default": -1}
+                  },
+                  "additionalProperties": false
+                },
+                "description": ""
+              }
+            },
+            "required": ["open_list"],
+            "additionalProperties": false
+          }
+        }
+      }
+    ]
+  },
+  {
+    "role": "user",
+    "content": "小柴胡冲剂和布洛芬能一起吃吗？"
+  },
+  {
+    "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": "热海大滚锅在中国乃至全球的地热奇观中占有重要地位，但“世界著名”的称号更侧重于它作为独特的地质现象和旅游景点。",
+    "mask": 1
+  },
+  {
+    "role": "user",
+    "content": "世界著名温泉有哪些",
+    "task": "action"
+  },
+  {
+    "role": "assistant",
+    "content": "Search"
+  }
+]

encoding/tests/test_input_5.json

@@ -0,0 +1,39 @@
+{
+    "thinking_mode": "thinking",
+    "reasoning_effort": "max",
+    "messages": [
+        {
+            "role": "system",
+            "content": "You are a helpful vision assistant."
+        },
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "text",
+                    "text": "请按“第一张、第二张”的顺序回答：第一张图"
+                },
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": "examples/images/carrots.jpeg"
+                    }
+                },
+                {
+                    "type": "text",
+                    "text": "和第二张图"
+                },
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": "examples/images/corn.jpeg"
+                    }
+                },
+                {
+                    "type": "text",
+                    "text": "中分别是什么食材？它们通常食用的部位分别是什么？"
+                }
+            ]
+        }
+    ]
+}

encoding/tests/test_output_1.txt

@@ -0,0 +1,38 @@
+<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)
+
+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｜ calls>" block like the following:
+
+<｜DSML｜ 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｜ 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 like in Beijing?<｜Assistant｜><think>The user wants the weather in Beijing. I should call get_weather.</think>
+
+<｜DSML｜ 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｜ 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｜><｜System｜>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 capital of France is Paris.<｜end▁of▁sentence｜>

encoding/tests/test_output_3.txt

@@ -0,0 +1,35 @@
+<｜begin▁of▁sentence｜><｜System｜>该助手为DeepSeek，由深度求索公司创造。<｜latest_reminder｜>2026-02-21,星期六,广州,App,中文<｜System｜>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｜ calls>" block like the following:
+
+<｜DSML｜ 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｜ 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}}
+{"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}, "loc": {"type": "integer", "description": "Start line", "default": -1}, "num_lines": {"type": "integer", "description": "", "default": -1}}, "additionalProperties": false}, "description": ""}}, "required": ["open_list"], "additionalProperties": false}}
+
+You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
+<｜User｜>小柴胡冲剂和布洛芬能一起吃吗？<｜Assistant｜></think>
+
+<｜DSML｜ calls>
+<｜DSML｜ invoke name="search">
+<｜DSML｜ parameter name="queries" string="true">小柴胡冲剂 布洛芬 相互作用 一起吃</｜DSML｜ parameter>
+</｜DSML｜ invoke>
+</｜DSML｜ calls><｜end▁of▁sentence｜><｜User｜><tool_result>[0]</tool_result><｜Assistant｜></think>请及时就医。<｜end▁of▁sentence｜>

encoding/tests/test_output_4.txt

@@ -0,0 +1,2 @@
+<｜begin▁of▁sentence｜><｜System｜>该助手为DeepSeek-V3，由深度求索公司创造。
+今天是2025年10月17日，星期五。<｜latest_reminder｜>2024-11-15,上海市,App,中文<｜User｜>热海大滚锅是世界著名温泉吗<｜Assistant｜></think>热海大滚锅在中国乃至全球的地热奇观中占有重要地位，但“世界著名”的称号更侧重于它作为独特的地质现象和旅游景点。<｜end▁of▁sentence｜><｜User｜>世界著名温泉有哪些<｜Assistant｜></think><｜action｜>Search<｜end▁of▁sentence｜>

encoding/tests/test_output_5.txt

@@ -0,0 +1,11 @@
+<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 100 (range 1-100, the higher the value, the more thorough the reasoning)
+
+You are a helpful vision assistant.<｜User｜>请按“第一张、第二张”的顺序回答：第一张图
+
+<｜deepseek_image｜>
+
+和第二张图
+
+<｜deepseek_image｜>
+
+中分别是什么食材？它们通常食用的部位分别是什么？<｜Assistant｜><think>

evaluation/README.md

@@ -0,0 +1,103 @@
+# Running DeepSWE with `dsh-minimal` and `mini-swe-agent`
+
+## 1. Prerequisites
+
+- Docker, running and able to pull images
+- Python 3.12+ and [uv](https://docs.astral.sh/uv/)
+- An endpoint and key for any DeepSeek-API-compatible service (the DeepSeek official API is used as the example below)
+
+```sh
+export DEEPSEEK_API_KEY=sk-your-key-here
+export DEEPSEEK_BASE_URL=https://api.deepseek.com
+```
+
+## 2. Get Pier and DeepSWE
+
+```sh
+git clone https://github.com/datacurve-ai/pier.git
+git -C pier checkout 0c802fc067a425345b24d1c69411aa98acf61a1d
+
+git clone https://github.com/datacurve-ai/deep-swe.git
+git -C deep-swe checkout 0b9fabbb63b9104d678fe965e1632f2dd9eaa2ea
+```
+
+## 3. Patch and install Pier
+
+`dsh-minimal.patch` ships next to this document. Treat it as a **reference patch** and adapt it to your own setup.
+
+```sh
+cd pier
+git apply /path/to/dsh-minimal.patch
+uv sync
+```
+
+What the patch changes:
+
+- **Adds the `dsh-minimal` agent**, which drives the Harness SDK and folds its event stream into a Pier ATIF trajectory. The SDK artifact is never installed into the image: step 4's `--mounts-json` bind-mounts it read-only into the sandbox, so no trial installs anything.
+- **Appends a runtime-constraints section to the task instruction for both agents**: work in `/app`, leave `/tests` alone, no network or package mirror.
+- **Passes test-runner concurrency caps into the container**: Docker's `--cpus` is only a quota, so `nproc` inside the container reports the host's core count and test runners size their worker pools from that rather than from the container's share.
+- **Enables IPv6 loopback in the container**: Docker disables it by default, so suites that bind `::1` are skipped and scored as failures.
+- **Makes `--mounts-json` additive instead of replacing the default mounts**, keeping the `/logs` binds that carry agent logs and collected patches.
+
+## 4. Run the suite
+
+Both agents take the same task set, concurrency, and `--no-delete` (which keeps the task images cached between trials). Repeat each run with a different `--job-name` and average the results.
+
+Each trial's container takes the 2 CPUs and 8 GB its task declares, so size `-n` against the host's cores and memory.
+
+### `mini-swe-agent`
+
+Pier installs it into each task image at trial time, so no host-side preparation is needed.
+
+```sh
+uv run pier run \
+  -p ../deep-swe/tasks \
+  --agent mini-swe-agent \
+  --model deepseek/deepseek-flash \
+  --ak reasoning_effort=max \
+  --ak cost_limit=0 \
+  --ae DEEPSEEK_API_KEY="$DEEPSEEK_API_KEY" \
+  --ae DEEPSEEK_BASE_URL="$DEEPSEEK_BASE_URL" \
+  -n 32 --no-delete -r 2 --job-name deepswe-mini-run1 -y
+```
+
+- `--model` takes a litellm-style `provider/model` string.
+
+### `dsh-minimal`
+
+Install the Harness SDK artifact once on the host, then bind-mount it read-only into every container.
+
+```sh
+mkdir -p ~/dsh-minimal && cd ~/dsh-minimal
+uv pip install --target dsh-dist \
+  --python-version 3.12 --python-platform x86_64-manylinux_2_28 \
+  'deepseek-harness-sdk==0.1.5.*'
+```
+
+```sh
+uv run pier run \
+  -p ../deep-swe/tasks \
+  --agent dsh-minimal \
+  --model deepseek-flash \
+  --ak reasoning_effort=max \
+  --ae DEEPSEEK_API_KEY="$DEEPSEEK_API_KEY" \
+  --ae DEEPSEEK_BASE_URL="$DEEPSEEK_BASE_URL" \
+  --mounts-json '[{"type":"bind","source":"'"$HOME"'/dsh-minimal/dsh-dist","target":"/opt/dsh-minimal","read_only":true}]' \
+  -n 32 --no-delete --job-name deepswe-dsh-run1 -y
+```
+
+- In `--mounts-json`, `source` is the absolute path of the `dsh-dist` directory above; `target` is always `/opt/dsh-minimal`.
+
+## 5. Read the results
+
+```
+jobs/<job-name>/
+  result.json              pass rate and token totals
+  <task>__<id>/
+    result.json            reward, fail-to-pass / pass-to-pass counts, tokens
+    agent/trajectory.json  full ATIF trajectory (dsh-minimal)
+    agent/mini-swe-agent.trajectory.json   mini-swe-agent trajectory
+    verifier/              reward.json and test output
+```
+
+Browse a job with `uv run pier view jobs/<job-name>`.

evaluation/dsh-minimal.patch

@@ -0,0 +1,711 @@
+diff --git a/src/pier/agents/factory.py b/src/pier/agents/factory.py
+index 665fadd..503a45d 100644
+--- a/src/pier/agents/factory.py
++++ b/src/pier/agents/factory.py
+@@ -6,6 +6,7 @@ from pier.agents.installed.antigravity_sdk import AntigravitySDK
+ from pier.agents.installed.claude_code import ClaudeCode
+ from pier.agents.installed.codex import Codex
+ from pier.agents.installed.cursor_cli import CursorCli
++from pier.agents.installed.dsh_minimal import DshMinimal
+ from pier.agents.installed.gemini_cli import GeminiCli
+ from pier.agents.installed.mini_swe_agent import MiniSweAgent
+ from pier.agents.installed.opencode import OpenCode
+@@ -24,6 +25,7 @@ class AgentFactory:
+         ClaudeCode,
+         Codex,
+         CursorCli,
++        DshMinimal,
+         GeminiCli,
+         MiniSweAgent,
+         OpenCode,
+diff --git a/src/pier/agents/installed/base.py b/src/pier/agents/installed/base.py
+index 16f6a84..1fd1050 100644
+--- a/src/pier/agents/installed/base.py
++++ b/src/pier/agents/installed/base.py
+@@ -15,6 +15,18 @@ from pier.utils.env import parse_bool_env_value
+ from pier.utils.templating import render_prompt_template
+ 
+ 
++RUNTIME_CONSTRAINTS = (
++    "\n## Runtime constraints\n"
++    "- Work in `/app`; do not modify files under `/tests`.\n"
++    "- No network or mirror access; use only dependencies already in the image.\n"
++)
++
++
++def with_runtime_constraints(instruction: str) -> str:
++    """Append the sandbox constraints to a task instruction."""
++    return instruction.rstrip("\n") + "\n" + RUNTIME_CONSTRAINTS
++
++
+ class NonZeroAgentExitCodeError(RuntimeError):
+     """Raised when the agent process exits with a non-zero exit code."""
+ 
+@@ -394,12 +406,14 @@ class BaseInstalledAgent(BaseAgent, ABC):
+         return instruction
+ 
+     @abstractmethod
+-    def install_spec(self) -> AgentInstallSpec:
+-        """Declarative install steps executed at setup and inlined into Dockerfile builds."""
++    def install_spec(self) -> AgentInstallSpec | None:
++        """Declarative install steps executed at setup and inlined into Dockerfile
++        builds, or ``None`` for an agent already present in the environment."""
+ 
+     async def install(self, environment: BaseEnvironment) -> None:
+         """Run each step from :meth:`install_spec` with matching privilege."""
+-        for step in self.install_spec().steps:
++        spec = self.install_spec()
++        for step in spec.steps if spec is not None else ():
+             if step.user == "root":
+                 await self.exec_as_root(environment, command=step.run, env=step.env)
+             else:
+diff --git a/src/pier/agents/installed/dsh_minimal.py b/src/pier/agents/installed/dsh_minimal.py
+new file mode 100644
+index 0000000..c84a1b1
+--- /dev/null
++++ b/src/pier/agents/installed/dsh_minimal.py
+@@ -0,0 +1,131 @@
++import shlex
++from pathlib import Path
++from typing import Any, override
++
++from pier.agents.installed.base import (
++    BaseInstalledAgent,
++    with_prompt_template,
++    with_runtime_constraints,
++)
++from pier.agents.network import allowlist_from_urls
++from pier.environments.base import BaseEnvironment
++from pier.models.agent.context import AgentContext
++from pier.models.agent.name import AgentName
++from pier.models.agent.network import NetworkAllowlist
++from pier.models.trajectories import Trajectory
++from pier.utils.trajectory_metrics import populate_context_from_final_metrics
++
++DIST = "/opt/dsh-minimal"
++DEFAULT_MODEL = "deepseek-flash"
++DEFAULT_BASE_URL = "https://api.deepseek.com"
++REASONING_EFFORTS = ("low", "high", "max")
++
++
++class DshMinimal(BaseInstalledAgent):
++    """DeepSeek Harness `sdk-minimal` profile driven through its Python SDK.
++
++    The SDK and its bundled runtime executable come from a `pip install --target`
++    tree bind-mounted read-only at ``DIST``, so no trial installs anything.
++    """
++
++    SUPPORTS_ATIF = True
++    _TRAJECTORY = "trajectory.json"
++    # The distribution mount is read-only, so the runner lands beside it.
++    _RUNNER = "/tmp/dsh-minimal-runner.py"
++
++    def __init__(
++        self, *args: Any, reasoning_effort: str = "max", **kwargs: Any
++    ) -> None:
++        if reasoning_effort not in REASONING_EFFORTS:
++            raise ValueError(
++                f"reasoning_effort must be one of {REASONING_EFFORTS}, "
++                f"got {reasoning_effort!r}"
++            )
++        self._reasoning_effort = reasoning_effort
++        kwargs["model_name"] = kwargs.get("model_name") or DEFAULT_MODEL
++        super().__init__(*args, **kwargs)
++
++    @staticmethod
++    @override
++    def name() -> str:
++        return AgentName.DSH_MINIMAL.value
++
++    @override
++    def install_spec(self) -> None:
++        """The distribution is bind-mounted, so no image layer is needed."""
++        return None
++
++    @override
++    def get_version_command(self) -> str:
++        return f"sed -n 's/^Version: //p' {DIST}/deepseek_harness_sdk-*.dist-info/METADATA"
++
++    @override
++    async def setup(self, environment: BaseEnvironment) -> None:
++        await super().setup(environment)
++        runner = Path(__file__).with_name("dsh_minimal_runner.py")
++        await environment.upload_file(runner, self._RUNNER)
++        await self.exec_as_root(environment, f"chmod a+r {self._RUNNER}")
++
++    def _base_url(self) -> str:
++        return self._get_env("DEEPSEEK_BASE_URL") or DEFAULT_BASE_URL
++
++    @override
++    def network_allowlist(self) -> NetworkAllowlist:
++        return allowlist_from_urls([self._base_url()])
++
++    @override
++    def populate_context_post_run(self, context: AgentContext) -> None:
++        path = self.logs_dir / self._TRAJECTORY
++        if not path.exists():
++            self.logger.debug("No dsh-minimal trajectory found at %s", path)
++            return
++        try:
++            trajectory = Trajectory.model_validate_json(path.read_text())
++        except (OSError, ValueError):
++            self.logger.exception("Failed to parse dsh-minimal trajectory")
++            return
++        if trajectory.final_metrics is not None:
++            populate_context_from_final_metrics(context, trajectory.final_metrics)
++        context.n_agent_steps = sum(step.source == "agent" for step in trajectory.steps)
++
++    @override
++    @with_prompt_template
++    async def run(
++        self,
++        instruction: str,
++        environment: BaseEnvironment,
++        context: AgentContext,
++    ) -> None:
++        api_key = self._get_env("DEEPSEEK_API_KEY")
++        if not api_key:
++            raise ValueError("DEEPSEEK_API_KEY environment variable must be set")
++
++        instruction = with_runtime_constraints(instruction)
++        agent_dir = environment.env_paths.agent_dir.as_posix()
++        env = self.build_process_env(
++            {
++                "DEEPSEEK_API_KEY": api_key,
++                "DEEPSEEK_BASE_URL": self._base_url(),
++                # The Node runtime reaches the model only through Pier's egress
++                # proxy, and its fetch ignores the proxy variables without this.
++                "NODE_USE_ENV_PROXY": "1",
++                "PYTHONPATH": DIST,
++                "DSH_MODEL": self._parsed_model_name,
++                "DSH_REASONING_EFFORT": self._reasoning_effort,
++                # Required, and kept off the /logs bind mount: it materializes many
++                # files the host cannot read anyway. `trajectory.json` is the record.
++                "DSH_HOME": "/tmp/dsh-home",
++                "SESSION_ID": environment.session_id,
++            }
++        )
++
++        await self.exec_as_agent(
++            environment,
++            command=(
++                f"python3 {self._RUNNER} "
++                f"--instruction {shlex.quote(instruction)} "
++                f"--trajectory-path {agent_dir}/{self._TRAJECTORY} "
++                f"2>&1 </dev/null | stdbuf -oL tee {agent_dir}/dsh-minimal.txt"
++            ),
++            env=env,
++        )
+diff --git a/src/pier/agents/installed/dsh_minimal_runner.py b/src/pier/agents/installed/dsh_minimal_runner.py
+new file mode 100644
+index 0000000..8b69d7d
+--- /dev/null
++++ b/src/pier/agents/installed/dsh_minimal_runner.py
+@@ -0,0 +1,307 @@
++"""Run the DeepSeek Harness `sdk-minimal` profile and persist its stream as ATIF v1.7."""
++
++from __future__ import annotations
++
++import argparse
++import json
++import os
++import sys
++import traceback
++from datetime import datetime, timezone
++from importlib.metadata import PackageNotFoundError, version
++from pathlib import Path
++from typing import Any
++
++from deepseek_harness import DeepSeekHarness
++from deepseek_harness.models import Notification
++
++PROVIDER = "deepseek-official"
++
++
++def _iso(epoch_ms: Any) -> str | None:
++    if not isinstance(epoch_ms, (int, float)):
++        return None
++    return datetime.fromtimestamp(epoch_ms / 1000, tz=timezone.utc).isoformat()
++
++
++def _text(blocks: Any, kind: str) -> str:
++    if not isinstance(blocks, list):
++        return ""
++    return "".join(
++        str(block.get("text") or "")
++        for block in blocks
++        if isinstance(block, dict) and block.get("type") == kind
++    )
++
++
++def _flatten(blocks: Any) -> str:
++    if isinstance(blocks, str):
++        return blocks
++    if not isinstance(blocks, list):
++        return json.dumps(blocks, ensure_ascii=False)
++    return "".join(
++        str(block.get("text") or "")
++        if isinstance(block, dict) and block.get("type") == "text"
++        else json.dumps(block, ensure_ascii=False)
++        for block in blocks
++    )
++
++
++def _arguments(raw: Any) -> dict[str, Any]:
++    """ATIF wants an object; the wire carries the model's raw argument string."""
++    if isinstance(raw, dict):
++        return raw
++    if isinstance(raw, str):
++        try:
++            parsed = json.loads(raw)
++        except json.JSONDecodeError:
++            return {"_unparsed_arguments": raw}
++        return parsed if isinstance(parsed, dict) else {"_arguments": parsed}
++    return {}
++
++
++def _metrics(usage: dict[str, Any]) -> dict[str, Any]:
++    """The adapter reports disjoint counts, so the cache legs rejoin prompt tokens."""
++    cache_read = usage.get("cacheReadTokens") or 0
++    cache_write = usage.get("cacheWriteTokens") or 0
++    metrics = {
++        "prompt_tokens": (usage.get("inputTokens") or 0) + cache_read + cache_write,
++        "completion_tokens": usage.get("outputTokens") or 0,
++        "cached_tokens": cache_read,
++    }
++    if usage.get("reasoningTokens") is not None:
++        metrics["extra"] = {"reasoning_tokens": usage["reasoningTokens"]}
++    return metrics
++
++
++class Collector:
++    """Fold root-session events into one ATIF step per model call."""
++
++    def __init__(self, instruction: str, model: str, effort: str, session: str) -> None:
++        self.model = model
++        self.effort = effort
++        self.session = session
++        self.steps: list[dict[str, Any]] = [
++            {
++                "step_id": 1,
++                "timestamp": datetime.now(tz=timezone.utc).isoformat(),
++                "source": "user",
++                "message": instruction,
++            }
++        ]
++        # (turn, step) identifies one model call; callId maps a later `tool/result`
++        # back to the step that issued it, which is never the step it arrives in.
++        self._by_call: dict[tuple[int, int], dict[str, Any]] = {}
++        self._by_call_id: dict[str, dict[str, Any]] = {}
++        self.turn_end_reasons: list[dict[str, Any]] = []
++        self.llm_retries = 0
++
++    def _step(self, turn: Any, step: Any, timestamp: str | None) -> dict[str, Any]:
++        key = (int(turn or 0), int(step or 0))
++        if key not in self._by_call:
++            created = {
++                "step_id": len(self.steps) + 1,
++                "timestamp": timestamp,
++                "source": "agent",
++                "message": "",
++                "model_name": self.model,
++                "reasoning_effort": self.effort,
++                "llm_call_count": 1,
++            }
++            self.steps.append(created)
++            self._by_call[key] = created
++        return self._by_call[key]
++
++    def _tool_call(self, step: dict[str, Any], call_id: Any, name: Any, raw: Any) -> None:
++        if not isinstance(call_id, str) or not call_id:
++            return
++        calls = step.setdefault("tool_calls", [])
++        entry = {
++            "tool_call_id": call_id,
++            "function_name": str(name),
++            "arguments": _arguments(raw),
++        }
++        for index, existing in enumerate(calls):
++            if existing["tool_call_id"] == call_id:
++                calls[index] = entry
++                break
++        else:
++            calls.append(entry)
++        self._by_call_id[call_id] = step
++
++    def record(self, event: dict[str, Any]) -> None:
++        kind = event.get("type")
++        data = event.get("data") if isinstance(event.get("data"), dict) else {}
++        timestamp = _iso(event.get("time"))
++
++        if kind == "assistant/message":
++            step = self._step(data.get("turn"), data.get("step"), timestamp)
++            message = data.get("message") if isinstance(data.get("message"), dict) else {}
++            content = message.get("content")
++            step["message"] = _text(content, "text")
++            if reasoning := _text(content, "reasoning"):
++                step["reasoning_content"] = reasoning
++            source = message.get("source")
++            if isinstance(source, dict) and isinstance(source.get("model"), str):
++                step["model_name"] = source["model"]
++            if isinstance(data.get("usage"), dict):
++                step["metrics"] = _metrics(data["usage"])
++            if data.get("interrupted"):
++                step.setdefault("extra", {})["interrupted"] = True
++            for block in content if isinstance(content, list) else []:
++                if isinstance(block, dict) and block.get("type") == "tool-call":
++                    self._tool_call(
++                        step, block.get("id"), block.get("name"), block.get("arguments")
++                    )
++        elif kind == "tool/call":
++            step = self._step(data.get("turn"), data.get("step"), timestamp)
++            self._tool_call(
++                step, data.get("callId"), data.get("name"), data.get("arguments")
++            )
++        elif kind == "tool/result":
++            self._result(data)
++        elif kind == "turn/end":
++            reason = data.get("reason")
++            self.turn_end_reasons.append(reason if isinstance(reason, dict) else {})
++        elif kind == "llm/retry":
++            self.llm_retries += 1
++
++    def _result(self, data: dict[str, Any]) -> None:
++        message = data.get("message") if isinstance(data.get("message"), dict) else {}
++        for block in message.get("content") or []:
++            if not isinstance(block, dict) or block.get("type") != "tool-result":
++                continue
++            step = self._by_call_id.get(block.get("toolCallId"))
++            if step is None:
++                continue
++            entry: dict[str, Any] = {
++                "source_call_id": block["toolCallId"],
++                "content": _flatten(block.get("content")),
++            }
++            if block.get("isError"):
++                entry["extra"] = {"is_error": True, "error": data.get("error")}
++            results = step.setdefault("observation", {"results": []})["results"]
++            for index, existing in enumerate(results):
++                if existing["source_call_id"] == entry["source_call_id"]:
++                    results[index] = entry
++                    break
++            else:
++                results.append(entry)
++
++    def trajectory(self, finish_reason: str | None, failure: str | None) -> dict[str, Any]:
++        prompt = completion = cached = peak = 0
++        for step in self.steps:
++            metrics = step.get("metrics")
++            if not isinstance(metrics, dict):
++                continue
++            prompt += metrics["prompt_tokens"]
++            completion += metrics["completion_tokens"]
++            cached += metrics["cached_tokens"]
++            peak = max(peak, metrics["prompt_tokens"])
++        extra = {
++            "finish_reason": finish_reason,
++            "turn_end_reasons": self.turn_end_reasons,
++            "llm_retry_count": self.llm_retries,
++        }
++        if failure is not None:
++            extra["failure"] = failure
++        return {
++            "schema_version": "ATIF-v1.7",
++            "session_id": self.session,
++            "agent": {"name": "dsh-minimal", "version": _version(), "model_name": self.model},
++            "steps": self.steps,
++            "final_metrics": {
++                "total_prompt_tokens": prompt,
++                "total_completion_tokens": completion,
++                "total_cached_tokens": cached,
++                "total_steps": len(self.steps),
++                # The sdk-minimal profile has no compaction, so no summary step
++                # can replace a transcript prefix.
++                "extra": {"peak_context_tokens": peak, "summarization_count": 0},
++            },
++            "extra": extra,
++        }
++
++
++def _version() -> str:
++    try:
++        return version("deepseek-harness-sdk")
++    except PackageNotFoundError:
++        return "unknown"
++
++
++def _write(path: Path, value: Any) -> None:
++    path.parent.mkdir(parents=True, exist_ok=True)
++    temporary = path.with_suffix(f"{path.suffix}.tmp")
++    temporary.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n")
++    temporary.replace(path)
++
++
++def main() -> None:
++    parser = argparse.ArgumentParser()
++    parser.add_argument("--instruction", required=True)
++    parser.add_argument("--trajectory-path", required=True)
++    args = parser.parse_args()
++
++    # Task images can leave the agent user on a 077 umask; the session log and
++    # trajectory have to stay readable to the host user that collects /logs.
++    os.umask(0o022)
++
++    model = os.environ["DSH_MODEL"]
++    effort = os.environ.get("DSH_REASONING_EFFORT", "max")
++    session = os.environ.get("SESSION_ID") or "pier-session"
++    path = Path(args.trajectory_path)
++    collector = Collector(args.instruction, model, effort, session)
++    state: dict[str, Any] = {"finish_reason": None, "failure": None}
++
++    def checkpoint() -> None:
++        _write(path, collector.trajectory(state["finish_reason"], state["failure"]))
++
++    def on_notification(notification: Notification) -> None:
++        payload = notification.payload
++        event = payload.get("event") if isinstance(payload, dict) else None
++        if (
++            notification.method != "session.event"
++            or not isinstance(event, dict)
++            # Token deltas repeat what their `assistant/message` already carries.
++            or event.get("type") == "assistant/chunk"
++            or payload.get("sessionId") != session
++        ):
++            return
++        collector.record(event)
++        checkpoint()
++
++    checkpoint()
++    print(f"dsh-minimal: model={model} effort={effort} session={session}")
++    try:
++        with DeepSeekHarness(
++            profile="sdk-minimal",
++            provider=PROVIDER,
++            model=model,
++            reasoning_effort=effort,
++            cwd=os.getcwd(),
++            dsh_home=os.environ["DSH_HOME"],
++        ) as harness:
++            result = harness.run(
++                args.instruction, session_id=session, on_notification=on_notification
++            )
++        state["finish_reason"] = result.finish_reason
++        print(f"dsh-minimal: finish_reason={result.finish_reason}")
++        print(f"dsh-minimal: final_response={result.final_response}")
++    except Exception as error:  # noqa: BLE001 - the trajectory must survive any failure
++        state["failure"] = f"{type(error).__name__}: {error}"
++        traceback.print_exc()
++    finally:
++        checkpoint()
++
++    # `max-tokens` is a bounded model outcome the verifier can still grade; anything
++    # else non-terminal means the harness never produced a usable turn.
++    if state["failure"] is not None or state["finish_reason"] not in (
++        "completed",
++        "max-tokens",
++    ):
++        sys.exit(1)
++
++
++if __name__ == "__main__":
++    main()
+diff --git a/src/pier/agents/installed/mini_swe_agent.py b/src/pier/agents/installed/mini_swe_agent.py
+index 75d25b6..d283f66 100644
+--- a/src/pier/agents/installed/mini_swe_agent.py
++++ b/src/pier/agents/installed/mini_swe_agent.py
+@@ -11,6 +11,7 @@ from pier.agents.installed.base import (
+     BaseInstalledAgent,
+     CliFlag,
+     with_prompt_template,
++    with_runtime_constraints,
+ )
+ from pier.agents.network import allowlist_from_urls, collect_url_values
+ from pier.agents.utils import get_api_key_var_names_from_model_name
+@@ -821,7 +822,7 @@ mini-swe-agent --help
+     async def run(
+         self, instruction: str, environment: BaseEnvironment, context: AgentContext
+     ) -> None:
+-        augmented_instruction = instruction
++        augmented_instruction = with_runtime_constraints(instruction)
+         if self.mcp_servers:
+             mcp_info = "\n\nMCP Servers:\nThe following MCP servers are available for this task.\n"
+             for s in self.mcp_servers:
+diff --git a/src/pier/environments/docker/__init__.py b/src/pier/environments/docker/__init__.py
+index 4bb8f11..04a8c6e 100644
+--- a/src/pier/environments/docker/__init__.py
++++ b/src/pier/environments/docker/__init__.py
+@@ -25,6 +25,9 @@ def write_resources_compose_file(
+     *,
+     cpu_request: int | None = None,
+     cpu_limit: int | None = None,
++    environment: dict[str, str] | None = None,
++    sysctls: dict[str, str] | None = None,
++    volumes: list[str] | None = None,
+     memory_request_mb: int | None = None,
+     memory_limit_mb: int | None = None,
+ ) -> Path:
+@@ -46,7 +49,13 @@ def write_resources_compose_file(
+     if reservations:
+         resources["reservations"] = reservations
+ 
+-    main = {"deploy": {"resources": resources}} if resources else {}
++    main: dict[str, object] = {"deploy": {"resources": resources}} if resources else {}
++    if environment:
++        main["environment"] = dict(environment)
++    if sysctls:
++        main["sysctls"] = dict(sysctls)
++    if volumes:
++        main["volumes"] = list(volumes)
+     compose = {"services": {"main": main}}
+     path.parent.mkdir(parents=True, exist_ok=True)
+     path.write_text(json.dumps(compose, indent=2))
+diff --git a/src/pier/environments/docker/docker.py b/src/pier/environments/docker/docker.py
+index 2ed4e92..bc671be 100644
+--- a/src/pier/environments/docker/docker.py
++++ b/src/pier/environments/docker/docker.py
+@@ -31,6 +31,11 @@ from pier.environments.docker import (
+     write_mounts_compose_file,
+     write_resources_compose_file,
+ )
++from pier.environments.docker.parallelism import (
++    CPU_CLAMP_PATH,
++    cpu_clamp_source,
++    parallelism_env,
++)
+ from pier.models.environment_type import EnvironmentType
+ from pier.models.task.config import EnvironmentConfig, TaskOS
+ from pier.models.trial.config import ResourceMode, ServiceVolumeConfig
+@@ -102,6 +107,11 @@ class DockerEnvironment(BaseEnvironment):
+ 
+     _DOCKER_COMPOSE_WINDOWS_KEEPALIVE_PATH = COMPOSE_WINDOWS_KEEPALIVE_PATH
+ 
++    # Docker leaves IPv6 off in the container netns, so loopback has no ::1 and
++    # suites that bind it skip where a real Linux host passes. These environments
++    # always get their own netns, which is what net.* sysctls require.
++    _LINUX_SYSCTLS = {"net.ipv6.conf.all.disable_ipv6": "0"}
++
+     # Class-level lock per image name to prevent parallel builds of the same image.
+     _image_build_locks: dict[str, asyncio.Lock] = {}
+ 
+@@ -160,6 +170,7 @@ class DockerEnvironment(BaseEnvironment):
+         task_env_config: EnvironmentConfig,
+         keep_containers: bool = False,
+         mounts_json: list[ServiceVolumeConfig] | None = None,
++        mounts_override: list[ServiceVolumeConfig] | None = None,
+         *args,
+         **kwargs,
+     ):
+@@ -193,8 +204,13 @@ class DockerEnvironment(BaseEnvironment):
+             self._windows_container_name: str | None = None
+             self._platform = UnixOps(self)
+ 
++        # Configured mounts are additive so the default /logs binds survive and
++        # agent logs and artifacts stay host-visible. A separate verifier
++        # environment must not share those directories, so it overrides the set.
+         self._mounts_json = (
+-            mounts_json if mounts_json is not None else self._default_log_mounts()
++            [*self._default_log_mounts(), *(mounts_json or [])]
++            if mounts_override is None
++            else mounts_override
+         )
+         self._mounts_compose_path: Path | None = None
+         self._resources_compose_temp_dir: tempfile.TemporaryDirectory | None = None
+@@ -427,12 +443,16 @@ class DockerEnvironment(BaseEnvironment):
+             Path(self._resources_compose_temp_dir.name)
+             / f"{self.session_id}-{RESOURCES_COMPOSE_NAME}"
+         )
++        cpu_limit = self._resource_limit_value("cpu", auto_mode=ResourceMode.LIMIT)
+         return write_resources_compose_file(
+             path,
++            environment=parallelism_env(cpu_limit),
++            volumes=self._write_cpu_clamp(path.parent, cpu_limit),
++            sysctls=None if self._is_windows_container else self._LINUX_SYSCTLS,
+             cpu_request=self._resource_request_value(
+                 "cpu", auto_mode=ResourceMode.LIMIT
+             ),
+-            cpu_limit=self._resource_limit_value("cpu", auto_mode=ResourceMode.LIMIT),
++            cpu_limit=cpu_limit,
+             memory_request_mb=self._resource_request_value(
+                 "memory", auto_mode=ResourceMode.LIMIT
+             ),
+@@ -441,6 +461,15 @@ class DockerEnvironment(BaseEnvironment):
+             ),
+         )
+ 
++    def _write_cpu_clamp(self, directory: Path, cpu_limit: int | None) -> list[str]:
++        """Mount the preload that `parallelism_env` points NODE_OPTIONS at."""
++        if not cpu_limit or self._is_windows_container:
++            return []
++        source = directory / "pier-node-cpu-clamp.js"
++        source.write_text(cpu_clamp_source(cpu_limit))
++        source.chmod(0o644)
++        return [f"{source}:{CPU_CLAMP_PATH}:ro"]
++
+     def _cleanup_resources_compose_file(self) -> None:
+         if self._resources_compose_temp_dir is None:
+             return
+diff --git a/src/pier/environments/docker/parallelism.py b/src/pier/environments/docker/parallelism.py
+new file mode 100644
+index 0000000..b21ae7c
+--- /dev/null
++++ b/src/pier/environments/docker/parallelism.py
+@@ -0,0 +1,34 @@
++"""Tell test runners how many CPUs the container actually gets.
++
++Docker's ``--cpus`` is a bandwidth quota, so ``nproc`` still reports every host
++core and runner worker pools oversubscribe the container.
++"""
++
++from __future__ import annotations
++
++CPU_CLAMP_PATH = "/opt/pier-node-cpu-clamp.js"
++
++
++def cpu_clamp_source(cpus: int) -> str:
++    """A ``node --require`` preload; the Node runners take no cap from the environment."""
++    return (
++        "const os = require('node:os')\n"
++        f"const limit = {cpus}\n"
++        "const cpus = os.cpus\n"
++        "os.cpus = () => cpus.call(os).slice(0, limit)\n"
++        "os.availableParallelism = () => limit\n"
++    )
++
++
++def parallelism_env(cpus: int | None) -> dict[str, str]:
++    """Worker-count caps for the test runners, derived from the CPU limit."""
++    if not cpus or cpus < 1:
++        return {}
++    n = str(cpus)
++    return {
++        "GOMAXPROCS": n,
++        "CARGO_BUILD_JOBS": n,
++        "NEXTEST_TEST_THREADS": n,
++        "PYTEST_XDIST_AUTO_NUM_WORKERS": n,
++        "NODE_OPTIONS": f"--require {CPU_CLAMP_PATH}",
++    }
+diff --git a/src/pier/models/agent/name.py b/src/pier/models/agent/name.py
+index d84516a..41ac147 100644
+--- a/src/pier/models/agent/name.py
++++ b/src/pier/models/agent/name.py
+@@ -8,6 +8,7 @@ class AgentName(str, Enum):
+     ANTIGRAVITY_SDK = "antigravity-sdk"
+     CODEX = "codex"
+     CURSOR_CLI = "cursor-cli"
++    DSH_MINIMAL = "dsh-minimal"
+     GEMINI_CLI = "gemini-cli"
+     MINI_SWE_AGENT = "mini-swe-agent"
+     SWE_AGENT = "swe-agent"
+diff --git a/src/pier/trial/trial.py b/src/pier/trial/trial.py
+index fd5d4a2..126f0fa 100644
+--- a/src/pier/trial/trial.py
++++ b/src/pier/trial/trial.py
+@@ -402,7 +402,7 @@ class Trial:
+             trial_paths=self._trial_paths,
+             task_env_config=env_config,
+             logger=self._logger,
+-            mounts_json=self._verifier_env_mounts(env_config),
++            mounts_override=self._verifier_env_mounts(env_config),
+             agent_install_spec=None,
+             network_allowlist=None,
+             default_user=(

inference/README.md

@@ -0,0 +1,71 @@
+# Minimal inference
+
+A readable reference implementation rather than a production serving engine. The
+model code covers the vision encoder and aligner, sliding-window plus compressed
+sparse attention with its two-level indexer, engram n-gram lookups, MoE,
+Hyper-Connections, and the DSpark forward path. Generation itself is plain
+autoregressive sampling.
+
+## Install
+
+```bash
+python -m pip install -r requirements.txt
+```
+
+## Convert Hugging Face weights
+
+The runtime uses one converted checkpoint file per tensor-parallel rank. From
+this directory:
+
+```bash
+export HF_CKPT_PATH=/path/to/DeepSeek-V4.1-Flash-HF
+export SAVE_PATH=/path/to/DeepSeek-V4.1-Flash-TP8
+export MP=8
+
+python convert.py \
+  --hf-ckpt-path "${HF_CKPT_PATH}" \
+  --save-path "${SAVE_PATH}" \
+  --model-parallel "${MP}" \
+  --expert-dtype fp4 \
+  --tokenizer-path "${HF_CKPT_PATH}"
+```
+
+Expert counts are inferred from the weight names, so they do not need to be
+passed. `--tokenizer-path` points at whichever directory holds `tokenizer.json`
+and `tokenizer_config.json`; they are copied into the converted checkpoint.
+
+## Run the equivalent TXT and JSON examples
+
+```bash
+export CKPT_PATH=/path/to/DeepSeek-V4.1-Flash-TP8
+export MP=8
+
+INPUT_FILE=examples/example.txt ./run.sh
+INPUT_FILE=examples/example_harmony.json ./run.sh
+```
+
+The two files express the same interleaved two-image prompt, so they produce
+identical encoded prompts and input token IDs.
+
+For interactive chat:
+
+```bash
+torchrun --nproc-per-node "${MP}" generate.py \
+  --ckpt-path "${CKPT_PATH}" \
+  --config config.json \
+  --interactive \
+  --temperature 0.6
+```
+
+For multi-node execution, pass the usual `torchrun --nnodes`, `--node-rank`,
+`--master-addr`, and `--master-port` arguments before `generate.py`.
+
+## Self-test
+
+`model.py` builds a small model from the `ModelArgs` defaults and runs a prefill
+plus 22 decode steps, exercising the real dense-fp8 / MoE-fp4 kernels. Weights
+are uninitialized, so it checks shapes and kernel plumbing, not numerics:
+
+```bash
+python model.py
+```

inference/config.json

@@ -0,0 +1,66 @@
+{
+    "vocab_size": 129280,
+    "dim": 5120,
+    "moe_inter_dim": 2304,
+    "n_layers": 40,
+    "n_mtp_layers": 3,
+    "dspark_block_size": 5,
+    "dspark_noise_token_id": 128799,
+    "dspark_target_layer_ids": [37, 38, 39],
+    "dspark_markov_rank": 256,
+    "dspark_n_routed_experts": 128,
+    "dspark_n_activated_experts": 3,
+    "n_heads": 64,
+    "n_routed_experts": 384,
+    "n_shared_experts": 1,
+    "n_activated_experts": 6,
+    "score_func": "sqrtsoftplus",
+    "route_scale": 1.5,
+    "swiglu_limit": 10.0,
+    "q_lora_rank": 1280,
+    "head_dim": 512,
+    "rope_head_dim": 64,
+    "norm_eps": 1e-20,
+    "o_groups": 8,
+    "o_lora_rank": 1024,
+    "window_size": 128,
+    "kv_source_layers": [2, 8, 14, 20],
+    "index_source_layers": [2, 8, 14, 20, 24, 28, 32, 36],
+    "original_seq_len": 65536,
+    "rope_theta": 10000,
+    "rope_factor": 16,
+    "beta_fast": 32,
+    "beta_slow": 1,
+    "index_n_heads": 32,
+    "index_head_dim": 128,
+    "index_topk": 512,
+    "candidate_source_layer": 20,
+    "candidate_topk_blocks": 2048,
+    "candidate_block_size": 8,
+    "hc_mult": 4,
+    "hc_sinkhorn_iters": 20,
+    "hc_eps": 1e-06,
+    "engram_layer_ids": [1, 14],
+    "engram_vocab_size": 16000000,
+    "engram_num_embeddings": [384006168, 384016682],
+    "engram_max_ngram_size": 4,
+    "engram_pad_id": 2,
+    "engram_compressed_vocab_size": 99092,
+    "dtype": "fp8",
+    "expert_dtype": "fp4",
+    "compress_rope_theta": 160000,
+    "compress_ratios": [0, 0, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
+    "vision_n_layers": 32,
+    "vision_dim": 1024,
+    "vision_n_heads": 16,
+    "vision_inter_dim": 2816,
+    "vision_patch_size": 14,
+    "vision_downsample_ratio": 3,
+    "vision_max_n_token": 1024,
+    "vision_min_pixels": 295936,
+    "vision_max_wh_ratio": null,
+    "image_token_id": 129264,
+    "engram_n_heads": 8,
+    "engram_head_dim": 256,
+    "vision_rope_theta": 10000
+}

inference/convert.py

@@ -0,0 +1,205 @@
+import json
+import os
+import re
+import shutil
+from argparse import ArgumentParser
+from glob import glob
+from tqdm import tqdm, trange
+
+import torch
+from safetensors.torch import safe_open, save_file
+
+
+FP4_TABLE = torch.tensor(
+    [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], dtype=torch.float32
+)
+
+
+def cast_e2m1fn_to_e4m3fn(x: torch.Tensor, scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Casts a tensor from e2m1fn to e4m3fn losslessly.
+    """
+    assert x.dtype == torch.int8
+    assert x.ndim == 2
+    out_dim, in_dim = x.size()
+    in_dim *= 2
+    fp8_block_size = 32
+    fp4_block_size = 32
+    assert in_dim % fp8_block_size == 0 and out_dim % fp8_block_size == 0
+    assert scale.size(0) == out_dim and scale.size(1) == in_dim // fp4_block_size
+
+    x = x.view(torch.uint8)
+    low = x & 0x0F
+    high = (x >> 4) & 0x0F
+    x = torch.stack([FP4_TABLE[low.long()], FP4_TABLE[high.long()]], dim=-1).flatten(2)
+
+    # max_fp4 (6.0) * MAX_OFFSET must fit in e4m3fn (max 448)
+    # 6.0 * 2^6 = 384 < 448; 6.0 * 2^7 = 768 > 448; so MAX_OFFSET_BITS = 6
+    MAX_OFFSET_BITS = 6
+
+    bOut = out_dim // fp8_block_size
+    bIn = in_dim // fp8_block_size
+    # bOut, bIn, fp8_block_size, fp8_block_size
+    x = x.view(bOut, fp8_block_size, bIn, fp8_block_size).transpose(1, 2)
+    # bOut, bIn, fp8_block_size * (fp8_block_size // fp4_block_size)
+    scale = scale.float().view(bOut, fp8_block_size, bIn, -1).transpose(1, 2).flatten(2)
+    ## bOut, bIn, 1
+    scale_max_offset_bits = scale.amax(dim=-1, keepdim=True) / (2**MAX_OFFSET_BITS)
+    # bOut, bIn, fp8_block_size * (fp8_block_size // fp4_block_size)
+    offset = scale / scale_max_offset_bits
+    # bOut, bIn, fp8_block_size, fp8_block_size
+    offset = offset.unflatten(-1, (fp8_block_size, -1)).repeat_interleave(fp4_block_size, dim=-1)
+    x = (x * offset).transpose(1, 2).reshape(out_dim, in_dim)
+    return x.to(torch.float8_e4m3fn), scale_max_offset_bits.squeeze(-1).to(torch.float8_e8m0fnu)
+
+
+mapping = {
+    "embed": ("embed", 0),
+    "wq_b": ("wq_b", 0),
+    "wo_a": ("wo_a", 0),
+    "wo_b": ("wo_b", 1),
+    "head": ("head", 0),
+    "attn_sink": ("attn_sink", 0),
+    "weights_proj": ("weights_proj", 0),
+}
+
+
+def infer_num_experts(names) -> tuple[int, int]:
+    """Number of routed experts in the backbone and in the MTP layers, from the weight names."""
+    counts = [0, 0]
+    for name in names:
+        name = name.removeprefix("model.")
+        match = re.search(r"(?:mlp|ffn)\.experts\.(\d+)\.", name)
+        if match:
+            is_mtp = name.startswith("mtp.")
+            counts[is_mtp] = max(counts[is_mtp], int(match.group(1)) + 1)
+    assert counts[0], "no routed experts found in the checkpoint"
+    return counts[0], counts[1] or counts[0]
+
+
+def main(hf_ckpt_path, save_path, mp, expert_dtype, tokenizer_path=None):
+    """Shard an exported HuggingFace checkpoint into `mp` files for this inference stack."""
+    torch.set_num_threads(8)
+    state_dicts = [{} for _ in range(mp)]
+    os.makedirs(save_path, exist_ok=True)
+
+    index_path = os.path.join(hf_ckpt_path, "model.safetensors.index.json")
+    expected_names = set(json.load(open(index_path))["weight_map"]) if os.path.exists(index_path) else None
+    seen_names = set()
+
+    all_names = expected_names
+    if all_names is None:
+        all_names = set()
+        for file_path in glob(os.path.join(hf_ckpt_path, "*.safetensors")):
+            with safe_open(file_path, framework="pt", device="cpu") as f:
+                all_names.update(f.keys())
+    n_experts, mtp_n_experts = infer_num_experts(all_names)
+    assert n_experts % mp == 0 and mtp_n_experts % mp == 0, (n_experts, mtp_n_experts, mp)
+    print(f"{n_experts=} {mtp_n_experts=}")
+
+    for file_path in tqdm(glob(os.path.join(hf_ckpt_path, "*.safetensors"))):
+        with safe_open(file_path, framework="pt", device="cpu") as f:
+            for source_name in f.keys():
+                seen_names.add(source_name)
+                name = source_name
+                if name.startswith("model."):
+                    name = name[len("model.") :]
+                param: torch.Tensor = f.get_tensor(source_name)
+                # an MTP layer ties its token embedding and output head to the backbone's
+                if name.startswith("mtp.") and name.split(".", 2)[-1] in ("embed.weight", "head.weight"):
+                    continue
+                name = name.replace("self_attn", "attn")
+                if not name.startswith("vision."):
+                    name = name.replace("mlp", "ffn")
+                name = name.replace("weight_scale_inv", "scale")
+                name = name.replace("e_score_correction_bias", "bias")
+                if any(
+                    x in name for x in ["hc", "attn_sink", "tie2eid", "tid2eid", "ape", "image_"]
+                ):  # without .weight
+                    key = name.split(".")[-1]
+                else:
+                    key = name.split(".")[-2]
+                if key in mapping:
+                    new_key, dim = mapping[key]
+                else:
+                    new_key, dim = key, None
+                name = name.replace(key, new_key)
+                for i in range(mp):
+                    new_param = param
+                    if "experts" in name and "shared_experts" not in name:
+                        current_n_experts = mtp_n_experts if name.startswith("mtp.") else n_experts
+                        n_local_experts = current_n_experts // mp
+                        idx = int(name.split(".")[-3])
+                        if idx < i * n_local_experts or idx >= (i + 1) * n_local_experts:
+                            continue
+                    elif ".engram.embed." in name:
+                        shard_size = (param.size(0) + mp - 1) // mp
+                        new_param = param[i * shard_size : (i + 1) * shard_size].contiguous()
+                        if new_param.size(0) < shard_size:
+                            pad_value = 1 if name.endswith(".scale") else 0
+                            padding = param.new_full((shard_size - new_param.size(0), param.size(1)), pad_value)
+                            new_param = torch.cat([new_param, padding])
+                    elif dim is not None:
+                        assert param.size(dim) % mp == 0, f"Dimension {dim} must be divisible by {mp}"
+                        shard_size = param.size(dim) // mp
+                        new_param = param.narrow(dim, i * shard_size, shard_size).contiguous()
+                    state_dicts[i][name] = new_param
+
+    if expected_names is not None:
+        assert seen_names == expected_names, (
+            f"checkpoint shards incomplete: {len(expected_names - seen_names)} tensors missing, "
+            f"{len(seen_names - expected_names)} unexpected (source may be mid-upload)"
+        )
+
+    for i in trange(mp):
+        names = list(state_dicts[i].keys())
+        for name in names:
+            if name.endswith("wo_a.weight"):
+                weight = state_dicts[i][name]
+                scale = state_dicts[i].pop(name.replace("weight", "scale"))
+                assert weight.size(0) % scale.size(0) == 0
+                assert weight.size(1) % scale.size(1) == 0
+                out_block_size = weight.size(0) // scale.size(0)
+                in_block_size = weight.size(1) // scale.size(1)
+                assert (out_block_size, in_block_size) in ((32, 32), (128, 128)), (
+                    name,
+                    weight.shape,
+                    scale.shape,
+                )
+                weight = (
+                    weight.unflatten(0, (-1, out_block_size)).unflatten(-1, (-1, in_block_size)).float()
+                    * scale[:, None, :, None].float()
+                )
+                state_dicts[i][name] = weight.flatten(2, 3).flatten(0, 1).bfloat16()
+            elif "experts" in name and state_dicts[i][name].dtype == torch.int8:
+                if expert_dtype == "fp8":
+                    scale_name = name.replace("weight", "scale")
+                    weight = state_dicts[i].pop(name)
+                    scale = state_dicts[i].pop(scale_name)
+                    state_dicts[i][name], state_dicts[i][scale_name] = cast_e2m1fn_to_e4m3fn(weight, scale)
+                else:
+                    state_dicts[i][name] = state_dicts[i][name].view(torch.float4_e2m1fn_x2)
+        save_file(state_dicts[i], os.path.join(save_path, f"model{i}-mp{mp}.safetensors"))
+
+    tokenizer_path = tokenizer_path or hf_ckpt_path
+    for file in ["tokenizer.json", "tokenizer_config.json"]:
+        old_file_path = os.path.join(tokenizer_path, file)
+        new_file_path = os.path.join(save_path, file)
+        if os.path.exists(old_file_path):
+            shutil.copyfile(old_file_path, new_file_path)
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--hf-ckpt-path", type=str, required=True)
+    parser.add_argument("--save-path", type=str, required=True)
+    parser.add_argument("--model-parallel", type=int, required=True)
+    parser.add_argument("--expert-dtype", type=str, choices=["fp8", "fp4"], default=None)
+    parser.add_argument(
+        "--tokenizer-path",
+        type=str,
+        default=None,
+        help="Optional tokenizer directory when the HF checkpoint does not contain tokenizer files",
+    )
+    args = parser.parse_args()
+    main(args.hf_ckpt_path, args.save_path, args.model_parallel, args.expert_dtype, args.tokenizer_path)

inference/engram.py

@@ -0,0 +1,184 @@
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+from sympy import isprime
+from torch import nn
+
+
+def find_next_prime(start: int, seen_primes: set[int]) -> int:
+    """The smallest prime above `start` that has not been handed out yet."""
+    candidate = start + 1
+    while not isprime(candidate) or candidate in seen_primes:
+        candidate += 1
+    return candidate
+
+
+def build_compressed_token_map(tokenizer) -> tuple[list[int], int]:
+    """Map every token id onto a smaller id space where tokens that normalize alike collapse together.
+
+    N-grams are hashed over these compressed ids, so " The", "the" and "THE" all hash the same way.
+    Returns the lookup plus the size of the compressed vocab -- and that size matters beyond bounds
+    checking, because every hash multiplier is derived from it.
+    """
+    from tokenizers import Regex, normalizers
+
+    # a private-use char, so a token that is exactly one space survives Strip() instead of
+    # collapsing to the empty string and merging with unrelated tokens
+    sentinel = "\ue000"
+    normalizer = normalizers.Sequence(
+        [
+            normalizers.NFKC(),
+            normalizers.NFD(),
+            normalizers.StripAccents(),
+            normalizers.Lowercase(),
+            normalizers.Replace(Regex(r"[ \t\r\n]+"), " "),
+            normalizers.Replace(Regex(r"^ $"), sentinel),
+            normalizers.Strip(),
+            normalizers.Replace(sentinel, " "),
+        ]
+    )
+
+    # the raw Rust tokenizer, matching what training decodes with (no clean_up_tokenization_spaces)
+    backend = tokenizer.backend_tokenizer
+    key_to_new: dict[str, int] = {}
+    lookup = [0] * len(tokenizer)
+    for token_id in range(len(tokenizer)):
+        text = backend.decode([token_id], skip_special_tokens=False)
+        if "\ufffd" in text:
+            # a partial UTF-8 byte token: nothing to normalize, so key it by its raw form
+            key = backend.id_to_token(token_id)
+        else:
+            normalized = normalizer.normalize_str(text)
+            key = normalized if normalized else text
+
+        new_id = key_to_new.get(key)
+        if new_id is None:
+            new_id = len(key_to_new)
+            key_to_new[key] = new_id
+        lookup[token_id] = new_id
+
+    return lookup, len(key_to_new)
+
+
+def compute_hash_multipliers(
+    layer_ids: tuple[int, ...], max_ngram_size: int, tokenizer_vocab_size: int
+) -> torch.Tensor:
+    """One multiplier per (layer, lookback), from a per-layer RNG so layers hash differently.
+
+    Kept odd, and bounded so that `token_id * multiplier` cannot overflow int64.
+    """
+    max_long = np.iinfo(np.int64).max
+    multiplier_bound = max(1, (max_long // tokenizer_vocab_size) // 2)
+    rows = []
+    for layer_id in layer_ids:
+        generator = np.random.default_rng(10007 * layer_id)
+        values = generator.integers(
+            low=0,
+            high=multiplier_bound,
+            size=(max_ngram_size,),
+            dtype=np.int64,
+        )
+        rows.append(torch.tensor(values * 2 + 1))
+    return torch.stack(rows)
+
+
+@dataclass(frozen=True)
+class EngramLayout:
+    """Bucket layout of the n-gram hash tables.
+
+    A position is hashed as `max_ngram_size - 1` n-grams (2-gram .. max_ngram_size-gram), each split
+    over `n_heads` heads. Every (n-gram size, head) pair owns its own prime-sized bucket range in the
+    layer's table; the primes are drawn in order and never reused, which keeps the ranges disjoint.
+    """
+
+    max_ngram_size: int
+    layer_ids: tuple[int, ...]
+    num_embeddings: tuple[int, ...]  # table rows, per engram layer
+    primes: tuple[tuple[tuple[int, ...], ...], ...]  # [layer][n-gram size][head] bucket modulus
+    n_heads: int
+    head_dim: int
+
+    @classmethod
+    def from_args(cls, args) -> "EngramLayout | None":
+        layer_ids = tuple(args.engram_layer_ids)
+        if not layer_ids:
+            return None
+        max_ngram_size, n_heads = args.engram_max_ngram_size, args.engram_n_heads
+        primes, seen = [], set()
+        for _ in layer_ids:
+            per_ngram = []
+            for _ in range(max_ngram_size - 1):
+                sizes, current = [], args.engram_vocab_size - 1
+                for _ in range(n_heads):
+                    current = find_next_prime(current, seen)
+                    seen.add(current)
+                    sizes.append(current)
+                per_ngram.append(tuple(sizes))
+            primes.append(tuple(per_ngram))
+        return cls(
+            max_ngram_size=max_ngram_size,
+            layer_ids=layer_ids,
+            num_embeddings=tuple(args.engram_num_embeddings),
+            primes=tuple(primes),
+            n_heads=n_heads,
+            head_dim=args.engram_head_dim,
+        )
+
+
+class NgramHashState(nn.Module):
+    """Maps each position to the hash ids of the n-grams ending there.
+
+    Ids go through the compressed table, then each position is hashed with the `max_ngram_size - 1`
+    tokens before it. Look-back stops at the start of the sequence and at any dead token (an image
+    span, cached as DEAD), so an n-gram never spans one. The cache carries all of this across the
+    prefill/decode split.
+    """
+
+    DEAD = -1
+
+    def __init__(self, args, layout: EngramLayout, tokenizer):
+        super().__init__()
+        self.layout = layout
+        # every hash multiplier derives from the compressed vocab size, so a mismatch there would
+        # silently rehash the whole table
+        token_map, vocab_size = build_compressed_token_map(tokenizer)
+        assert vocab_size == args.engram_compressed_vocab_size, (vocab_size, args.engram_compressed_vocab_size)
+        self.pad_id = token_map[args.engram_pad_id]
+        flat = [[p for per_ngram in layer for p in per_ngram] for layer in layout.primes]
+        offsets = [np.cumsum([0, *sizes[:-1]]) for sizes in flat]
+        multipliers = compute_hash_multipliers(layout.layer_ids, layout.max_ngram_size, vocab_size)
+        self.register_buffer("primes", torch.tensor(layout.primes), persistent=False)
+        self.register_buffer("offsets", torch.tensor(np.array(offsets)), persistent=False)
+        self.register_buffer("multipliers", multipliers, persistent=False)
+        self.register_buffer("token_map", torch.tensor(token_map), persistent=False)
+        self.register_buffer(
+            "cache", torch.empty(args.max_batch_size, args.max_seq_len, dtype=torch.int64), persistent=False
+        )
+
+    @torch.inference_mode()
+    def forward(self, input_ids: torch.Tensor, start_pos: int, token_mask: torch.Tensor | None = None) -> torch.Tensor:
+        """token_mask: [B, L], False for tokens that take no part in an n-gram (image spans).
+        Returns the hash ids, shaped [B, L, n_engram_layers, n_hash_cols]."""
+        batch, seqlen = input_ids.shape
+        compressed = self.token_map[input_ids]
+        if token_mask is not None:
+            compressed = torch.where(token_mask, compressed, self.DEAD)
+        self.cache[:batch, start_pos : start_pos + seqlen] = compressed
+
+        positions = torch.arange(start_pos, start_pos + seqlen, device=input_ids.device).expand(batch, seqlen)
+        tokens, blocked = [], torch.zeros_like(positions, dtype=torch.bool)
+        for shift in range(self.layout.max_ngram_size):
+            source = self.cache[:batch].gather(1, (positions - shift).clamp_min(0))
+            blocked = blocked | (positions < shift) | (source == self.DEAD)
+            tokens.append(torch.where(blocked, self.pad_id, source))
+        tokens = torch.stack(tokens, dim=-1)  # [B, L, max_ngram_size]
+
+        # XOR the multiplied ids together one lookback at a time, so the running value after step i
+        # is the hash of the (i+1)-gram; each lands in its own prime-sized bucket range
+        products = tokens.unsqueeze(2) * self.multipliers  # [B, L, n_engram_layers, max_ngram_size]
+        rolling, hashes = products[..., 0], []
+        for i in range(1, self.layout.max_ngram_size):
+            rolling = torch.bitwise_xor(rolling, products[..., i])
+            hashes.append(rolling.unsqueeze(-1) % self.primes[:, i - 1])
+        return torch.cat(hashes, dim=-1) + self.offsets

inference/examples/example.txt

@@ -0,0 +1,7 @@
+中国的首都是哪里？
+
+列出100以内的所有素数。
+
+DeepSeek是做什么的公司？
+
+请按“第一张、第二张”的顺序回答：第一张图<image>examples/images/carrots.jpeg</image>和第二张图<image>examples/images/corn.jpeg</image>中分别是什么食材？它们通常食用的部位分别是什么？

inference/examples/example_harmony.json

@@ -0,0 +1,96 @@
+[
+  {
+    "messages": [
+      {
+        "role": "user",
+        "content": [
+          {
+            "type": "text",
+            "text": "请按“第一张、第二张”的顺序回答：第一张图"
+          },
+          {
+            "type": "image_url",
+            "image_url": {
+              "url": "examples/images/carrots.jpeg"
+            }
+          },
+          {
+            "type": "text",
+            "text": "和第二张图"
+          },
+          {
+            "type": "image_url",
+            "image_url": {
+              "url": "examples/images/corn.jpeg"
+            }
+          },
+          {
+            "type": "text",
+            "text": "中分别是什么食材？它们通常食用的部位分别是什么？"
+          }
+        ]
+      }
+    ]
+  },
+  {
+    "messages": [
+      {
+        "role": "system",
+        "content": "You are a helpful assistant."
+      },
+      {
+        "role": "user",
+        "content": "中国的首都是哪里？"
+      }
+    ]
+  },
+  {
+    "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"]}
+            },
+            "required": ["location"]
+          }
+        }
+      }
+    ],
+    "messages": [
+      {
+        "role": "system",
+        "content": "You are a helpful assistant."
+      },
+      {
+        "role": "user",
+        "content": "What's the weather like in Beijing?"
+      }
+    ]
+  },
+  {
+    "messages": [
+      {
+        "role": "system",
+        "content": "You are a helpful assistant."
+      },
+      {
+        "role": "user",
+        "content": "Hello"
+      },
+      {
+        "role": "assistant",
+        "content": "Hi there! How can I help you?"
+      },
+      {
+        "role": "system",
+        "content": "Mid-conversation instruction update: reply in Chinese only. (deepseek_v41 only)"
+      }
+    ]
+  }
+]

inference/generate.py

@@ -0,0 +1,218 @@
+import os
+import json
+import sys
+from argparse import ArgumentParser
+from typing import List
+
+import torch
+import torch.distributed as dist
+from transformers import AutoTokenizer
+from safetensors.torch import load_model
+
+from model import Transformer, ModelArgs
+
+current_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, os.path.abspath(os.path.join(current_dir, "../encoding")))
+
+from encoding import (
+    encode_case,
+    encode_messages,
+    load_cases,
+    parse_message_from_completion_text,
+    parse_tagged_text,
+    to_json,
+)
+from image_processor import TEXT, prepare_vl_inputs
+
+
+@torch.inference_mode()
+def generate(
+    model: Transformer,
+    prompt_tokens: List[List[int]],
+    max_new_tokens: int,
+    eos_id: int,
+    prompt_token_types: List[List[int]] | None = None,
+    images=None,
+) -> List[List[int]]:
+    """Batch generation with right-padded prompts.
+
+    The first forward pass processes [:min_prompt_len] tokens (prefill phase).
+    Subsequent passes generate one token at a time (decode phase). For positions
+    still within a prompt, the ground-truth token overrides the model's prediction.
+
+    `prompt_token_types` and `images` come from image_processor.prepare_vl_inputs. Image spans are
+    only visible to the prefill pass, so they must end before the shortest prompt does.
+    """
+    prompt_lens = [len(t) for t in prompt_tokens]
+    assert max(prompt_lens) <= model.max_seq_len, (
+        f"Prompt length exceeds model maximum sequence length (max_seq_len={model.max_seq_len})"
+    )
+    total_len = min(model.max_seq_len, max_new_tokens + max(prompt_lens))
+    tokens = torch.full((len(prompt_tokens), total_len), -1, dtype=torch.long)
+    for i, t in enumerate(prompt_tokens):
+        tokens[i, : len(t)] = torch.tensor(t, dtype=torch.long)
+
+    token_types = None
+    if images is not None:
+        token_types = torch.full((len(prompt_tokens), total_len), TEXT, dtype=torch.long)
+        for i, types in enumerate(prompt_token_types):
+            token_types[i, : len(types)] = torch.tensor(types, dtype=torch.long)
+        for sample in images:
+            for img in sample or ():
+                assert img.start + img.types.numel() <= min(prompt_lens), "image spans must fit in the prefill chunk"
+
+    prev_pos = 0
+    finished = torch.tensor([False] * len(prompt_tokens))
+    prompt_mask = tokens != -1
+    for cur_pos in range(min(prompt_lens), total_len):
+        with_images = images is not None and prev_pos == 0
+        next_token = model.forward(
+            tokens[:, prev_pos:cur_pos],
+            prev_pos,
+            images=images if with_images else None,
+            token_types=token_types[:, prev_pos:cur_pos] if with_images else None,
+        )[0]
+        next_token = torch.where(prompt_mask[:, cur_pos], tokens[:, cur_pos], next_token)
+        tokens[:, cur_pos] = next_token
+        finished |= torch.logical_and(~prompt_mask[:, cur_pos], next_token == eos_id)
+        prev_pos = cur_pos
+        if finished.all():
+            break
+    completion_tokens = []
+    for i, toks in enumerate(tokens.tolist()):
+        toks = toks[prompt_lens[i] : prompt_lens[i] + max_new_tokens]
+        if eos_id in toks:
+            toks = toks[: toks.index(eos_id)]
+        completion_tokens.append(toks)
+    return completion_tokens
+
+
+def prepare_case(case, thinking_mode, tokenizer, args):
+    """Encode one message case and expand any image placeholders."""
+    if case.get("context"):
+        raise ValueError("Standalone inference does not support context without a prefilled KV cache")
+    prompt, image_records = encode_case(case, thinking_mode)
+    tokens, token_types, images = prepare_vl_inputs(prompt, image_records, tokenizer, args)
+    return prompt, tokens, token_types, images
+
+
+def main(
+    ckpt_path: str,
+    config: str,
+    input_file: str = "",
+    interactive: bool = True,
+    max_new_tokens: int = 100,
+    temperature: float = 1.0,
+    thinking_mode: str = "chat",
+) -> None:
+    world_size = int(os.getenv("WORLD_SIZE", "1"))
+    rank = int(os.getenv("RANK", "0"))
+    local_rank = int(os.getenv("LOCAL_RANK", "0"))
+    if world_size > 1:
+        dist.init_process_group("nccl")
+    global print
+    if rank != 0:
+        print = lambda *_, **__: None
+    torch.cuda.set_device(local_rank)
+    torch.cuda.memory._set_allocator_settings("expandable_segments:True")
+    torch.set_default_dtype(torch.bfloat16)
+    torch.set_num_threads(8)
+    torch.manual_seed(33377335)
+    with open(config) as f:
+        args = ModelArgs(**json.load(f))
+        args.temperature = temperature
+    if interactive:
+        args.max_batch_size = 1
+        args.max_seq_len = 64 * 1024
+    print(args)
+    tokenizer = AutoTokenizer.from_pretrained(ckpt_path)
+    print("build model")
+    with torch.device("cuda"):
+        model = Transformer(args, tokenizer)
+    print("load model")
+    load_model(model, os.path.join(ckpt_path, f"model{rank}-mp{world_size}.safetensors"))
+    torch.set_default_device("cuda")
+    print("I'm DeepSeek 👋")
+
+    if interactive:
+        messages = []
+        while True:
+            if world_size == 1:
+                prompt = input(">>> ")
+            elif rank == 0:
+                prompt = input(">>> ")
+                objects = [prompt]
+                dist.broadcast_object_list(objects, 0)
+            else:
+                objects = [None]
+                dist.broadcast_object_list(objects, 0)
+                prompt = objects[0]
+            if prompt == "/exit":
+                break
+            elif prompt == "/clear":
+                messages.clear()
+                continue
+            messages.append({"role": "user", "content": prompt})
+            prompt_tokens = tokenizer.encode(encode_messages(messages, thinking_mode=thinking_mode))
+            completion_tokens = generate(model, [prompt_tokens], max_new_tokens, tokenizer.eos_token_id)
+            completion = tokenizer.decode(completion_tokens[0])
+            print(completion)
+            messages.append(parse_message_from_completion_text(completion, thinking_mode=thinking_mode))
+    else:
+        if input_file.endswith(".json"):
+            # Harmony input: a JSON file with one or more OpenAI-format cases
+            # ({"messages": [...], "tools": [...]} or a bare message list).
+            cases = load_cases(input_file)
+            raw_prompts = [to_json(case["messages"]) for case in cases]
+        else:
+            # Plain-text input: blank-line-separated prompts, optionally with
+            # <image>path</image> tags.
+            with open(input_file) as f:
+                raw_prompts = f.read().rstrip("\n").split("\n\n")
+            cases = [{"messages": [{"role": "user", "content": parse_tagged_text(prompt)}]} for prompt in raw_prompts]
+
+        prompt_tokens, prompt_token_types, images = [], [], []
+        for case in cases:
+            _, tokens, token_types, image_inputs = prepare_case(case, thinking_mode, tokenizer, args)
+            prompt_tokens.append(tokens)
+            prompt_token_types.append(token_types)
+            images.append(image_inputs)
+
+        if any(images):
+            # image spans must be prefilled in one chunk, so VL prompts are generated one at a time
+            completion_tokens = [
+                generate(model, [tok], max_new_tokens, tokenizer.eos_token_id, [types], [image])[0]
+                for tok, types, image in zip(prompt_tokens, prompt_token_types, images)
+            ]
+        else:
+            completion_tokens = generate(model, prompt_tokens, max_new_tokens, tokenizer.eos_token_id)
+        completions = tokenizer.batch_decode(completion_tokens)
+        for raw_prompt, completion in zip(raw_prompts, completions):
+            print("Prompt:", raw_prompt)
+            print("Completion:", completion)
+            print()
+
+    if world_size > 1:
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--ckpt-path", type=str, required=True)
+    parser.add_argument("--config", type=str, required=True)
+    parser.add_argument("--input-file", type=str, default="")
+    parser.add_argument("--interactive", action="store_true")
+    parser.add_argument("--max-new-tokens", type=int, default=200)
+    parser.add_argument("--temperature", type=float, default=1.0)
+    parser.add_argument("--thinking-mode", type=str, default="chat", choices=["chat", "thinking"])
+    args = parser.parse_args()
+    assert args.input_file or args.interactive, "Either input-file or interactive mode must be specified"
+    main(
+        args.ckpt_path,
+        args.config,
+        args.input_file,
+        args.interactive,
+        args.max_new_tokens,
+        args.temperature,
+        args.thinking_mode,
+    )

inference/image_processor.py

@@ -0,0 +1,173 @@
+"""Image preprocessing.
+
+An image becomes a `n_vit_h x n_vit_w` patch grid for the ViT and a `n_llm_h x n_llm_w` token grid
+after the 3x3 aligner downsample, which the LLM sees as
+
+    [IMAGE_START] + ([IMAGE] * n_llm_w + [IMAGE_NEW_LINE]) * n_llm_h + [IMAGE_END]
+
+Every one of those positions carries `image_token_id` in `input_ids`; only the token type tells them
+apart. The IMAGE slots are filled with aligner rows in reading order.
+"""
+
+import base64
+import io
+import math
+from dataclasses import dataclass
+from urllib.request import urlopen
+
+import numpy as np
+import torch
+from PIL import Image, ImageOps
+
+TEXT = -1
+IMAGE_START, IMAGE, IMAGE_NEW_LINE, IMAGE_END = range(4)
+
+
+@dataclass
+class ImageInput:
+    start: int
+    patches: torch.Tensor
+    n_vit_h: int
+    n_vit_w: int
+    types: torch.Tensor
+
+
+def num_image_tokens(n_llm_h: int, n_llm_w: int) -> int:
+    return n_llm_h * (n_llm_w + 1) + 2
+
+
+def llm_grid(best_height: int, best_width: int, patch_size: int, downsample_ratio: int):
+    """Token grid the aligner produces from a patch grid of this pixel size."""
+    return math.ceil((best_height // patch_size) / downsample_ratio), math.ceil(
+        (best_width // patch_size) / downsample_ratio
+    )
+
+
+def solve_resize_ratio(height, width, patch_size, downsample_ratio, max_n_token):
+    """Largest aspect-preserving pixel size whose token grid still fits in max_n_token."""
+    r = height / width
+    max_w_float = math.sqrt((max_n_token - 2) / r + 0.25) - 0.5
+    max_h_float = max_w_float * r
+    cell = patch_size * downsample_ratio
+    if max_w_float < 1.0:  # very tall: collapse to a single column
+        return (max_n_token - 2) // 2 * cell, cell
+    if max_h_float < 1.0:  # very wide: collapse to a single row
+        return cell, (max_n_token - 3) * cell
+    beta = min(math.floor(max_w_float) * cell / width, math.floor(max_h_float) * cell / height)
+    return math.floor(height * beta / patch_size) * patch_size, math.floor(width * beta / patch_size) * patch_size
+
+
+def safe_resize(height, width, best_height, best_width, patch_size, downsample_ratio, max_n_token):
+    """Shrink the pixel size until the image costs at most max_n_token LLM tokens."""
+    n_llm_h, n_llm_w = llm_grid(best_height, best_width, patch_size, downsample_ratio)
+    if num_image_tokens(n_llm_h, n_llm_w) > max_n_token:
+        best_height, best_width = solve_resize_ratio(height, width, patch_size, downsample_ratio, max_n_token)
+        n_llm_h, n_llm_w = llm_grid(best_height, best_width, patch_size, downsample_ratio)
+        assert num_image_tokens(n_llm_h, n_llm_w) <= max_n_token
+    return n_llm_h, n_llm_w, best_height, best_width
+
+
+def load_image_bytes(record) -> bytes:
+    """Load image bytes from raw/base64 data, an Anthropic source, URL, or path."""
+    data = record.get("data")
+    if isinstance(data, bytes):
+        return data
+    if isinstance(data, str):
+        return base64.b64decode(data)
+
+    source = record.get("source")
+    if isinstance(source, dict):
+        if source.get("data") is not None:
+            return base64.b64decode(source["data"])
+        if source.get("url"):
+            return load_image_bytes({"url": source["url"]})
+
+    url = record.get("url")
+    if isinstance(url, str) and url:
+        if url.startswith("data:"):
+            header, _, payload = url.partition(",")
+            if ";base64" not in header:
+                raise ValueError(f"Unsupported data URL encoding: {header}")
+            return base64.b64decode(payload)
+        if url.startswith(("http://", "https://")):
+            with urlopen(url, timeout=30) as response:
+                return response.read()
+        with open(url, "rb") as file:
+            return file.read()
+
+    raise ValueError(f"Cannot load image from record: {list(record.keys())}")
+
+
+def plan_image_grid(width: int, height: int, args):
+    """Resize plan for an image of the given original size; a pure function of its arguments."""
+    p = args.vision_patch_size
+    if args.vision_max_wh_ratio is not None and width > height * args.vision_max_wh_ratio:
+        width = height * args.vision_max_wh_ratio
+    if 0 < width * height < args.vision_min_pixels:
+        ratio = (args.vision_min_pixels / (width * height)) ** 0.5
+        width = int(width * ratio)
+        height = int(height * ratio)
+    best_width = math.ceil(width / p) * p
+    best_height = math.ceil(height / p) * p
+    return safe_resize(height, width, best_height, best_width, p, args.vision_downsample_ratio, args.vision_max_n_token)
+
+
+def load_image(record, args):
+    """Load and transform one image record into ViT patches."""
+    p = args.vision_patch_size
+    with Image.open(io.BytesIO(load_image_bytes(record))) as source:
+        image = source.convert("RGB")
+    n_llm_h, n_llm_w, best_height, best_width = plan_image_grid(image.width, image.height, args)
+    n_vit_h, n_vit_w = best_height // p, best_width // p
+    if args.vision_max_wh_ratio is not None and image.width >= args.vision_max_wh_ratio * image.height:
+        image = image.resize((best_width, best_height))
+    else:
+        image = ImageOps.pad(image, (best_width, best_height), color=(127, 127, 127))
+    x = torch.from_numpy(np.asarray(image, dtype=np.float32)).permute(2, 0, 1) / 255
+    x = ((x - 0.5) / 0.5).to(torch.bfloat16)
+    patches = x.reshape(3, n_vit_h, p, n_vit_w, p).permute(1, 3, 0, 2, 4).reshape(n_vit_h * n_vit_w, 3, p, p)
+    return patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w
+
+
+def image_token_types(n_llm_h: int, n_llm_w: int) -> torch.Tensor:
+    """Default layout: the aligner grid in reading order, one IMAGE_NEW_LINE per row."""
+    types = [IMAGE_START]
+    types += ([IMAGE] * n_llm_w + [IMAGE_NEW_LINE]) * n_llm_h
+    types.append(IMAGE_END)
+    return torch.tensor(types, dtype=torch.int64)
+
+
+def prepare_vl_inputs(prompt, images, tokenizer, args):
+    """Tokenize `prompt`, expanding each image placeholder token into its image span.
+
+    Returns (tokens, token_types, image_inputs). Image-span positions carry `args.image_token_id` in
+    `tokens` and are distinguished only by `token_types` (TEXT elsewhere). `image_inputs` is None when
+    the prompt has no images."""
+    from encoding import IMAGE_PLACEHOLDER
+
+    # The placeholder is spelled differently across tokenizer revisions, so the id comes from the
+    # config; only cross-check it when this tokenizer does know the training-time spelling.
+    image_token_id = args.image_token_id
+    placeholder_id = tokenizer.convert_tokens_to_ids(IMAGE_PLACEHOLDER)
+    if placeholder_id is not None and placeholder_id != tokenizer.unk_token_id:
+        assert placeholder_id == image_token_id, (placeholder_id, image_token_id)
+    prompt_tokens = tokenizer.encode(prompt)
+    num_placeholders = sum(token == image_token_id for token in prompt_tokens)
+    if num_placeholders != len(images):
+        raise ValueError(f"Found {num_placeholders} image tokens but got {len(images)} images")
+    if num_placeholders and not args.vision_enabled:
+        raise ValueError("The model config has no vision tower (vision_n_layers == 0) but the prompt contains images")
+
+    tokens, token_types, image_inputs = [], [], []
+    image_iter = iter(images)
+    for tok in prompt_tokens:
+        if tok != image_token_id:
+            tokens.append(tok)
+            token_types.append(TEXT)
+            continue
+        patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w = load_image(next(image_iter), args)
+        types = image_token_types(n_llm_h, n_llm_w)
+        image_inputs.append(ImageInput(len(tokens), patches, n_vit_h, n_vit_w, types))
+        tokens += [image_token_id] * types.numel()
+        token_types += types.tolist()
+    return tokens, token_types, image_inputs or None

inference/kernel.py

@@ -0,0 +1,591 @@
+import torch
+import tilelang
+import tilelang.language as T
+from typing import Tuple, Optional
+
+
+tilelang.set_log_level("WARNING")
+
+pass_configs = {
+    tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
+    tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
+}
+
+FP8 = "float8_e4m3"
+FP4 = "float4_e2m1fn"
+FE8M0 = "float8_e8m0fnu"
+BF16 = "bfloat16"
+FP32 = "float32"
+INT32 = "int32"
+
+
+def fast_log2_ceil(x):
+    """Compute ceil(log2(x)) via IEEE 754 bit manipulation. Avoids slow log/ceil intrinsics."""
+    bits_x = T.reinterpret("uint32", x)
+    exp_x = (bits_x >> 23) & 0xFF
+    man_bits = bits_x & ((1 << 23) - 1)
+    return T.Cast("int32", exp_x - 127 + T.if_then_else(man_bits != 0, 1, 0))
+
+
+def fast_pow2(x):
+    """Compute 2^x for integer x via IEEE 754 bit manipulation."""
+    bits_x = (x + 127) << 23
+    return T.reinterpret("float32", bits_x)
+
+
+def fast_round_scale(amax, fp8_max_inv):
+    return fast_pow2(fast_log2_ceil(amax * fp8_max_inv))
+
+
+@tilelang.jit(pass_configs=pass_configs)
+def act_quant_kernel(N, block_size=128, in_dtype=BF16, out_dtype=FP8, scale_dtype=FP32, round_scale=False, inplace=False):
+    """Block-wise FP8 quantization. inplace=True does fused quant+dequant back to BF16."""
+    M = T.symbolic("M")
+    fp8_min = -448.0
+    fp8_max = 448.0
+    fp8_max_inv = 1 / fp8_max
+    num_stages = 0 if round_scale or inplace else 2
+    blk_m = 32
+    group_size = block_size
+    # Internal computation in FP32; scale_dtype controls output storage format.
+    compute_dtype = FP32
+    out_dtype = in_dtype if inplace else out_dtype
+
+    @T.prim_func
+    def act_quant_kernel_(
+        X: T.Tensor[(M, N), in_dtype],
+        Y: T.Tensor[(M, N), out_dtype],
+        S: T.Tensor[(M, T.ceildiv(N, group_size)), scale_dtype],
+    ):
+        with T.Kernel(T.ceildiv(M, blk_m), T.ceildiv(N, group_size), threads=128) as (
+            pid_m,
+            pid_n,
+        ):
+            x_shared = T.alloc_shared((blk_m, group_size), in_dtype)
+            x_local = T.alloc_fragment((blk_m, group_size), in_dtype)
+            amax_local = T.alloc_fragment((blk_m,), compute_dtype)
+            s_local = T.alloc_fragment((blk_m,), compute_dtype)
+            y_local = T.alloc_fragment((blk_m, group_size), out_dtype)
+            y_shared = T.alloc_shared((blk_m, group_size), out_dtype)
+
+            for _ in T.Pipelined(1, num_stages=num_stages):
+                T.copy(X[pid_m * blk_m, pid_n * group_size], x_shared)
+                T.copy(x_shared, x_local)
+                T.reduce_absmax(x_local, amax_local, dim=1)
+                for i in T.Parallel(blk_m):
+                    amax_local[i] = T.max(amax_local[i], 1e-4)
+                    if round_scale:
+                        s_local[i] = fast_round_scale(amax_local[i], fp8_max_inv)
+                    else:
+                        s_local[i] = amax_local[i] * fp8_max_inv
+                if inplace:
+                    for i, j in T.Parallel(blk_m, group_size):
+                        y_local[i, j] = T.Cast(
+                            out_dtype,
+                            T.Cast(compute_dtype, T.Cast(FP8, T.clamp(x_local[i, j] / s_local[i], fp8_min, fp8_max))) * s_local[i],
+                        )
+                else:
+                    for i, j in T.Parallel(blk_m, group_size):
+                        y_local[i, j] = T.clamp(x_local[i, j] / s_local[i], fp8_min, fp8_max)
+                for i in T.Parallel(blk_m):
+                    S[pid_m * blk_m + i, pid_n] = T.Cast(scale_dtype, s_local[i])
+                T.copy(y_local, y_shared)
+                T.copy(y_shared, Y[pid_m * blk_m, pid_n * group_size])
+
+    return act_quant_kernel_
+
+
+def act_quant(
+    x: torch.Tensor,
+    block_size: int = 128,
+    scale_fmt: Optional[str] = None,
+    scale_dtype: torch.dtype = torch.float32,
+    inplace: bool = False,
+) -> torch.Tensor:
+    """Block-wise FP8 quantization. inplace=True does fused quant+dequant back to BF16.
+    When scale_fmt is set, scales are rounded to power-of-2 (MXFP)."""
+    N = x.size(-1)
+    assert N % block_size == 0
+    tl_dtype = FE8M0 if scale_dtype == torch.float8_e8m0fnu else FP32
+    z = x.contiguous()
+    y = torch.empty_like(z) if inplace else torch.empty_like(z, dtype=torch.float8_e4m3fn)
+    s = z.new_empty(*z.size()[:-1], N // block_size, dtype=scale_dtype)
+    kernel = act_quant_kernel(
+        N,
+        block_size,
+        scale_dtype=tl_dtype,
+        round_scale=scale_fmt is not None,
+        inplace=inplace,
+    )
+    kernel(z.view(-1, N), y.view(-1, N), s.view(-1, N // block_size))
+    if inplace:
+        x.copy_(y)
+        return x
+    return y, s
+
+
+@tilelang.jit(pass_configs=pass_configs)
+def fp4_quant_kernel(N, block_size=32, in_dtype=BF16, scale_dtype=FE8M0, inplace=False):
+    """Block-wise FP4 with power-of-2 or E4M3 scales; optionally dequantize in place."""
+    M = T.symbolic("M")
+    fp4_max = 6.0
+    fp4_max_inv = 1.0 / fp4_max
+    blk_m = 32
+    group_size = block_size
+    compute_dtype = FP32
+    out_dtype = in_dtype if inplace else FP4
+
+    @T.prim_func
+    def fp4_quant_kernel_(
+        X: T.Tensor[(M, N), in_dtype],
+        Y: T.Tensor[(M, N), out_dtype],
+        S: T.Tensor[(M, T.ceildiv(N, group_size)), scale_dtype],
+    ):
+        with T.Kernel(T.ceildiv(M, blk_m), T.ceildiv(N, group_size), threads=128) as (
+            pid_m,
+            pid_n,
+        ):
+            x_shared = T.alloc_shared((blk_m, group_size), in_dtype)
+            x_local = T.alloc_fragment((blk_m, group_size), in_dtype)
+            amax_local = T.alloc_fragment((blk_m,), compute_dtype)
+            s_local = T.alloc_fragment((blk_m,), compute_dtype)
+            y_local = T.alloc_fragment((blk_m, group_size), out_dtype)
+            y_shared = T.alloc_shared((blk_m, group_size), out_dtype)
+
+            for _ in T.Pipelined(1, num_stages=2):
+                T.copy(X[pid_m * blk_m, pid_n * group_size], x_shared)
+                T.copy(x_shared, x_local)
+                T.reduce_absmax(x_local, amax_local, dim=1)
+                for i in T.Parallel(blk_m):
+                    if scale_dtype == FP8:
+                        # Training's compressed KV: keep even an all-zero group's scale nonzero.
+                        amax_local[i] = T.max(amax_local[i], 6 * (2**-9))
+                        s_local[i] = T.Cast(compute_dtype, T.Cast(FP8, amax_local[i] / fp4_max))
+                    else:
+                        amax_local[i] = T.max(amax_local[i], 6 * (2**-126))
+                        s_local[i] = fast_round_scale(amax_local[i], fp4_max_inv)
+                if inplace:
+                    for i, j in T.Parallel(blk_m, group_size):
+                        y_local[i, j] = T.Cast(
+                            out_dtype,
+                            T.Cast(compute_dtype, T.Cast(FP4, T.clamp(x_local[i, j] / s_local[i], -fp4_max, fp4_max))) * s_local[i],
+                        )
+                else:
+                    for i, j in T.Parallel(blk_m, group_size):
+                        y_local[i, j] = T.clamp(x_local[i, j] / s_local[i], -fp4_max, fp4_max)
+                for i in T.Parallel(blk_m):
+                    S[pid_m * blk_m + i, pid_n] = T.Cast(scale_dtype, s_local[i])
+                T.copy(y_local, y_shared)
+                T.copy(y_shared, Y[pid_m * blk_m, pid_n * group_size])
+
+    return fp4_quant_kernel_
+
+
+def fp4_act_quant(
+    x: torch.Tensor,
+    block_size: int = 32,
+    inplace: bool = False,
+    scale_dtype: torch.dtype = torch.float8_e8m0fnu,
+) -> torch.Tensor:
+    """FP4 with E8M0 scales for the indexer or E4M3 scales for compressed KV.
+    inplace=True writes the dequantized values back to x."""
+    assert scale_dtype in (torch.float8_e8m0fnu, torch.float8_e4m3fn)
+    N = x.size(-1)
+    assert N % block_size == 0
+    z = x.contiguous()
+    y = torch.empty_like(z) if inplace else z.new_empty(*z.shape[:-1], N // 2, dtype=torch.float4_e2m1fn_x2)
+    s = z.new_empty(*z.size()[:-1], N // block_size, dtype=scale_dtype)
+    tl_dtype = FP8 if scale_dtype == torch.float8_e4m3fn else FE8M0
+    kernel = fp4_quant_kernel(N, block_size, scale_dtype=tl_dtype, inplace=inplace)
+    kernel(z.view(-1, N), y.view(-1, y.size(-1)), s.view(-1, N // block_size))
+    if inplace:
+        x.copy_(y)
+        return x
+    return y, s
+
+
+@tilelang.jit(pass_configs=pass_configs)
+def fp8_gemm_kernel(
+    N,
+    K,
+    block_size=128,
+    out_dtype=BF16,
+    accum_dtype=FP32,
+    scale_dtype=FP32,
+):
+    assert out_dtype in [BF16, FP32]
+    assert block_size in [32, 128]
+
+    M = T.symbolic("M")
+    group_size = block_size
+    block_M = 32
+    block_N = 128
+    block_K = group_size
+    num_stages = 4 if block_K == 128 else 2
+
+    @T.prim_func
+    def fp8_gemm_kernel_(
+        A: T.Tensor[(M, K), FP8],
+        B: T.Tensor[(N, K), FP8],
+        C: T.Tensor[(M, N), out_dtype],
+        scales_a: T.Tensor[(M, T.ceildiv(K, group_size)), scale_dtype],
+        scales_b: T.Tensor[(T.ceildiv(N, group_size), T.ceildiv(K, group_size)), scale_dtype],
+    ):
+        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (
+            bx,
+            by,
+        ):
+            A_shared = T.alloc_shared((block_M, block_K), FP8)
+            B_shared = T.alloc_shared((block_N, block_K), FP8)
+            C_shared = T.alloc_shared((block_M, block_N), out_dtype)
+            C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
+            C_local_accum = T.alloc_fragment((block_M, block_N), accum_dtype)
+            scale_a_frag = T.alloc_fragment((block_M,), FP32)
+            scale_b_frag = T.alloc_fragment((block_N,), FP32)
+
+            # Improve L2 Cache
+            T.use_swizzle(panel_size=10)
+            T.clear(C_local)
+            T.clear(C_local_accum)
+
+            K_iters = T.ceildiv(K, block_K)
+            for k in T.Pipelined(K_iters, num_stages=num_stages):
+                T.copy(A[by * block_M, k * block_K], A_shared)
+                T.copy(B[bx * block_N, k * block_K], B_shared)
+
+                # Activation scales are per row and K block. Weight scales are
+                # per (N block, K block), where the block can be 32 or 128.
+                for i in T.Parallel(block_M):
+                    scale_a_frag[i] = T.Cast(FP32, scales_a[by * block_M + i, k])
+                for j in T.Parallel(block_N):
+                    scale_b_frag[j] = T.Cast(
+                        FP32,
+                        scales_b[(bx * block_N + j) // group_size, k],
+                    )
+
+                T.gemm(A_shared, B_shared, C_local, transpose_B=True)
+                # Separate accumulator for scale-corrected results.
+                for i, j in T.Parallel(block_M, block_N):
+                    C_local_accum[i, j] += C_local[i, j] * scale_a_frag[i] * scale_b_frag[j]
+                T.clear(C_local)
+            T.copy(C_local_accum, C_shared)
+            T.copy(C_shared, C[by * block_M, bx * block_N])
+
+    return fp8_gemm_kernel_
+
+
+def fp8_gemm(
+    a: torch.Tensor,
+    a_s: torch.Tensor,
+    b: torch.Tensor,
+    b_s: torch.Tensor,
+    scale_dtype: torch.dtype = torch.float32,
+    block_size: int = 128,
+) -> torch.Tensor:
+    """C[M,N] = A[M,K] @ B[N,K]^T with per-block FP8 scaling."""
+    assert a.is_contiguous() and b.is_contiguous(), "Input tensors must be contiguous"
+    assert a_s.is_contiguous() and b_s.is_contiguous(), "Scaling factor tensors must be contiguous"
+    assert block_size in (32, 128)
+    tl_dtype = FE8M0 if scale_dtype == torch.float8_e8m0fnu else FP32
+    K = a.size(-1)
+    M = a.numel() // K
+    N = b.size(0)
+    assert K % block_size == 0
+    assert a_s.numel() == M * (K // block_size)
+    assert b_s.shape == (
+        (N + block_size - 1) // block_size,
+        K // block_size,
+    )
+    c = a.new_empty(*a.size()[:-1], N, dtype=torch.get_default_dtype())
+    kernel = fp8_gemm_kernel(
+        N,
+        K,
+        block_size=block_size,
+        scale_dtype=tl_dtype,
+    )
+    kernel(a.view(M, K), b, c.view(M, N), a_s.view(M, -1), b_s)
+    return c
+
+
+@tilelang.jit(pass_configs=pass_configs)
+def sparse_attn_kernel(h: int, d: int, scale=None):
+    """Sparse multi-head attention via index gathering + online softmax (FlashAttention-style).
+    For each (batch, seq_pos), gathers top-k KV positions by index, computes attention
+    with numerically stable running max/sum, and includes a learnable attn_sink bias."""
+    b = T.symbolic("b")
+    m = T.symbolic("m")
+    n = T.symbolic("n")
+    topk = T.symbolic("topk")
+    if scale is None:
+        scale = (1.0 / d) ** 0.5
+
+    num_stages = 2
+    threads = 256
+    block = 64
+    num_blocks = tilelang.cdiv(topk, block)
+
+    @T.prim_func
+    def sparse_attn_kernel_(
+        q: T.Tensor[(b, m, h, d), BF16],
+        kv: T.Tensor[(b, n, d), BF16],
+        o: T.Tensor[(b, m, h, d), BF16],
+        attn_sink: T.Tensor[(h,), FP32],
+        topk_idxs: T.Tensor[(b, m, topk), INT32],
+    ):
+        with T.Kernel(m, b, threads=threads) as (bx, by):
+            q_shared = T.alloc_shared((h, d), BF16)
+            kv_shared = T.alloc_shared((block, d), BF16)
+            o_shared = T.alloc_shared((h, d), BF16)
+            acc_s_cast = T.alloc_shared((h, block), BF16)
+
+            idxs = T.alloc_fragment(block, INT32)
+            acc_s = T.alloc_fragment((h, block), FP32)
+            acc_o = T.alloc_fragment((h, d), FP32)
+            scores_max = T.alloc_fragment(h, FP32)
+            scores_max_prev = T.alloc_fragment(h, FP32)
+            scores_scale = T.alloc_fragment(h, FP32)
+            scores_sum = T.alloc_fragment(h, FP32)
+            sum_exp = T.alloc_fragment(h, FP32)
+
+            T.clear(acc_o)
+            T.clear(sum_exp)
+            # Finite lower bound instead of -inf: a row with no valid index (all -1) would otherwise
+            # produce exp(-inf - (-inf)) = NaN. With a finite bound such rows yield an all-zero output,
+            # matching the training kernel's convention.
+            T.fill(scores_max, -1e30)
+            T.copy(q[by, bx, :, :], q_shared)
+
+            for t in T.Pipelined(num_blocks, num_stages=num_stages):
+                for i in T.Parallel(block):
+                    idxs[i] = T.if_then_else(t * block + i < topk, topk_idxs[by, bx, t * block + i], -1)
+                for i, j in T.Parallel(block, d):
+                    kv_shared[i, j] = T.if_then_else(idxs[i] != -1, kv[by, idxs[i], j], 0)
+                for i, j in T.Parallel(h, block):
+                    acc_s[i, j] = T.if_then_else(idxs[j] != -1, 0, -T.infinity(FP32))
+                T.gemm(q_shared, kv_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
+                for i, j in T.Parallel(h, block):
+                    acc_s[i, j] *= scale
+                T.copy(scores_max, scores_max_prev)
+                T.reduce_max(acc_s, scores_max, dim=1, clear=False)
+                for i in T.Parallel(h):
+                    scores_scale[i] = T.exp(scores_max_prev[i] - scores_max[i])
+                for i, j in T.Parallel(h, block):
+                    acc_s[i, j] = T.exp(acc_s[i, j] - scores_max[i])
+                T.reduce_sum(acc_s, scores_sum, dim=1)
+                for i in T.Parallel(h):
+                    sum_exp[i] = sum_exp[i] * scores_scale[i] + scores_sum[i]
+                T.copy(acc_s, acc_s_cast)
+                for i, j in T.Parallel(h, d):
+                    acc_o[i, j] *= scores_scale[i]
+                T.gemm(acc_s_cast, kv_shared, acc_o, policy=T.GemmWarpPolicy.FullRow)
+
+            for i in T.Parallel(h):
+                sum_exp[i] += T.exp(attn_sink[i] - scores_max[i])
+            for i, j in T.Parallel(h, d):
+                acc_o[i, j] /= sum_exp[i]
+            T.copy(acc_o, o_shared)
+            T.copy(o_shared, o[by, bx, :, :])
+
+    return sparse_attn_kernel_
+
+
+def sparse_attn(q: torch.Tensor, kv: torch.Tensor, attn_sink: torch.Tensor, topk_idxs: torch.Tensor, softmax_scale: float) -> torch.Tensor:
+    b, s, h, d = q.size()
+    # Pad heads to 16 for kernel efficiency (stripped after)
+    if h < 16:
+        q = torch.cat([q, q.new_zeros(b, s, 16 - h, d)], dim=2)
+        attn_sink = torch.cat([attn_sink, attn_sink.new_zeros(16 - h)])
+    o = torch.empty_like(q)
+    kernel = sparse_attn_kernel(q.size(2), d, softmax_scale)
+    kernel(q, kv, o, attn_sink, topk_idxs)
+    if h < 16:
+        o = o.narrow(2, 0, h).contiguous()
+    return o
+
+
+@tilelang.jit(pass_configs=pass_configs)
+def hc_split_sinkhorn_kernel(hc: int, sinkhorn_iters: int, eps: float):
+    n = T.symbolic("n")
+    mix_hc = (2 + hc) * hc
+    threads = 64
+
+    @T.prim_func
+    def hc_split_sinkhorn_kernel_(
+        mixes: T.Tensor[(n, mix_hc), FP32],
+        hc_scale: T.Tensor[(3,), FP32],
+        hc_base: T.Tensor[(mix_hc,), FP32],
+        pre: T.Tensor[(n, hc), FP32],
+        post: T.Tensor[(n, hc), FP32],
+        comb: T.Tensor[(n, hc, hc), FP32],
+    ):
+        with T.Kernel(n, threads=threads) as i:
+            mixes_shared = T.alloc_shared(mix_hc, FP32)
+            comb_frag = T.alloc_fragment((hc, hc), FP32)
+            T.copy(mixes[i, :], mixes_shared)
+
+            for j in T.Parallel(hc):
+                pre[i, j] = T.sigmoid(mixes_shared[j] * hc_scale[0] + hc_base[j]) + eps
+            for j in T.Parallel(hc):
+                post[i, j] = 2 * T.sigmoid(mixes_shared[j + hc] * hc_scale[1] + hc_base[j + hc])
+            for j, k in T.Parallel(hc, hc):
+                comb_frag[j, k] = mixes_shared[j * hc + k + hc * 2] * hc_scale[2] + hc_base[j * hc + k + hc * 2]
+
+            row_sum = T.alloc_fragment(hc, FP32)
+            col_sum = T.alloc_fragment(hc, FP32)
+
+            # comb = comb.softmax(-1) + eps
+            row_max = T.alloc_fragment(hc, FP32)
+            T.reduce_max(comb_frag, row_max, dim=1)
+            for j, k in T.Parallel(hc, hc):
+                comb_frag[j, k] = T.exp(comb_frag[j, k] - row_max[j])
+            T.reduce_sum(comb_frag, row_sum, dim=1)
+            for j, k in T.Parallel(hc, hc):
+                comb_frag[j, k] = comb_frag[j, k] / row_sum[j] + eps
+
+            # comb = comb / (comb.sum(-2) + eps)
+            T.reduce_sum(comb_frag, col_sum, dim=0)
+            for j, k in T.Parallel(hc, hc):
+                comb_frag[j, k] = comb_frag[j, k] / (col_sum[k] + eps)
+
+            for _ in T.serial(sinkhorn_iters - 1):
+                # comb = comb / (comb.sum(-1) + eps)
+                T.reduce_sum(comb_frag, row_sum, dim=1)
+                for j, k in T.Parallel(hc, hc):
+                    comb_frag[j, k] = comb_frag[j, k] / (row_sum[j] + eps)
+                # comb = comb / (comb.sum(-2) + eps)
+                T.reduce_sum(comb_frag, col_sum, dim=0)
+                for j, k in T.Parallel(hc, hc):
+                    comb_frag[j, k] = comb_frag[j, k] / (col_sum[k] + eps)
+
+            T.copy(comb_frag, comb[i, :, :])
+
+    return hc_split_sinkhorn_kernel_
+
+
+def hc_split_sinkhorn(
+    mixes: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor, hc_mult: int = 4, sinkhorn_iters: int = 20, eps: float = 1e-6
+):
+    b, s, _ = mixes.size()
+    pre = mixes.new_empty(b, s, hc_mult)
+    post = mixes.new_empty(b, s, hc_mult)
+    comb = mixes.new_empty(b, s, hc_mult, hc_mult)
+    kernel = hc_split_sinkhorn_kernel(hc_mult, sinkhorn_iters, eps)
+    kernel(mixes.view(-1, (2 + hc_mult) * hc_mult), hc_scale, hc_base, pre.view(-1, hc_mult), post.view(-1, hc_mult), comb.view(-1, hc_mult, hc_mult))
+    return pre, post, comb
+
+
+@tilelang.jit(pass_configs=pass_configs)
+def fp4_gemm_kernel(
+    N,
+    K,
+    act_block_size=128,
+    out_dtype=BF16,
+    accum_dtype=FP32,
+    scale_dtype=FP32,
+):
+    """FP8 act x FP4 weight GEMM kernel.
+
+    C[M, N] = A_fp8[M, K] @ B_fp4[N, K]^T
+
+    Act: 1x32 or 1x128 quant on K, FP8 with configurable scale dtype
+    Weight: 1x32 quant on K (reduce dim), FP4 with E8M0 scale
+
+    B is stored as [N, K//2] in float4_e2m1fn_x2, logical [N, K] in fp4.
+    The FP4 values are packed along the K (last) dimension.
+
+    Strategy: load FP4 sub-blocks of size [block_N, sub_K] (sub_K=32),
+    cast FP4 to FP8 via float, then do FP8xFP8 GEMM.
+    Apply activation and weight scales to the accumulator.
+    """
+    M = T.symbolic("M")
+    assert act_block_size in [32, 128]
+    act_group_size = act_block_size
+    weight_group_size = 32
+    block_M = 32
+    block_N = 128
+    block_K = 32  # matches weight_group_size for simple scale handling
+    n_sub = act_group_size // block_K
+
+    @T.prim_func
+    def fp4_gemm_kernel_(
+        A: T.Tensor[(M, K), FP8],
+        B: T.Tensor[(N, K), FP4],
+        C: T.Tensor[(M, N), out_dtype],
+        scales_a: T.Tensor[(M, T.ceildiv(K, act_group_size)), scale_dtype],
+        scales_b: T.Tensor[(N, T.ceildiv(K, weight_group_size)), scale_dtype],
+    ):
+        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (
+            bx,
+            by,
+        ):
+            A_shared = T.alloc_shared((block_M, block_K), FP8)
+            B_fp4_shared = T.alloc_shared((block_N, block_K), FP4)
+            B_shared = T.alloc_shared((block_N, block_K), FP8)
+            C_shared = T.alloc_shared((block_M, block_N), out_dtype)
+            C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
+            C_local_accum = T.alloc_fragment((block_M, block_N), accum_dtype)
+            scale_a_frag = T.alloc_fragment((block_M,), FP32)
+            scale_b_frag = T.alloc_fragment((block_N,), FP32)
+
+            T.use_swizzle(panel_size=10)
+            T.clear(C_local)
+            T.clear(C_local_accum)
+
+            K_iters = T.ceildiv(K, block_K)
+            for k in T.Pipelined(K_iters, num_stages=2):
+                T.copy(A[by * block_M, k * block_K], A_shared)
+                T.copy(B[bx * block_N, k * block_K], B_fp4_shared)
+                # FP4->FP8 cast must go through FP32 to avoid ambiguous C++ overload
+                for i, j in T.Parallel(block_N, block_K):
+                    B_shared[i, j] = T.Cast(FP8, T.Cast(FP32, B_fp4_shared[i, j]))
+
+                # Weight scale: per 32 on K, indexed by k (each k is one block_K=32)
+                for i in T.Parallel(block_N):
+                    scale_b_frag[i] = T.Cast(FP32, scales_b[bx * block_N + i, k])
+
+                # Activation scale: one value per act_group_size elements.
+                for i in T.Parallel(block_M):
+                    scale_a_frag[i] = T.Cast(FP32, scales_a[by * block_M + i, k // n_sub])
+
+                T.gemm(A_shared, B_shared, C_local, transpose_B=True)
+
+                for i, j in T.Parallel(block_M, block_N):
+                    C_local_accum[i, j] += C_local[i, j] * scale_a_frag[i] * scale_b_frag[j]
+                T.clear(C_local)
+
+            T.copy(C_local_accum, C_shared)
+            T.copy(C_shared, C[by * block_M, bx * block_N])
+
+    return fp4_gemm_kernel_
+
+
+def fp4_gemm(
+    a: torch.Tensor,
+    a_s: torch.Tensor,
+    b: torch.Tensor,
+    b_s: torch.Tensor,
+    scale_dtype: torch.dtype = torch.float32,
+    act_block_size: int = 128,
+) -> torch.Tensor:
+    """C[M,N] = A_fp8[M,K] @ B_fp4[N,K]^T.
+    A has per-32 or per-128 activation scale; B has per-32 E8M0 weight scale.
+    B is stored as [N, K//2] in float4_e2m1fn_x2 (2 FP4 values per byte, packed along K)."""
+    assert a.is_contiguous() and b.is_contiguous(), "Input tensors must be contiguous"
+    assert a_s.is_contiguous() and b_s.is_contiguous(), "Scaling factor tensors must be contiguous"
+    tl_dtype = FE8M0 if scale_dtype == torch.float8_e8m0fnu else FP32
+    K = a.size(-1)
+    M = a.numel() // K
+    N = b.size(0)
+    assert act_block_size in (32, 128)
+    assert K % act_block_size == 0
+    assert a_s.numel() == M * (K // act_block_size)
+    assert b_s.shape == (N, K // 32)
+    c = a.new_empty(*a.size()[:-1], N, dtype=torch.get_default_dtype())
+    kernel = fp4_gemm_kernel(
+        N,
+        K,
+        act_block_size=act_block_size,
+        scale_dtype=tl_dtype,
+    )
+    kernel(a.view(M, K), b, c.view(M, N), a_s.view(M, -1), b_s)
+    return c

inference/model.py

@@ -0,0 +1,1309 @@
+import math
+from contextlib import contextmanager
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Literal
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from engram import EngramLayout, NgramHashState
+from image_processor import IMAGE, IMAGE_END, IMAGE_NEW_LINE, IMAGE_START
+from kernel import (
+    act_quant,
+    fp4_act_quant,
+    fp4_gemm,
+    fp8_gemm,
+    hc_split_sinkhorn,
+    sparse_attn,
+)
+from torch import nn
+from vision import Aligner, ViT
+
+# Set once by Transformer.__init__; one model per process, so layers just read them.
+world_size = 1
+rank = 0
+default_dtype = torch.float8_e4m3fn  # storage dtype for Linear weights, from ModelArgs.dtype
+fp8_block_size = 32  # one fp8 scale per 32x32 weight block / 32 activations
+fp4_block_size = 32  # one fp4 scale per 32 elements along K
+scale_fmt = "ue8m0"
+scale_dtype = torch.float8_e8m0fnu
+
+
+@contextmanager
+def set_dtype(dtype):
+    """Temporarily override torch's default dtype, restoring it even if the body raises."""
+    prev = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    try:
+        yield
+    finally:
+        torch.set_default_dtype(prev)
+
+
+@dataclass
+class ModelArgs:
+    """Field names are exactly the config JSON keys. The defaults are a small model that
+    `python model.py` can run, not the released shapes -- though the scale-independent
+    values (norm_eps, score_func, hc_*, engram_*) do match it."""
+
+    # runtime limits rather than model shape: they size the KV caches
+    max_batch_size: int = 4
+    max_seq_len: int = 4096
+    temperature: float = 1
+    dtype: Literal["bf16", "fp8"] = "fp8"
+    expert_dtype: Literal["fp4"] | None = "fp4"
+    vocab_size: int = 129280
+    dim: int = 1024
+    moe_inter_dim: int = 1024
+    n_layers: int = 5
+    n_mtp_layers: int = 1  # extra draft layers appended after the backbone, indices n_layers..
+    n_heads: int = 16
+    # moe
+    n_routed_experts: int = 8
+    n_shared_experts: int = 1
+    n_activated_experts: int = 2
+    score_func: Literal["softmax", "sigmoid", "sqrtsoftplus"] = "sqrtsoftplus"
+    gate_temp: float = 1.0
+    norm_topk_prob: bool = True
+    route_scale: float = 1.0
+    swiglu_limit: float = 0.0
+    # attention: latent q/kv projections, plus a LoRA-factorised output projection over o_groups
+    q_lora_rank: int = 256
+    head_dim: int = 128
+    rope_head_dim: int = 32
+    norm_eps: float = 1e-20
+    o_groups: int = 8
+    o_lora_rank: int = 256
+    # sparse attention: every layer attends over a sliding window, and may add compressed KV on top
+    window_size: int = 128
+    # one entry per layer, MTP layers included: 0 = sliding window only, r = KV compressed r-to-1
+    compress_ratios: tuple[int, ...] = (0, 2, 2, 1, 1, 0)
+    # layers sharing a ratio also share one compressed KV and one indexer, produced by the first
+    kv_source_layers: tuple[int, ...] = (1, 3)
+    index_source_layers: tuple[int, ...] = (1, 3)
+    # rope, with YaRN extrapolation when original_seq_len > 0. Compressed KV rotates at its own
+    # theta because one latent stands for compress_ratio tokens, so its positions are further apart.
+    compress_rope_theta: float = 40000.0
+    original_seq_len: int = 0
+    rope_theta: float = 10000.0
+    rope_factor: float = 40
+    beta_fast: int = 32
+    beta_slow: int = 1
+    # the indexer: a small extra attention that scores compressed positions, so each query can keep
+    # just `index_topk` of them. Names match DeepSeek-V3.2-Exp, where this mechanism first appeared.
+    index_n_heads: int = 16
+    index_head_dim: int = 64
+    index_topk: int = 64
+    # candidate pre-filtering: candidate_source_layer < 0 turns it off and the other two are unused
+    candidate_source_layer: int = -1
+    candidate_topk_blocks: int = 0
+    candidate_block_size: int = 0
+    # hyper-connections: the residual stream is carried as hc_mult parallel copies
+    hc_mult: int = 4
+    hc_sinkhorn_iters: int = 20
+    hc_eps: float = 1e-6
+    # engram: n-gram hash lookups added into the residual stream at a few layers
+    engram_layer_ids: tuple[int, ...] = ()
+    engram_num_embeddings: tuple[int, ...] = ()  # unpadded table rows; each rank allocates ceil(rows / world_size)
+    engram_max_ngram_size: int = 1
+    engram_vocab_size: int = 0  # bucket size each (n-gram size, head) starts searching primes from
+    engram_n_heads: int = 0
+    engram_head_dim: int = 0
+    engram_pad_id: int = 2  # token that fills n-gram slots with no history; matches training
+    # size of the compressed tokenizer vocab; every hash multiplier is derived from it
+    engram_compressed_vocab_size: int = 0
+    # vision (VL); vision_n_layers == 0 disables the vision path
+    vision_n_layers: int = 0
+    vision_dim: int = 1024
+    vision_n_heads: int = 16
+    vision_inter_dim: int = 2816
+    vision_patch_size: int = 14
+    vision_rope_theta: float = 10000.0
+    vision_downsample_ratio: int = 3
+    vision_max_n_token: int = 1024
+    vision_min_pixels: int = 544 * 544
+    vision_max_wh_ratio: int | None = None
+    # raw id of <｜deepseek_image｜>; every position of an image span carries this id in input_ids
+    image_token_id: int = 129264
+    # dspark draft head. Only the forward pass is implemented here -- nothing calls forward_spec,
+    # so these are read but the speculative-decoding loop itself is out of scope for this repo.
+    dspark_block_size: int = 0
+    dspark_noise_token_id: int = 0
+    dspark_target_layer_ids: tuple[int, ...] = ()
+    dspark_markov_rank: int = 256
+    dspark_n_routed_experts: int = 0
+    dspark_n_activated_experts: int = 0
+
+    @property
+    def vision_enabled(self) -> bool:
+        return self.vision_n_layers > 0
+
+    def get_moe_config(self, layer_id: int) -> tuple[int, int]:
+        """Return the routed/activated expert counts for a given layer."""
+        if layer_id < self.n_layers:
+            return self.n_routed_experts, self.n_activated_experts
+        return (
+            self.dspark_n_routed_experts or self.n_routed_experts,
+            self.dspark_n_activated_experts or self.n_activated_experts,
+        )
+
+
+class ParallelEmbedding(nn.Module):
+    """Embedding sharded along the vocab dimension. Each rank holds vocab_size // world_size rows.
+    Out-of-range indices are zero-masked before all_reduce to combine partial embeddings."""
+
+    def __init__(self, vocab_size: int, dim: int):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.dim = dim
+        assert vocab_size % world_size == 0, (
+            f"Vocabulary size must be divisible by world size (world_size={world_size})"
+        )
+        self.part_vocab_size = vocab_size // world_size
+        self.vocab_start_idx = rank * self.part_vocab_size
+        self.vocab_end_idx = self.vocab_start_idx + self.part_vocab_size
+        self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if world_size > 1:
+            # ids off this rank read row 0 then get zeroed, so the all_reduce sums one real row
+            mask = (x < self.vocab_start_idx) | (x >= self.vocab_end_idx)
+            x = x - self.vocab_start_idx
+            x[mask] = 0
+        y = F.embedding(x, self.weight)
+        if world_size > 1:
+            y[mask] = 0
+            dist.all_reduce(y)
+        return y
+
+
+def linear(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
+    """Pick a GEMM from the weight dtype. Quantized weights need a quantized activation, and both
+    fp4 and fp8 weights take an fp8 one -- for fp4 the kernel handles the mixed precision."""
+    assert bias is None
+
+    if weight.dtype == torch.float4_e2m1fn_x2:
+        x, s = act_quant(x, fp8_block_size, scale_fmt, scale_dtype)
+        return fp4_gemm(
+            x,
+            s,
+            weight,
+            weight.scale,
+            scale_dtype,
+            act_block_size=fp8_block_size,
+        )
+    elif weight.dtype == torch.float8_e4m3fn:
+        x, s = act_quant(x, fp8_block_size, scale_fmt, scale_dtype)
+        return fp8_gemm(
+            x,
+            s,
+            weight,
+            weight.scale,
+            scale_dtype,
+            block_size=fp8_block_size,
+        )
+    else:
+        return F.linear(x, weight)
+
+
+class Linear(nn.Module):
+    """bf16, fp8 or fp4 weights. Quantized ones get a `scale`, also attached to `.weight` so that
+    `linear()` can reach it."""
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        dtype = dtype or default_dtype
+        if dtype == torch.float4_e2m1fn_x2:
+            # two values per byte: [out, in] logically, [out, in//2] stored
+            self.weight = nn.Parameter(torch.empty(out_features, in_features // 2, dtype=torch.float4_e2m1fn_x2))
+            self.weight.scale = self.scale = nn.Parameter(
+                torch.empty(out_features, in_features // fp4_block_size, dtype=torch.float8_e8m0fnu)
+            )
+        elif dtype == torch.float8_e4m3fn:
+            self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype))
+            self.weight.scale = self.scale = nn.Parameter(
+                torch.empty(
+                    (out_features + fp8_block_size - 1) // fp8_block_size,
+                    (in_features + fp8_block_size - 1) // fp8_block_size,
+                    dtype=torch.float8_e8m0fnu,
+                )
+            )
+        else:
+            self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype))
+            self.register_parameter("scale", None)
+        if bias:
+            self.bias = nn.Parameter(torch.empty(out_features))
+        else:
+            self.register_parameter("bias", None)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return linear(x, self.weight, self.bias)
+
+
+class ColumnParallelLinear(Linear):
+    """Splits the output dim across ranks; each rank's slice of the output is already complete."""
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None):
+        assert out_features % world_size == 0, (
+            f"Output features must be divisible by world size (world_size={world_size})"
+        )
+        self.part_out_features = out_features // world_size
+        super().__init__(in_features, self.part_out_features, bias, dtype)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return linear(x, self.weight, self.bias)
+
+
+class RowParallelLinear(Linear):
+    """Splits the reduction dim, so each rank holds a partial sum: hence the fp32 all_reduce, with
+    the bias added only after it."""
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None):
+        assert in_features % world_size == 0, (
+            f"Input features must be divisible by world size (world_size={world_size})"
+        )
+        self.part_in_features = in_features // world_size
+        super().__init__(self.part_in_features, out_features, bias, dtype)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = linear(x, self.weight, None)
+        if world_size > 1:
+            y = y.float()
+            dist.all_reduce(y)
+        if self.bias is not None:
+            y += self.bias
+        return y.type_as(x)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor):
+        dtype = x.dtype
+        x = x.float()
+        var = x.square().mean(-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.eps)
+        return (self.weight * x).to(dtype)
+
+
+class ParallelEngramEmbedding(nn.Module):
+    """The n-gram hash table, sharded over its rows. Stays fp8: rows are dequantized on lookup."""
+
+    def __init__(self, num_embeddings: int, dim: int):
+        super().__init__()
+        self.num_embeddings = num_embeddings
+        self.dim = dim
+        self.part_num_embeddings = (num_embeddings + world_size - 1) // world_size
+        self.vocab_start_idx = rank * self.part_num_embeddings
+        self.vocab_end_idx = self.vocab_start_idx + self.part_num_embeddings
+
+        self.block_size = fp8_block_size
+        # the table stays fp8 as stored: rows are dequantized with `scale` on lookup
+        self.weight = nn.Parameter(torch.empty(self.part_num_embeddings, dim, dtype=torch.float8_e4m3fn))
+        self.scale = nn.Parameter(torch.empty(self.part_num_embeddings, dim // self.block_size, dtype=scale_dtype))
+
+    def forward(self, indices: torch.Tensor) -> torch.Tensor:
+        mask = (indices < self.vocab_start_idx) | (indices >= self.vocab_end_idx)
+        local_indices = indices - self.vocab_start_idx
+        local_indices = local_indices.masked_fill(mask, 0)
+
+        values = F.embedding(local_indices, self.weight)
+        scales = F.embedding(local_indices, self.scale)
+        values = values.float().unflatten(-1, (-1, self.block_size)) * scales.float().unsqueeze(-1)
+        values = values.flatten(-2).to(torch.bfloat16)
+        values = values.masked_fill(mask.unsqueeze(-1), 0)
+
+        if world_size > 1:
+            dist.all_reduce(values)
+        return values
+
+
+class Engram(nn.Module):
+    """Writes an n-gram lookup into the residual stream, gated by how well it matches that stream.
+
+    The hash ids fetch `n_hash_cols` rows; `wkv` turns them into one key per hc copy plus a shared
+    value. The gate is a normalized dot product of stream against key.
+    """
+
+    def __init__(self, args: ModelArgs, layer_id: int, layout: EngramLayout):
+        super().__init__()
+        self.layer_id = layer_id
+        self.layer_hash_index = layout.layer_ids.index(layer_id)
+        self.dim = args.dim
+        self.hc_mult = args.hc_mult
+        self.clamp_value = 1e-6
+
+        self.embed = ParallelEngramEmbedding(layout.num_embeddings[self.layer_hash_index], layout.head_dim)
+        n_hash_cols = (layout.max_ngram_size - 1) * layout.n_heads
+        self.wkv = Linear(n_hash_cols * layout.head_dim, args.dim * (args.hc_mult + 1))
+        self.eps = args.norm_eps
+        self.q_weight = nn.Parameter(torch.ones(args.hc_mult, args.dim))
+        self.k_weight = nn.Parameter(torch.ones(args.hc_mult, args.dim))
+
+    def forward(self, x: torch.Tensor, hash_ids: torch.Tensor, token_mask: torch.Tensor | None = None) -> torch.Tensor:
+        """x: [B, L, hc_mult, dim]; hash_ids: [B, L, n_hash_cols]; token_mask: [B, L], False shuts
+        the gate so those positions pass through untouched."""
+        kv = self.wkv(self.embed(hash_ids).flatten(-2))
+        key, value = kv.split([self.hc_mult * self.dim, self.dim], dim=-1)
+        key = key.float().unflatten(-1, (self.hc_mult, self.dim))
+        weight = self.q_weight.float() * self.k_weight.float()  # only ever used as a product
+        h, eps = x.float(), self.eps
+        # normalized per (token, hc copy) over `dim`, NOT jointly over the copies
+        rstd = torch.rsqrt(h.square().mean(-1) + eps) * torch.rsqrt(key.square().mean(-1) + eps)
+        dot = (h * weight * key).sum(-1) * rstd * self.dim**-0.5
+        # signed sqrt before the sigmoid, matching the training kernel
+        gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(self.clamp_value).sqrt(), dot))
+        if token_mask is not None:
+            gate = gate.masked_fill(~token_mask.unsqueeze(-1), 0)
+        return (h + gate.unsqueeze(-1) * value.float().unsqueeze(-2)).to(x.dtype)
+
+
+@lru_cache(2)
+def precompute_freqs_cis(dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow) -> torch.Tensor:
+    """Rotary frequencies as complex exponentials, one row per position.
+
+    With original_seq_len > 0 this applies YaRN: dimensions whose wavelength already fits inside the
+    training context keep their frequency, those far beyond it are divided by `factor`, and the
+    `beta_fast`..`beta_slow` band in between is faded across with a linear ramp.
+    """
+    freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+    if original_seq_len > 0:
+        # the dim whose wavelength completes `rotations` turns over the training context
+        def corrected_dim(rotations):
+            return dim * math.log(original_seq_len / (rotations * 2 * math.pi)) / (2 * math.log(base))
+
+        low = max(math.floor(corrected_dim(beta_fast)), 0)
+        high = min(math.ceil(corrected_dim(beta_slow)), dim - 1)
+        ramp = ((torch.arange(dim // 2, dtype=torch.float32) - low) / max(high - low, 1e-3)).clamp(0, 1)
+        smooth = 1 - ramp
+        freqs = freqs / factor * (1 - smooth) + freqs * smooth
+
+    freqs = torch.outer(torch.arange(seqlen), freqs)
+    return torch.polar(torch.ones_like(freqs), freqs)
+
+
+def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False) -> torch.Tensor:
+    """Rotate `x` in place, taking adjacent element pairs as complex numbers. Accepts [b, s, d] and
+    [b, s, h, d]; `inverse` conjugates the rotation, which is how the attention output gets the
+    query's rotation removed again so the cache can stay in one shared rotated form."""
+    y = x
+    x = torch.view_as_complex(x.float().unflatten(-1, (-1, 2)))
+    if inverse:
+        freqs_cis = freqs_cis.conj()
+    if x.ndim == 3:
+        freqs_cis = freqs_cis.view(1, x.size(1), x.size(-1))
+    else:
+        freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1))
+    x = torch.view_as_real(x * freqs_cis).flatten(-2)
+    y.copy_(x)
+    return y
+
+
+@lru_cache(1)
+def get_window_topk_idxs(window_size: int, bsz: int, seqlen: int, start_pos: int):
+    """Which sliding-window cache slots each query attends to; -1 marks a slot holding nothing.
+
+    The cache is a ring of `window_size` slots. Prefill needs one row per query, each seeing its own
+    causal window. A decode step has a single query that sees the whole ring, listed oldest first.
+    Order within a row does not matter to `sparse_attn`, which handles every slot independently.
+    """
+    if start_pos == 0:
+        end = torch.arange(seqlen).unsqueeze(1)
+        idxs = (end - window_size + 1).clamp(0) + torch.arange(min(seqlen, window_size))
+        idxs = torch.where(idxs > end, -1, idxs)  # before the sequence started
+    else:
+        oldest = start_pos % window_size + 1
+        idxs = torch.cat([torch.arange(oldest, window_size), torch.arange(oldest)])
+        idxs = torch.where(idxs > start_pos, -1, idxs)  # ring still filling
+    # sparse_attn needs real [b, m, topk] int32 memory, hence the materializing expand
+    return idxs.int().unsqueeze(0).expand(bsz, -1, -1).contiguous()
+
+
+class Compressor(nn.Module):
+    """Pools `compress_ratio` consecutive tokens into one KV latent with a learned softmax gate.
+
+    Returns the latent before RoPE, or None while a group is still filling up -- so during decode it
+    only yields every `compress_ratio` steps, holding the partial group in `kv_state`/`score_state`.
+    Pre-RoPE is deliberate: the indexer needs the unrotated form, so Attention rotates afterwards.
+    """
+
+    def __init__(self, args: ModelArgs, layer_id: int):
+        super().__init__()
+        compress_ratio = args.compress_ratios[layer_id]
+        head_dim = args.head_dim
+        self.compress_ratio = compress_ratio
+        self.head_dim = head_dim
+        self.norm = RMSNorm(head_dim, args.norm_eps)
+        # ratio 1 is a plain projection, so it stays in the checkpoint's bf16; the softmax pooling
+        # above ratio 1 runs in fp32, so those weights are promoted to fp32 to match
+        self.wkv = Linear(args.dim, head_dim, dtype=torch.float32 if compress_ratio > 1 else torch.bfloat16)
+        if compress_ratio > 1:
+            self.wgate = Linear(args.dim, head_dim, dtype=torch.float32)
+            # tail of an incomplete group, carried across decode steps
+            self.kv_state: torch.Tensor
+            self.score_state: torch.Tensor
+            state_shape = (args.max_batch_size, compress_ratio, head_dim)
+            self.register_buffer("kv_state", torch.zeros(state_shape, dtype=torch.float32), persistent=False)
+            self.register_buffer(
+                "score_state", torch.full(state_shape, -torch.inf, dtype=torch.float32), persistent=False
+            )
+
+    def forward(self, x: torch.Tensor, start_pos: int) -> torch.Tensor | None:
+        bsz, seqlen, _ = x.size()
+        ratio, dtype = self.compress_ratio, x.dtype
+        if ratio == 1:  # one token per group: nothing to pool, so no gate and no fp32
+            return self.norm(self.wkv(x))
+
+        x = x.float()
+        kv, score = self.wkv(x), self.wgate(x)
+        if start_pos == 0:
+            should_compress = seqlen >= ratio
+            remainder = seqlen % ratio
+            cutoff = seqlen - remainder
+            if remainder:  # trailing partial group waits in the state
+                kv, self.kv_state[:bsz, :remainder] = kv.split([cutoff, remainder], dim=1)
+                score, self.score_state[:bsz, :remainder] = score.split([cutoff, remainder], dim=1)
+            kv = kv.unflatten(1, (-1, ratio))
+            score = score.unflatten(1, (-1, ratio))
+            kv = (kv * score.softmax(dim=2)).sum(dim=2)
+        else:  # one token per step: fill a slot, and pool only when the group just completed
+            should_compress = (start_pos + 1) % ratio == 0
+            slot = start_pos % ratio
+            self.kv_state[:bsz, slot] = kv.squeeze(1)
+            self.score_state[:bsz, slot] = score.squeeze(1)
+            if should_compress:
+                kv = (self.kv_state[:bsz] * self.score_state[:bsz].softmax(dim=1)).sum(dim=1, keepdim=True)
+        if not should_compress:
+            return None
+        return self.norm(kv.to(dtype))
+
+
+class Indexer(torch.nn.Module):
+    """Keeps the `index_topk` best compressed positions per query.
+
+    A small side attention: fp4 query heads against one shared key per compressed position, scores
+    rectified then combined by `weights_proj`. With a candidate source this is the second of two
+    levels; `select_candidate_blocks` is the first.
+    """
+
+    def __init__(self, args: ModelArgs, layer_id: int):
+        super().__init__()
+        # the index keys are derived from the compressor's latent, so only a layer that compresses
+        # its own KV can produce them; every other indexer reads them from that layer's cache
+        self.owns_k = layer_id in args.kv_source_layers
+        self.compress_ratio = args.compress_ratios[layer_id]
+        self.is_candidate_source = layer_id == args.candidate_source_layer
+        self.uses_candidates = 0 <= args.candidate_source_layer < layer_id
+        self.candidate_topk_blocks = args.candidate_topk_blocks
+        self.candidate_block_size = args.candidate_block_size
+        self.dim = args.dim
+        self.n_heads = args.index_n_heads
+        self.n_local_heads = args.index_n_heads // world_size
+        self.index_head_dim = args.index_head_dim
+        self.rope_head_dim = args.rope_head_dim
+        self.index_topk = args.index_topk
+        self.q_lora_rank = args.q_lora_rank
+        self.softmax_scale = self.index_head_dim**-0.5
+        self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.index_head_dim)
+        self.weights_proj = ColumnParallelLinear(self.dim, self.n_heads, dtype=torch.bfloat16)
+        self.freqs_cis: torch.Tensor | None = None
+        if self.owns_k:
+            self.wk = Linear(args.head_dim, self.index_head_dim, dtype=torch.bfloat16)
+            self.k_norm = RMSNorm(self.index_head_dim, args.norm_eps)
+            self.k_cache: torch.Tensor
+            self.register_buffer(
+                "k_cache",
+                torch.zeros(args.max_batch_size, args.max_seq_len // self.compress_ratio, args.index_head_dim),
+                persistent=False,
+            )
+
+    def forward(self, x: torch.Tensor, qr: torch.Tensor, latent: torch.Tensor, start_pos: int, offset: int):
+        """`latent` is this layer's RoPE-free compressed latent, None when this layer does not
+        compress or when its current group is still incomplete. An index-key owner turns it into
+        index keys here, which has to happen before Attention overwrites that same storage with
+        the RoPE'd, quantized values."""
+        assert self.freqs_cis is not None
+        bsz, seqlen, _ = x.size()
+        ratio, rd, end_pos = self.compress_ratio, self.rope_head_dim, start_pos + seqlen
+
+        # latent is None while a group is still filling up, so there is nothing to publish yet
+        if self.owns_k and latent is not None:
+            # a latent stands for the first token of its group, so group j takes position j * ratio
+            freqs = (
+                self.freqs_cis[: seqlen - seqlen % ratio : ratio]
+                if start_pos == 0
+                else self.freqs_cis[start_pos + 1 - ratio].unsqueeze(0)
+            )
+            k = self.k_norm(self.wk(latent))
+            apply_rotary_emb(k[..., -rd:], freqs)
+            fp4_act_quant(k, fp4_block_size, True)
+            self.k_cache[:bsz, start_pos // ratio : start_pos // ratio + k.size(1)] = k
+            shared_attn.index_k = self.k_cache
+
+        q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.index_head_dim))
+        apply_rotary_emb(q[..., -rd:], self.freqs_cis[start_pos:end_pos])
+        fp4_act_quant(q, fp4_block_size, True)
+
+        index_k = shared_attn.index_k[:bsz, : end_pos // ratio]
+        weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads**-0.5)
+        index_score = torch.einsum("bshd,btd->bsht", q, index_k)
+        index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2)
+        if world_size > 1:
+            dist.all_reduce(index_score)
+
+        # how many compressed positions each query can see: a block becomes visible once the query
+        # has passed its last token. One query per decode step, so there it is just a number.
+        if start_pos == 0:
+            compress_lens = (torch.arange(1, seqlen + 1, device=x.device) // ratio).unsqueeze(-1)
+            index_score.masked_fill_(torch.arange(seqlen // ratio, device=x.device) >= compress_lens, -torch.inf)
+        else:
+            compress_lens = end_pos // ratio
+
+        if self.is_candidate_source:
+            shared_attn.candidates = select_candidate_blocks(
+                index_score, compress_lens, self.candidate_topk_blocks, self.candidate_block_size
+            )
+        elif self.uses_candidates:
+            # level two: score with our own weights, but only inside the source's candidate blocks
+            index_score = index_score.masked_fill(~shared_attn.candidates, -torch.inf)
+
+        # top-k by score, re-sorted into position order; unreachable -> -1, rest shifted by offset
+        topk = min(self.index_topk, end_pos // ratio)
+        idxs = index_score.topk(topk, dim=-1, sorted=False).indices.sort(dim=-1).values
+        return torch.where(idxs < compress_lens, idxs + offset, -1).int()
+
+
+def select_candidate_blocks(
+    logits: torch.Tensor,
+    compress_lens: torch.Tensor | int,
+    topk_blocks: int,
+    block_size: int,
+) -> torch.Tensor:
+    """Level one of the two-level top-k: keep the `topk_blocks` highest-scoring blocks per query.
+
+    `logits` is [..., n_positions] with positions the query cannot reach already at -inf, which is
+    what makes a block score of -inf mean "not reachable yet". `compress_lens` is a plain int during
+    decode, or broadcasts against logits' leading dims during prefill. Returns a bool mask shaped
+    like `logits`, so the layers consuming it just mask and never think about blocks again.
+    """
+    width = logits.size(-1)
+    # score each block by its best position; -inf pads the last one out to block_size
+    scores = F.pad(logits, (0, -width % block_size), value=-torch.inf)
+    scores = scores.unflatten(-1, (-1, block_size)).amax(dim=-1)
+    num_blocks = scores.size(-1)
+
+    # the block with this query's newest position is only partly filled, so pin it in: it holds the
+    # most recent tokens but could otherwise be outscored by an older, full block
+    last = (compress_lens - 1) // block_size
+    scores = scores.masked_fill(torch.arange(num_blocks, device=logits.device) == last, torch.inf)
+
+    top = scores.topk(min(topk_blocks, num_blocks), dim=-1)
+    # fewer reachable blocks than topk_blocks means leftover picks came back -inf: drop them
+    keep = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, top.indices, top.values > -torch.inf)
+    return keep.repeat_interleave(block_size, dim=-1)[..., :width]
+
+
+class Attention(nn.Module):
+    """Latent attention over two KV sources at once, concatenated into one `sparse_attn` call: a
+    sliding window of raw KV, plus -- when compress_ratio > 0 -- `index_topk` compressed positions
+    reaching further back. Q and the output projection are both low-rank, the latter grouped.
+
+    compress_ratio > 0 does not mean the layer compresses its own KV: only kv_source_layers do,
+    the rest read that same cache.
+    """
+
+    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.n_local_heads = args.n_heads // world_size
+        self.q_lora_rank = args.q_lora_rank
+        self.o_lora_rank = args.o_lora_rank
+        self.head_dim = args.head_dim
+        self.rope_head_dim = args.rope_head_dim
+        self.nope_head_dim = args.head_dim - args.rope_head_dim
+        self.n_groups = args.o_groups
+        self.n_local_groups = self.n_groups // world_size
+        self.window_size = args.window_size
+        self.compress_ratio = args.compress_ratios[layer_id]
+        self.eps = args.norm_eps
+
+        self.attn_sink = nn.Parameter(torch.empty(self.n_local_heads, dtype=torch.float32))
+        self.wq_a = Linear(self.dim, self.q_lora_rank)
+        self.q_norm = RMSNorm(self.q_lora_rank, self.eps)
+        self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.head_dim)
+        self.wkv = Linear(self.dim, self.head_dim)
+        self.kv_norm = RMSNorm(self.head_dim, self.eps)
+        self.wo_a = ColumnParallelLinear(
+            self.n_heads * self.head_dim // self.n_groups,
+            self.n_groups * args.o_lora_rank,
+            dtype=torch.bfloat16,
+        )
+        self.wo_b = RowParallelLinear(self.n_groups * args.o_lora_rank, self.dim)
+        self.softmax_scale = self.head_dim**-0.5
+
+        is_backbone = layer_id < args.n_layers
+        self.is_kv_source = is_backbone and layer_id in args.kv_source_layers
+        self.is_index_source = is_backbone and layer_id in args.index_source_layers
+        self.compressor: Compressor | None = None
+        self.indexer: Indexer | None = None
+        if self.is_kv_source:
+            self.compressor = Compressor(args, layer_id)
+        if self.is_index_source:
+            self.indexer = Indexer(args, layer_id)
+
+        self.window_kv_cache: torch.Tensor
+        self.register_buffer(
+            "window_kv_cache",
+            torch.zeros(args.max_batch_size, args.window_size, self.head_dim),
+            persistent=False,
+        )
+        if self.is_kv_source:
+            self.compress_kv_cache: torch.Tensor
+            self.register_buffer(
+                "compress_kv_cache",
+                torch.zeros(
+                    args.max_batch_size,
+                    args.max_seq_len // self.compress_ratio,
+                    self.head_dim,
+                ),
+                persistent=False,
+            )
+        if self.compress_ratio:
+            original_seq_len, rope_theta = (
+                args.original_seq_len,
+                args.compress_rope_theta,
+            )
+        else:
+            # disable YaRN and use base rope_theta in pure sliding-window attention
+            original_seq_len, rope_theta = 0, args.rope_theta
+        freqs_cis = precompute_freqs_cis(
+            self.rope_head_dim,
+            args.max_seq_len,
+            original_seq_len,
+            rope_theta,
+            args.rope_factor,
+            args.beta_fast,
+            args.beta_slow,
+        )
+        self.freqs_cis: torch.Tensor
+        self.register_buffer("freqs_cis", freqs_cis, persistent=False)
+
+    def _window_kv(self, x, freqs_cis, start_pos):
+        """This layer's sliding-window K and the window positions every query may attend to. The K
+        stays fp8, quantized over the whole post-RoPE vector, RoPE tail included."""
+        bsz, seqlen, _ = x.size()
+        win = self.window_size
+        kv = self.kv_norm(self.wkv(x))
+        apply_rotary_emb(kv[..., -self.rope_head_dim :], freqs_cis)
+        act_quant(kv, fp8_block_size, scale_fmt, scale_dtype, True)
+        if start_pos == 0:  # prefill: attend over this chunk, seeding the ring buffer for decode
+            if seqlen <= win:
+                self.window_kv_cache[:bsz, :seqlen] = kv
+            else:
+                cutoff = seqlen % win
+                self.window_kv_cache[:bsz, cutoff:win], self.window_kv_cache[:bsz, :cutoff] = kv[:, -win:].split(
+                    [win - cutoff, cutoff], dim=1
+                )
+            window_kv = kv
+        else:  # decode: one token into the ring buffer, attend over the whole window
+            self.window_kv_cache[:bsz, start_pos % win] = kv.squeeze(1)
+            window_kv = self.window_kv_cache[:bsz]
+        return window_kv, get_window_topk_idxs(win, bsz, seqlen, start_pos)
+
+    def _compress_topk_idxs(self, x, qr, latent, start_pos, offset, compress_len):
+        """Which compressed positions each query attends to. Index sources run their own indexer;
+        the layers in between reuse the result their source published."""
+        if not self.is_index_source:
+            return shared_attn.topk_idxs
+
+        bsz, seqlen, _ = x.size()
+        if compress_len == 0:
+            idxs = torch.empty(bsz, seqlen, 0, dtype=torch.int32, device=x.device)
+        else:
+            assert self.indexer is not None
+            if self.indexer.freqs_cis is None:
+                self.indexer.freqs_cis = self.freqs_cis
+            idxs = self.indexer(x, qr, latent, start_pos, offset)
+        shared_attn.topk_idxs = idxs
+        return idxs
+
+    def _compress_kv(self, x, qr, start_pos, offset):
+        """The shared compressed KV and the compressed positions every query may attend to. This
+        layer compresses its own KV only when it is a source; otherwise it just reads the cache."""
+        bsz, seqlen, _ = x.size()
+        ratio = self.compress_ratio
+        compress_len = (start_pos + seqlen) // ratio
+        latent = None
+        if self.is_kv_source:
+            latent = self.compressor(x, start_pos)
+            shared_attn.compress_kv = self.compress_kv_cache
+        # the indexer needs the latent before RoPE, so it runs before the cache is written
+        idxs = self._compress_topk_idxs(x, qr, latent, start_pos, offset, compress_len)
+        if latent is not None:
+            # a latent stands for the first token of its group, so group j takes position j * ratio
+            freqs = (
+                self.freqs_cis[: seqlen - seqlen % ratio : ratio]
+                if start_pos == 0
+                else self.freqs_cis[start_pos + 1 - ratio].unsqueeze(0)
+            )
+            apply_rotary_emb(latent[..., -self.rope_head_dim :], freqs)
+            # Compressed KV uses groups of 16 with E4M3 scales; the indexer uses 32 with E8M0.
+            fp4_act_quant(latent, 16, True, scale_dtype=torch.float8_e4m3fn)
+            self.compress_kv_cache[:bsz, start_pos // ratio : start_pos // ratio + latent.size(1)] = latent
+        # read after the write, so this does not depend on the slice aliasing the cache
+        return shared_attn.compress_kv[:bsz, :compress_len], idxs
+
+    def forward(self, x: torch.Tensor, start_pos: int):
+        bsz, seqlen, _ = x.size()
+        freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]
+        rd = self.rope_head_dim
+
+        qr = self.q_norm(self.wq_a(x))
+        q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.head_dim))
+        apply_rotary_emb(q[..., -rd:], freqs_cis)
+
+        kv, topk_idxs = self._window_kv(x, freqs_cis, start_pos)
+        if self.compress_ratio:
+            compress_kv, compress_idxs = self._compress_kv(x, qr, start_pos, kv.size(1))
+            kv = torch.cat([kv, compress_kv], dim=1)
+            topk_idxs = torch.cat([topk_idxs, compress_idxs], dim=-1)
+
+        o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale)
+        apply_rotary_emb(o[..., -rd:], freqs_cis, True)
+
+        # wo_a is block-diagonal over groups (each projects only its own heads), hence einsum not
+        # Linear. convert.py dequantizes it to bf16; an fp8 grouped GEMM would halve the memory.
+        o = o.view(bsz, seqlen, self.n_local_groups, -1)
+        wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1)
+        o = torch.einsum("bsgd,grd->bsgr", o, wo_a)
+        x = self.wo_b(o.flatten(2))
+        return x
+
+
+class Gate(nn.Module):
+    """MoE gating. The correction bias steers expert selection only; the routing weights come from the
+    unbiased scores. Image-span tokens use a separate bias (training `noaux_tc_for_vl`)."""
+
+    def __init__(self, layer_id: int, args: ModelArgs):
+        super().__init__()
+        n_routed_experts, n_activated_experts = args.get_moe_config(layer_id)
+        self.dim = args.dim
+        self.topk = n_activated_experts
+        self.score_func = args.score_func
+        self.gate_temp = args.gate_temp
+        self.norm_topk_prob = args.norm_topk_prob
+        self.route_scale = args.route_scale
+        self.weight = nn.Parameter(torch.empty(n_routed_experts, args.dim))
+        self.bias = nn.Parameter(torch.empty(n_routed_experts, dtype=torch.float32))
+        self.bias_vl = nn.Parameter(torch.empty(n_routed_experts, dtype=torch.float32)) if args.vision_enabled else None
+
+    def forward(self, x: torch.Tensor, image_mask: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]:
+        """x: [n, dim]; image_mask: [n] bool, True for tokens inside an image span."""
+        scores = linear(x.float(), self.weight.float()) / self.gate_temp
+        if self.score_func == "softmax":
+            scores = scores.softmax(dim=-1)
+        elif self.score_func == "sigmoid":
+            scores = scores.sigmoid()
+        else:
+            scores = F.softplus(scores).sqrt()
+        bias = self.bias
+        if image_mask is not None and self.bias_vl is not None:
+            bias = torch.where(image_mask.unsqueeze(-1), self.bias_vl, bias)
+        # the bias picks experts but does not scale them: weights come from the raw scores
+        indices = (scores + bias).topk(self.topk, dim=-1)[1]
+        weights = scores.gather(1, indices)
+        if self.norm_topk_prob and self.topk > 1:
+            weights /= weights.sum(dim=-1, keepdim=True) + 1e-20  # not norm_eps, matches training
+        weights *= self.route_scale
+        return weights, indices
+
+
+class Expert(nn.Module):
+    """One SwiGLU FFN. The clamps come straight from training, where they keep fp8/fp4 activations in
+    range: the up branch is clamped on both sides, the gate branch only from above."""
+
+    def __init__(self, dim: int, inter_dim: int, dtype=None, swiglu_limit=0.0):
+        super().__init__()
+        self.w1 = Linear(dim, inter_dim, dtype=dtype)
+        self.w2 = Linear(inter_dim, dim, dtype=dtype)
+        self.w3 = Linear(dim, inter_dim, dtype=dtype)
+        self.swiglu_limit = swiglu_limit
+
+    def forward(self, x: torch.Tensor, weights: torch.Tensor | None = None) -> torch.Tensor:
+        dtype = x.dtype
+        gate = self.w1(x).float()
+        up = self.w3(x).float()
+        if self.swiglu_limit > 0:
+            up = torch.clamp(up, min=-self.swiglu_limit, max=self.swiglu_limit)
+            gate = torch.clamp(gate, max=self.swiglu_limit)
+        x = F.silu(gate) * up
+        if weights is not None:
+            x = weights * x
+        return self.w2(x.to(dtype))
+
+
+class MoE(nn.Module):
+    """Top-k routed experts plus one shared expert every token goes through. Experts are split
+    across ranks, so `self.experts` is None for those another rank owns."""
+
+    def __init__(self, layer_id: int, args: ModelArgs):
+        super().__init__()
+        n_routed_experts, n_activated_experts = args.get_moe_config(layer_id)
+        self.layer_id = layer_id
+        self.dim = args.dim
+        assert n_routed_experts % world_size == 0, (
+            f"Number of experts must be divisible by world size (world_size={world_size})"
+        )
+        self.n_routed_experts = n_routed_experts
+        self.n_local_experts = n_routed_experts // world_size
+        self.n_activated_experts = n_activated_experts
+        self.experts_start_idx = rank * self.n_local_experts
+        self.experts_end_idx = self.experts_start_idx + self.n_local_experts
+        self.gate = Gate(layer_id, args)
+        expert_dtype = torch.float4_e2m1fn_x2 if args.expert_dtype == "fp4" else None
+        self.experts = nn.ModuleList(
+            [
+                Expert(
+                    args.dim,
+                    args.moe_inter_dim,
+                    dtype=expert_dtype,
+                    swiglu_limit=args.swiglu_limit,
+                )
+                if self.experts_start_idx <= i < self.experts_end_idx
+                else None
+                for i in range(self.n_routed_experts)
+            ]
+        )
+        assert args.n_shared_experts == 1
+        self.shared_experts = Expert(args.dim, args.moe_inter_dim, swiglu_limit=args.swiglu_limit)
+
+    def forward(self, x: torch.Tensor, image_mask: torch.Tensor | None = None) -> torch.Tensor:
+        shape = x.size()
+        x = x.view(-1, self.dim)
+        weights, indices = self.gate(x, None if image_mask is None else image_mask.flatten())
+        y = torch.zeros_like(x, dtype=torch.float32)
+        counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist()
+        for i in range(self.experts_start_idx, self.experts_end_idx):
+            if counts[i] == 0:
+                continue
+            expert = self.experts[i]
+            idx, top = torch.where(indices == i)
+            y[idx] += expert(x[idx], weights[idx, top, None])
+        if world_size > 1:
+            dist.all_reduce(y)
+        y += self.shared_experts(x)
+        return y.type_as(x).view(shape)
+
+
+class Block(nn.Module):
+    """A block whose residual stream is `hc_mult` parallel copies (Hyper-Connections).
+
+    Attention and FFN each sit between `hc_pre` (collapse the copies into one sublayer input) and
+    `hc_post` (expand back out, mixing the residual in through `comb`). `hc_mixes` derives all three
+    coefficient sets from the stream itself, `comb` made doubly stochastic by Sinkhorn.
+
+    The coefficients a sublayer computes are used by the *next* one -- see `forward`.
+    """
+
+    attention_cls = Attention
+
+    def __init__(
+        self,
+        layer_id: int,
+        args: ModelArgs,
+        engram_layout: EngramLayout | None = None,
+    ):
+        super().__init__()
+        self.layer_id = layer_id
+        self.norm_eps = args.norm_eps
+        self.attn = self.attention_cls(layer_id, args)
+        self.ffn = MoE(layer_id, args)
+        self.engram = None
+        if engram_layout is not None and layer_id in engram_layout.layer_ids:
+            self.engram = Engram(args, layer_id, engram_layout)
+        self.attn_norm = RMSNorm(args.dim, self.norm_eps)
+        self.ffn_norm = RMSNorm(args.dim, self.norm_eps)
+        self.hc_mult = hc_mult = args.hc_mult
+        self.hc_sinkhorn_iters = args.hc_sinkhorn_iters
+        self.hc_eps = args.hc_eps
+        mix_hc = (2 + hc_mult) * hc_mult
+        hc_dim = hc_mult * args.dim
+        with set_dtype(torch.float32):
+            self.hc_attn_fn = nn.Parameter(torch.empty(mix_hc, hc_dim))
+            self.hc_ffn_fn = nn.Parameter(torch.empty(mix_hc, hc_dim))
+            self.hc_attn_base = nn.Parameter(torch.empty(mix_hc))
+            self.hc_ffn_base = nn.Parameter(torch.empty(mix_hc))
+            self.hc_attn_scale = nn.Parameter(torch.empty(3))
+            self.hc_ffn_scale = nn.Parameter(torch.empty(3))
+
+    def hc_mixes(self, x: torch.Tensor, hc_fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor):
+        """x: [b,s,hc,d], hc_fn: [mix_hc, hc*d], hc_scale: [3], hc_base: [mix_hc]. Returns the
+        pre / post / comb coefficients, split out of one projection of the flattened stream."""
+        # normalized over the whole flattened hc*d stream, one statistic per token
+        x = x.flatten(2).float()
+        rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + self.norm_eps)
+        mixes = F.linear(x, hc_fn) * rsqrt
+        return hc_split_sinkhorn(mixes, hc_scale, hc_base, self.hc_mult, self.hc_sinkhorn_iters, self.hc_eps)
+
+    def hc_pre(self, x: torch.Tensor, pre_mix: torch.Tensor):
+        """Collapse the hc copies into one, weighted by pre_mix. [b,s,hc,d] x [b,s,hc] -> [b,s,d]"""
+        y = torch.sum(pre_mix.unsqueeze(-1) * x.float(), dim=2)
+        return y.to(x.dtype)
+
+    def hc_post(self, x: torch.Tensor, residual: torch.Tensor, post: torch.Tensor, comb: torch.Tensor):
+        """Expand the sublayer output back to hc copies and mix the residual in through `comb`.
+        x: [b,s,d], residual: [b,s,hc,d], post: [b,s,hc], comb: [b,s,hc,hc] -> [b,s,hc,d]"""
+        y = post.unsqueeze(-1) * x.unsqueeze(-2) + torch.sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim=2)
+        return y.type_as(x)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        pre_mix: torch.Tensor,
+        image_mask: torch.Tensor | None,
+        *attn_args,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """`pre_mix` collapses the hc_mult copies down to one input for this block's attention. Each
+        sub-block's own `hc_mixes` produces the mix for the *next* one, so attention uses what the
+        previous layer's FFN produced and the FFN uses what this attention produced.
+
+        image_mask: [b, s] bool, True inside image spans (selects the VL routing bias)."""
+        residual = x
+        attn_pre, attn_post, attn_comb = self.hc_mixes(x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base)
+        x = self.hc_pre(x, pre_mix)
+        x = self.attn_norm(x)
+        x = self.attn(x, start_pos, *attn_args)
+        x = self.hc_post(x, residual, attn_post, attn_comb)
+
+        residual = x
+        ffn_pre, ffn_post, ffn_comb = self.hc_mixes(x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base)
+        x = self.hc_pre(x, attn_pre)
+        x = self.ffn_norm(x)
+        x = self.ffn(x, image_mask)
+        x = self.hc_post(x, residual, ffn_post, ffn_comb)
+        return x, ffn_pre
+
+
+class ParallelHead(nn.Module):
+    def __init__(self, vocab_size: int, dim: int, norm_eps: float = 1e-6, hc_eps: float = 1e-6):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.dim = dim
+        self.norm_eps = norm_eps
+        self.hc_eps = hc_eps
+        self.part_vocab_size = vocab_size // world_size
+        # bf16 in the checkpoint, kept as fp32 here so the logits come out in fp32 directly
+        self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim, dtype=torch.float32))
+
+    def forward(self, x: torch.Tensor, full_logits=False):
+        """x: [b, s, d]. Generation only needs the last position, so that is the default."""
+        if not full_logits:
+            x = x[:, -1]
+        logits = F.linear(x.float(), self.weight)
+        if world_size > 1:
+            all_logits = [torch.empty_like(logits) for _ in range(world_size)]
+            dist.all_gather(all_logits, logits)
+            logits = torch.cat(all_logits, dim=-1)
+        return logits
+
+
+@lru_cache(1)
+def get_dspark_topk_idxs(window_size: int, bsz: int, block_size: int, start_pos: int):
+    assert start_pos > 0
+    matrix = torch.cat(
+        [
+            torch.arange(min(window_size, start_pos + 1)),
+            window_size + torch.arange(block_size),
+        ]
+    )
+    return matrix.int().view(1, 1, -1).expand(bsz, block_size, -1).contiguous()
+
+
+class DSparkAttention(Attention):
+    def forward(self, x: torch.Tensor, start_pos: int, main_x: torch.Tensor):
+        assert self.compress_ratio == 0
+        bsz, seqlen, _ = main_x.size()
+        win = self.window_size
+        rd = self.rope_head_dim
+
+        main_freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]
+        main_kv = self.kv_norm(self.wkv(main_x))
+        apply_rotary_emb(main_kv[..., -rd:], main_freqs_cis)
+        act_quant(main_kv, fp8_block_size, scale_fmt, scale_dtype, True)
+
+        if start_pos == 0:
+            if seqlen <= win:
+                self.window_kv_cache[:bsz, :seqlen] = main_kv
+            else:
+                cutoff = seqlen % win
+                self.window_kv_cache[:bsz, cutoff:win], self.window_kv_cache[:bsz, :cutoff] = main_kv[:, -win:].split(
+                    [win - cutoff, cutoff], dim=1
+                )
+            return x
+
+        bsz, block_size, _ = x.size()
+        freqs_cis = self.freqs_cis[start_pos + seqlen : start_pos + seqlen + block_size]
+
+        qr = self.q_norm(self.wq_a(x))
+        q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.head_dim))
+        apply_rotary_emb(q[..., -rd:], freqs_cis)
+        kv = self.kv_norm(self.wkv(x))
+        apply_rotary_emb(kv[..., -rd:], freqs_cis)
+        act_quant(kv, fp8_block_size, scale_fmt, scale_dtype, True)
+
+        topk_idxs = get_dspark_topk_idxs(win, bsz, block_size, start_pos)
+        self.window_kv_cache[:bsz, start_pos % win] = main_kv.squeeze(1)
+        kv = torch.cat([self.window_kv_cache[:bsz], kv], dim=1)
+        o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale)
+        apply_rotary_emb(o[..., -rd:], freqs_cis, True)
+
+        o = o.view(bsz, block_size, self.n_local_groups, -1)
+        wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1)
+        o = torch.einsum("bsgd,grd->bsgr", o, wo_a)
+        x = self.wo_b(o.flatten(2))
+        return x
+
+
+class DSparkMarkovHead(nn.Module):
+    def __init__(self, vocab_size: int, dspark_markov_rank: int):
+        super().__init__()
+        self.embed = ParallelEmbedding(vocab_size, dspark_markov_rank)
+        self.head = ParallelHead(vocab_size, dspark_markov_rank)
+
+    def forward(self, token_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        embed = self.embed(token_ids)
+        logits = self.head(embed, full_logits=True)
+        return logits, embed
+
+
+class DSparkConfidenceHead(nn.Module):
+    def __init__(self, input_dim: int):
+        super().__init__()
+        # proj in the checkpoint is stored in bf16, while the parameter here is stored in fp32 for fp32 confidence score.
+        self.proj = Linear(input_dim, 1, dtype=torch.float32)
+
+    def forward(self, hidden: torch.Tensor, markov_embed: torch.Tensor):
+        hidden = torch.cat([hidden, markov_embed], dim=-1)
+        return self.proj(hidden.float()).squeeze(-1)
+
+
+class DSparkBlock(Block):
+    """DSpark stage stored under the mtp.* checkpoint namespace."""
+
+    attention_cls = DSparkAttention
+
+    def __init__(self, layer_id: int, args: ModelArgs):
+        super().__init__(layer_id, args)
+        stage_id = layer_id - args.n_layers
+        self.block_size = args.dspark_block_size
+        self.noise_token_id = args.dspark_noise_token_id
+        self.temperature = args.temperature
+        if stage_id == 0:
+            assert len(args.dspark_target_layer_ids) > 0, "DSpark needs target layers"
+            self.main_proj = Linear(args.dim * len(args.dspark_target_layer_ids), args.dim)
+            self.main_norm = RMSNorm(args.dim, args.norm_eps)
+        if stage_id == args.n_mtp_layers - 1:
+            self.norm = RMSNorm(args.dim, args.norm_eps)
+            self.markov_head = DSparkMarkovHead(args.vocab_size, args.dspark_markov_rank)
+            self.confidence_head = DSparkConfidenceHead(args.dim + args.dspark_markov_rank)
+        self.embed: ParallelEmbedding | None = None
+        self.head: ParallelHead | None = None
+
+    def forward(self, x: torch.Tensor, start_pos: int, pre_mix: torch.Tensor, main_x: torch.Tensor):
+        if start_pos == 0:
+            self.attn(x, start_pos, main_x)  # prefill only seeds the window KV cache
+            return x, pre_mix
+        return super().forward(x, start_pos, pre_mix, None, main_x)  # drafts are text: no VL bias
+
+    def forward_embed(self, main_hidden: torch.Tensor, input_ids: torch.Tensor):
+        assert self.embed is not None
+        main_x = self.main_norm(self.main_proj(main_hidden))
+        draft_input_ids = input_ids.new_full([input_ids.size(0), self.block_size], self.noise_token_id)
+        draft_input_ids[:, 0] = input_ids
+        x = self.embed(draft_input_ids)
+        x = x.unsqueeze(2).repeat(1, 1, self.hc_mult, 1)
+        return x, main_x
+
+    def forward_head(
+        self,
+        x: torch.Tensor,
+        pre_mix: torch.Tensor,
+        input_ids: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        assert self.head is not None
+        x = self.hc_pre(x, pre_mix)
+        logits = self.head(self.norm(x), full_logits=True)
+        output_ids = input_ids.new_empty(input_ids.size(0), self.block_size + 1)
+        output_ids[:, 0] = input_ids
+        markov_embeds = []
+        for i in range(self.block_size):
+            logits_bias, markov_embed = self.markov_head(output_ids[:, i])
+            logits[:, i].add_(logits_bias)
+            markov_embeds.append(markov_embed)
+            output_ids[:, i + 1] = sample(logits[:, i], self.temperature)
+        markov_embed = torch.stack(markov_embeds, dim=1)
+        confidence = self.confidence_head(x, markov_embed)
+        return output_ids, logits, confidence
+
+
+def make_identity_pre_mix(x: torch.Tensor, hc_mult: int) -> torch.Tensor:
+    """initial one-hot mix"""
+    pre_mix = x.new_zeros(x.size(0), x.size(1), hc_mult, dtype=torch.float32)
+    pre_mix[:, :, 0] = 1.0
+    return pre_mix
+
+
+class SharedAttentionRuntime:
+    """What attention layers hand down the stack instead of recomputing. Layers run in order and
+    every source writes before its consumers read, so one slot each is enough and nothing needs
+    resetting between forwards. Sources: compress_kv and index_k from kv_source_layers, topk_idxs
+    from index_source_layers, candidates from candidate_source_layer."""
+
+    def __init__(self):
+        self.compress_kv: torch.Tensor | None = None
+        self.index_k: torch.Tensor | None = None
+        self.topk_idxs: torch.Tensor | None = None
+        self.candidates: torch.Tensor | None = None
+
+
+# Only ever one model per process, same as world_size / rank / default_dtype above.
+shared_attn = SharedAttentionRuntime()
+
+
+class Transformer(nn.Module):
+    """DeepSeek-V4.1: embed -> expand to hc_mult copies -> blocks -> collapse -> logits. Building
+    this sets the globals at the top of the file. The tokenizer only feeds the engram token map."""
+
+    def __init__(self, args: ModelArgs, tokenizer=None):
+        global world_size, rank, default_dtype
+        world_size = dist.get_world_size() if dist.is_initialized() else 1
+        rank = dist.get_rank() if dist.is_initialized() else 0
+        default_dtype = torch.float8_e4m3fn if args.dtype == "fp8" else torch.bfloat16
+        super().__init__()
+        self.max_seq_len = args.max_seq_len
+        self.temperature = args.temperature
+        self.norm_eps = args.norm_eps
+        self.hc_eps = args.hc_eps
+        self.engram_layout = EngramLayout.from_args(args)
+        self.engram_hash = (
+            NgramHashState(args, self.engram_layout, tokenizer) if self.engram_layout is not None else None
+        )
+        self.embed = ParallelEmbedding(args.vocab_size, args.dim)
+        self.layers = torch.nn.ModuleList()
+        for layer_id in range(args.n_layers):
+            self.layers.append(Block(layer_id, args, self.engram_layout))
+        self.norm = RMSNorm(args.dim, self.norm_eps)
+        self.head = ParallelHead(args.vocab_size, args.dim, self.norm_eps, self.hc_eps)
+        self.mtp = torch.nn.ModuleList()
+        self.target_layer_ids = args.dspark_target_layer_ids
+        if args.dspark_block_size:
+            for layer_id in range(args.n_mtp_layers):
+                self.mtp.append(DSparkBlock(args.n_layers + layer_id, args))
+                self.mtp[-1].embed = self.embed
+                self.mtp[-1].head = self.head
+        self.hc_mult = args.hc_mult
+        self.vision = None
+        if args.vision_enabled:
+            self.vision = ViT(args)
+            self.aligner = Aligner(args)
+            # learned embeddings for the image span delimiters
+            self.image_start = nn.Parameter(torch.empty(args.dim))
+            self.image_end = nn.Parameter(torch.empty(args.dim))
+            self.image_newline = nn.Parameter(torch.empty(args.dim))
+
+    @torch.inference_mode()
+    def encode_image(self, patches: torch.Tensor, n_vit_h: int, n_vit_w: int) -> torch.Tensor:
+        return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w)
+
+    def merge_image_embeddings(self, images, h: torch.Tensor):
+        """Overwrite each image's token span in h with its ViT/aligner features. The IMAGE slots take
+        the aligner rows in row-major order; the span delimiters take learned embeddings."""
+        for i, sample in enumerate(images):
+            for img in sample or ():
+                types = img.types.to(h.device)
+                span = h[i, img.start : img.start + types.numel()]
+                span[types == IMAGE_START] = self.image_start.to(h.dtype)
+                span[types == IMAGE_END] = self.image_end.to(h.dtype)
+                span[types == IMAGE_NEW_LINE] = self.image_newline.to(h.dtype)
+                embeds = self.encode_image(img.patches.to(h.device), img.n_vit_h, img.n_vit_w)
+                span[types == IMAGE] = embeds.to(h.dtype)
+
+    @torch.inference_mode()
+    def forward(
+        self, input_ids: torch.Tensor, start_pos: int = 0, images=None, token_types: torch.Tensor | None = None
+    ):
+        """input_ids: [b, s], every entry a real token id -- generate.py only ever passes positions it
+        has already filled, so the padding it uses internally never reaches here. token_types /
+        images carry the VL inputs built by image_processor.prepare_vl_inputs; image spans must lie
+        inside the first (start_pos 0) chunk."""
+        image_mask = None if token_types is None else token_types >= 0  # TEXT is -1
+        # image tokens take no part in an n-gram and get no engram contribution; text-only needs no mask
+        engram_mask = None if image_mask is None else ~image_mask
+        engram_hashes = self.engram_hash(input_ids, start_pos, engram_mask) if self.engram_hash is not None else None
+        h = self.embed(input_ids)
+        if images is not None:
+            assert start_pos == 0, "image spans must be prefilled in a single chunk"
+            self.merge_image_embeddings(images, h)
+        # Expand to hc_mult copies for Hyper-Connections
+        h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1)
+        main_hiddens = []
+        pre_mix = make_identity_pre_mix(h, self.hc_mult)
+        for i, layer in enumerate(self.layers):
+            if layer.engram is not None:
+                h = layer.engram(h, engram_hashes[:, :, layer.engram.layer_hash_index, :], engram_mask)
+            # the MTP head reads the attention input of its target layers, not their output
+            if i in self.target_layer_ids:
+                main_hiddens.append(h.mean(dim=2))
+            h, pre_mix = layer(h, start_pos, pre_mix, image_mask)
+        h = layer.hc_pre(h, pre_mix)
+        logits = self.head(self.norm(h))
+        output_ids = sample(logits, self.temperature)
+        main_hidden = torch.cat(main_hiddens, dim=-1) if main_hiddens else None
+        return output_ids, logits, main_hidden
+
+    @torch.inference_mode()
+    def forward_spec(self, input_ids: torch.Tensor, main_hidden: torch.Tensor, start_pos: int = 0):
+        h, main_x = self.mtp[0].forward_embed(main_hidden, input_ids)
+        pre_mix = make_identity_pre_mix(h, self.hc_mult)
+        for layer in self.mtp:
+            h, pre_mix = layer(h, start_pos, pre_mix, main_x)
+        if start_pos == 0:
+            return None
+        return self.mtp[-1].forward_head(h, pre_mix, input_ids)
+
+
+def sample(logits, temperature: float = 1.0):
+    """Gumbel-max trick: equivalent to multinomial sampling but faster on GPU,
+    since it avoids the GPU-to-CPU sync in torch.multinomial."""
+    if temperature == 0:
+        return logits.argmax(dim=-1)
+    logits = logits / max(temperature, 1e-5)
+    probs = torch.softmax(logits, dim=-1, dtype=torch.float32)
+    return probs.div_(torch.empty_like(probs).exponential_(1)).argmax(dim=-1)
+
+
+if __name__ == "__main__":
+    torch.set_default_dtype(torch.bfloat16)
+    torch.set_default_device("cuda")
+    torch.manual_seed(0)
+    args = ModelArgs(dspark_block_size=6, dspark_target_layer_ids=(3, 4))
+    x = torch.randint(0, args.vocab_size, (2, 150))
+    model = Transformer(args)
+
+    output_ids, logits, main_hidden = model(x[:, :128])
+    model.forward_spec(output_ids, main_hidden)
+    for i in range(128, 150):
+        output_ids, logits, main_hidden = model(x[:, i : i + 1], i)
+        result = model.forward_spec(output_ids, main_hidden, i)
+        assert result is not None
+        output_ids, logits, confidence = result

inference/requirements.txt

@@ -0,0 +1,9 @@
+torch>=2.10.0
+transformers
+tokenizers
+safetensors>=0.7.0
+numpy
+sympy
+Pillow
+tilelang==0.1.8
+tqdm

inference/run.sh

@@ -0,0 +1,54 @@
+#!/usr/bin/env bash
+#
+# Run the reference inference on a converted checkpoint.
+#
+#   ./run.sh /path/to/DeepSeek-V4.1-Exp-TP8
+#   ./run.sh /path/to/DeepSeek-V4.1-Exp-TP8 examples/example_harmony.json
+#   MP=4 ./run.sh /path/to/DeepSeek-V4.1-Exp-TP4
+#
+# Paths inside an example are resolved from this directory, so run it from anywhere.
+
+set -euo pipefail
+cd "$(dirname "$0")"
+
+CKPT_PATH="${1:-${CKPT_PATH:-}}"
+INPUT_FILE="${2:-${INPUT_FILE:-examples/example_harmony.json}}"
+MP="${MP:-8}"
+CONFIG="${CONFIG:-config.json}"
+
+usage() {
+    echo "usage: $0 <checkpoint-dir> [input-file]" >&2
+    echo >&2
+    echo "  checkpoint-dir  holds model0-mp${MP}.safetensors .. model$((MP - 1))-mp${MP}.safetensors," >&2
+    echo "                  as produced by convert.py --model-parallel ${MP}" >&2
+    echo "  input-file      TXT or JSON prompts (default: examples/example.txt)" >&2
+    echo >&2
+    echo "  MP=${MP}  CONFIG=${CONFIG}   override with environment variables" >&2
+    exit 1
+}
+
+[ -n "${CKPT_PATH}" ] || usage
+
+if [ ! -d "${CKPT_PATH}" ]; then
+    echo "error: checkpoint directory not found: ${CKPT_PATH}" >&2
+    usage
+fi
+
+missing=0
+for rank in $(seq 0 $((MP - 1))); do
+    if [ ! -f "${CKPT_PATH}/model${rank}-mp${MP}.safetensors" ]; then
+        missing=$((missing + 1))
+    fi
+done
+if [ "${missing}" -ne 0 ]; then
+    echo "error: ${CKPT_PATH} is missing ${missing} of the ${MP} shards MP=${MP} needs" >&2
+    echo "       expected model0-mp${MP}.safetensors .. model$((MP - 1))-mp${MP}.safetensors" >&2
+    usage
+fi
+
+[ -f "${INPUT_FILE}" ] || { echo "error: input file not found: ${INPUT_FILE}" >&2; usage; }
+
+torchrun --nproc-per-node "${MP}" generate.py \
+    --ckpt-path "${CKPT_PATH}" \
+    --config "${CONFIG}" \
+    --input-file "${INPUT_FILE}"

inference/vision.py

@@ -0,0 +1,119 @@
+from functools import lru_cache
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+@lru_cache(8)
+def get_vision_cos_sin(n_h: int, n_w: int, dim: int, theta: float):
+    inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+    hpos = torch.arange(n_h).unsqueeze(1).expand(n_h, n_w)
+    wpos = torch.arange(n_w).unsqueeze(0).expand(n_h, n_w)
+    freqs = torch.stack([hpos, wpos], dim=-1).reshape(-1, 2, 1).float() * inv_freq
+    freqs = freqs.flatten(1)
+    return freqs.cos().unsqueeze(1), freqs.sin().unsqueeze(1)
+
+
+def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    dtype = x.dtype
+    x1, x2 = x.float().chunk(2, dim=-1)
+    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1).to(dtype)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        x = x.float()
+        x = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps)
+        return (self.weight * x).to(dtype)
+
+
+class PatchEmbed(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.proj = nn.Linear(3 * args.vision_patch_size**2, args.vision_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.proj(x.flatten(1))
+
+
+class Attention(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.n_heads = args.vision_n_heads
+        self.head_dim = args.vision_dim // args.vision_n_heads
+        self.wqkv = nn.Linear(args.vision_dim, 3 * args.vision_dim)
+        self.wo = nn.Linear(args.vision_dim, args.vision_dim)
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+        n = x.size(0)
+        q, k, v = (t.view(n, self.n_heads, self.head_dim) for t in self.wqkv(x).chunk(3, dim=-1))
+        q = apply_rotary(q, cos, sin)
+        k = apply_rotary(k, cos, sin)
+        o = F.scaled_dot_product_attention(q.transpose(0, 1), k.transpose(0, 1), v.transpose(0, 1))
+        return self.wo(o.transpose(0, 1).reshape(n, -1))
+
+
+class MLP(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.w1 = nn.Linear(args.vision_dim, 2 * args.vision_inter_dim, bias=False)
+        self.w2 = nn.Linear(args.vision_inter_dim, args.vision_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate, up = self.w1(x).chunk(2, dim=-1)
+        return self.w2(F.silu(gate) * up)
+
+
+class Block(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.norm1 = RMSNorm(args.vision_dim)
+        self.attn = Attention(args)
+        self.norm2 = RMSNorm(args.vision_dim)
+        self.mlp = MLP(args)
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.norm1(x), cos, sin)
+        return x + self.mlp(self.norm2(x))
+
+
+class ViT(nn.Module):
+    """DeepSeek ViT: full bidirectional attention over one image with 2D RoPE."""
+
+    def __init__(self, args):
+        super().__init__()
+        self.rope_dim = args.vision_dim // args.vision_n_heads // 2
+        self.rope_theta = args.vision_rope_theta
+        self.patch_embed = PatchEmbed(args)
+        self.blocks = nn.ModuleList([Block(args) for _ in range(args.vision_n_layers)])
+        self.norm = RMSNorm(args.vision_dim)
+
+    def forward(self, patches: torch.Tensor, n_h: int, n_w: int) -> torch.Tensor:
+        x = self.patch_embed(patches)
+        cos, sin = get_vision_cos_sin(n_h, n_w, self.rope_dim, self.rope_theta)
+        for block in self.blocks:
+            x = block(x, cos, sin)
+        return self.norm(x)
+
+
+class Aligner(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.downsample_ratio = args.vision_downsample_ratio
+        in_dim = args.vision_dim * self.downsample_ratio**2
+        self.w1 = nn.Linear(in_dim, args.dim)
+        self.w2 = nn.Linear(args.dim, args.dim)
+
+    def forward(self, x: torch.Tensor, n_h: int, n_w: int) -> torch.Tensor:
+        r = self.downsample_ratio
+        x = x.view(n_h, n_w, -1).permute(2, 0, 1)
+        x = F.pad(x, (0, -n_w % r, 0, -n_h % r))
+        x = F.unfold(x.unsqueeze(0), r, stride=r).squeeze(0).transpose(0, 1)
+        return self.w2(F.gelu(self.w1(x)))

model-00001-of-00048.safetensors

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

model-00002-of-00048.safetensors

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

model-00003-of-00048.safetensors

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

model-00004-of-00048.safetensors

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

model-00005-of-00048.safetensors

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

model-00006-of-00048.safetensors

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

model-00007-of-00048.safetensors

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

model-00008-of-00048.safetensors

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

model-00009-of-00048.safetensors

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

model-00010-of-00048.safetensors

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

model-00011-of-00048.safetensors

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

model-00012-of-00048.safetensors

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

model-00013-of-00048.safetensors

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