init

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* 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

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,44 @@
+---
+license: mit
+library_name: transformers
+---
+
+# DeepSeek-V4.1-Exp
+
+Prompt encoding reference and a minimal PyTorch inference implementation for
+DeepSeek-V4.1. The inference 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.
+
+## Repository layout
+
+```text
+.
+├── encoding/                  # OpenAI-style messages -> model prompt
+│   └── tests/                 # golden prompt fixtures
+├── inference/                 # weight conversion and minimal inference
+│   └── examples/              # equivalent TXT and JSON prompts
+└── config.json                # Hugging Face model metadata
+```
+
+`encoding/` and `inference/` deliberately remain separate: prompt formatting
+does not depend on PyTorch, while inference imports the sibling encoding module
+with an explicit Python path. No symlinks are required.
+
+Weights, `model.safetensors.index.json`, `generation_config.json`, and the
+tokenizer files live in the Hugging Face weight repository, not here.
+
+## Prompt encoding
+
+See [`encoding/README.md`](encoding/README.md). Both OpenAI-style JSON content
+blocks and the compact `<image>path</image>` TXT notation are supported. The two
+examples under `inference/examples/` encode to identical prompts and token IDs.
+
+## Minimal inference
+
+See [`inference/README.md`](inference/README.md) for dependency installation,
+checkpoint conversion, and TXT/JSON inference commands.
+
+## License
+
+This repository is licensed under the [MIT License](LICENSE).

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,56 @@
+# DeepSeek-V4.1 text and vision encoding
+
+`encoding.py` is the standalone prompt-format reference. It supports multi-turn
+conversations, tool calls, thinking modes, and interleaved image content blocks
+without importing the inference implementation.
+
+## 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,
+)
+```
+
+Images are represented in the prompt by `<｜deepseek_image｜>`. `media["images"]`
+contains the corresponding image records in exactly the same order. 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.
+
+## Tests
+
+From this directory:
+
+```bash
+python -m pytest -q test_encoding.py
+```
+
+The tests 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,934 @@
+"""
+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"/"xhigh"/"max"
+   (mapped to 25/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 from OpenAI-format tool list."""
+    return [tool["function"] for tool in tools]
+
+
+def tool_calls_from_openai_format(tool_calls):
+    """Convert OpenAI-format tool calls to internal format."""
+    return [
+        {
+            "name": tool_call["function"]["name"],
+            "arguments": tool_call["function"]["arguments"],
+        }
+        for tool_call in tool_calls
+    ]
+
+
+def tool_calls_to_openai_format(tool_calls):
+    """Convert internal tool calls to OpenAI format."""
+    return [
+        {
+            "type": "function",
+            "function": {
+                "name": tool_call["name"],
+                "arguments": tool_call["arguments"],
+            }
+        }
+        for tool_call in tool_calls
+    ]
+
+
+def decode_dsml_to_arguments(tool_name: str, tool_args: Dict[str, Tuple[str, str]]) -> Dict[str, str]:
+    """
+    Decode DSML parameters back to a tool call dict.
+
+    Args:
+        tool_name: Name of the tool.
+        tool_args: Dict mapping param_name -> (value, is_string_flag).
+
+    Returns:
+        Dict with "name" and "arguments" (JSON string) keys.
+    """
+    def _decode_value(key: str, value: str, string: str):
+        if string == "true":
+            value = to_json(value)
+        return f"{to_json(key)}: {value}"
+
+    tool_args_json = "{" + ", ".join([_decode_value(k, v, string=is_str) for k, (v, is_str) in tool_args.items()]) + "}"
+    return dict(name=tool_name, arguments=tool_args_json)
+
+
+# ============================================================
+# 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": 25,
+    "high": 50,
+    "xhigh": 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=tc.get("name"),
+                    arguments=encode_arguments_to_dsml(tc)
+                )
+                for tc in tool_calls
+            ]
+            tc_content += '\n\n' + tool_calls_template.format(
+                dsml_token=dsml_token,
+                tool_calls="\n".join(tc_list),
+                tc_block_name=tool_calls_block_name,
+            )
+
+        summary_content = content or ""
+        rc = reasoning_content or ""
+
+        # Check if previous message has a task - if so, this is a task output (no thinking)
+        prev_has_task = index - 1 >= 0 and messages[index - 1].get("task") is not None
+
+        if thinking_mode == "thinking" and not prev_has_task:
+            if not drop_thinking or index > last_user_idx:
+                thinking_part = thinking_template.format(reasoning_content=rc) + thinking_end_token
+            else:
+                thinking_part = ""
+
+        if wo_eos:
+            prompt += assistant_msg_wo_eos_template.format(
+                reasoning=thinking_part,
+                content=summary_content,
+                tool_calls=tc_content,
+            )
+        else:
+            prompt += assistant_msg_template.format(
+                reasoning=thinking_part,
+                content=summary_content,
+                tool_calls=tc_content,
+            )
+    else:
+        raise NotImplementedError(f"Unknown role: {role}")
+
+    # Append transition tokens based on what follows
+    if index + 1 < len(messages) and messages[index + 1].get("role") not in ["assistant", "latest_reminder"]:
+        return prompt
+
+    task = messages[index].get("task")
+    if task is not None:
+        # Task special token for internal classification tasks
+        assert task in VALID_TASKS, f"Invalid task: '{task}'. Valid tasks are: {list(VALID_TASKS)}"
+        task_sp_token = DS_TASK_SP_TOKENS[task]
+
+        if task != "action":
+            # Non-action tasks: append task sp token directly after the message
+            prompt += task_sp_token
+        else:
+            # Action task: append Assistant + thinking token + action sp token
+            prompt += ASSISTANT_SP_TOKEN
+            prompt += thinking_end_token if thinking_mode != "thinking" else thinking_start_token
+            prompt += task_sp_token
+
+    elif 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,430 @@
+"""
+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 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, 50),
+        ("low", 25),
+        ("high", 50),
+        ("xhigh", 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")
+
+
+# ============================================================
+# 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: 50 (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>

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-Exp-HF
+export SAVE_PATH=/path/to/DeepSeek-V4.1-Exp-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-Exp-TP8
+export MP=8
+
+INPUT_FILE=examples/example_vl.txt ./run.sh
+INPUT_FILE=examples/example_vl_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)))

tokenizer_config.json

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