Temporary reasoning proxy just to test

.mlx_typings/mlx_lm/generate.pyi

@@ -383,11 +383,12 @@ class GenerationBatch:
     state_machines: List[SequenceStateMachine]
     max_tokens: List[int]
     _current_tokens: Optional[mx.array]
-    _current_logprobs: List[mx.array]
-    _next_tokens: mx.array
-    _next_logprobs: List[mx.array]
-    _token_context: List[mx.array]
+    _current_logprobs: mx.array | List[mx.array]
+    _next_tokens: Optional[mx.array]
+    _next_logprobs: mx.array | List[mx.array]
+    _token_context: List[Any]
     _num_tokens: List[int]
+    _matcher_states: List[Any]
 
     def __init__(
         self,

.mlx_typings/mlx_lm/models/deepseek_v4.pyi

@@ -0,0 +1,267 @@
+"""Type stubs for mlx_lm.models.deepseek_v4"""
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .cache import ArraysCache, RotatingKVCache
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    moe_intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    n_shared_experts: Optional[int]
+    n_routed_experts: int
+    num_experts_per_tok: int
+    head_dim: int
+    qk_rope_head_dim: int
+    q_lora_rank: int
+    o_lora_rank: int
+    o_groups: int
+    sliding_window: int
+    hc_mult: int
+    hc_sinkhorn_iters: int
+    hc_eps: float
+    compress_ratios: Optional[List[int]]
+    compress_rope_theta: float
+    rope_theta: float
+    rope_scaling: Optional[Dict[str, Any]]
+    rms_norm_eps: float
+    swiglu_limit: float
+    attention_bias: bool
+    max_position_embeddings: int
+
+class DeepseekV4RoPE(nn.Module):
+    dims: int
+    freqs: mx.array
+
+    def __init__(
+        self,
+        dims: int,
+        base: float,
+        scaling_config: Optional[Dict[str, Any]] = None,
+    ) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        offset: int = 0,
+        inverse: bool = False,
+    ) -> mx.array: ...
+
+class HyperConnection(nn.Module):
+    dim: int
+    hc_mult: int
+    norm_eps: float
+
+    def __init__(
+        self,
+        dim: int,
+        hc_mult: int,
+        norm_eps: float,
+        sinkhorn_iters: int,
+        hc_eps: float,
+    ) -> None: ...
+
+class HyperHead(nn.Module):
+    dim: int
+    hc_mult: int
+
+    def __init__(
+        self,
+        dim: int,
+        hc_mult: int,
+        norm_eps: float,
+        hc_eps: float,
+    ) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class Compressor(nn.Module):
+    dim: int
+    head_dim: int
+    rope_head_dim: int
+    compress_ratio: int
+    overlap: bool
+    wkv_gate: nn.Linear
+    ape: mx.array
+    norm: nn.RMSNorm
+    rope: DeepseekV4RoPE
+
+    def __init__(
+        self,
+        dim: int,
+        compress_ratio: int,
+        head_dim: int,
+        rope_head_dim: int,
+        rms_norm_eps: float,
+        rope: DeepseekV4RoPE,
+    ) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        state: ArraysCache,
+        offset: Any,
+        slot_compressed: int,
+        slot_kv_state: int,
+        slot_score_state: int,
+    ) -> Optional[mx.array]: ...
+
+class Indexer(nn.Module):
+    def __init__(
+        self,
+        args: ModelArgs,
+        compress_ratio: int,
+        rope: DeepseekV4RoPE,
+    ) -> None: ...
+
+class DeepseekV4Cache:
+    local: Any
+    offset: int
+    keys: Optional[mx.array]
+    values: Optional[mx.array]
+    state: Any
+    meta_state: Any
+    nbytes: int
+
+    def __init__(self, sliding_window: int) -> None: ...
+    def update_and_fetch(
+        self, keys: mx.array, values: mx.array
+    ) -> tuple[mx.array, mx.array]: ...
+    def is_trimmable(self) -> bool: ...
+    def trim(self, n: int) -> int: ...
+    def empty(self) -> bool: ...
+    def size(self) -> int: ...
+    def prepare(
+        self,
+        *,
+        left_padding: Optional[List[int]] = None,
+        lengths: Optional[List[int]] = None,
+        right_padding: Optional[List[int]] = None,
+    ) -> None: ...
+    def finalize(self) -> None: ...
+    def filter(self, batch_indices: mx.array) -> None: ...
+    def extend(self, other: "DeepseekV4Cache") -> None: ...
+    def extract(self, idx: int) -> "DeepseekV4Cache": ...
+    @classmethod
+    def merge(cls, caches: List["DeepseekV4Cache"]) -> "DeepseekV4Cache": ...
+
+class V4Attention(nn.Module):
+    args: ModelArgs
+    layer_id: int
+    dim: int
+    n_heads: int
+    head_dim: int
+    rope_head_dim: int
+    nope_head_dim: int
+    n_groups: int
+    q_lora_rank: int
+    o_lora_rank: int
+    window: int
+    eps: float
+    scale: float
+    compress_ratio: int
+    wqkv_a: nn.Linear
+    q_norm: nn.RMSNorm
+    wq_b: nn.Linear
+    kv_norm: nn.RMSNorm
+    attn_sink: mx.array
+    wo_a: nn.Linear
+    wo_b: nn.Linear
+    rope: DeepseekV4RoPE
+    compressor: Compressor
+    indexer: Indexer
+
+    def __init__(self, args: ModelArgs, layer_id: int) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class DeepseekV4MLP(nn.Module):
+    gate_proj: nn.Linear
+    up_proj: nn.Linear
+    down_proj: nn.Linear
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        swiglu_limit: float = 0.0,
+    ) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class MoEGate(nn.Module):
+    weight: mx.array
+
+    def __init__(self, args: ModelArgs, layer_id: int) -> None: ...
+    def __call__(
+        self, x: mx.array, input_ids: mx.array
+    ) -> tuple[mx.array, mx.array]: ...
+
+class DeepseekV4MoE(nn.Module):
+    num_experts_per_tok: int
+    switch_mlp: SwitchGLU
+    gate: MoEGate
+    shared_experts: DeepseekV4MLP
+
+    def __init__(self, args: ModelArgs, layer_id: int) -> None: ...
+    def __call__(self, x: mx.array, input_ids: mx.array) -> mx.array: ...
+
+class DeepseekV4Block(nn.Module):
+    attn_norm: nn.RMSNorm
+    attn: V4Attention
+    hc_attn: HyperConnection
+    ffn_norm: nn.RMSNorm
+    ffn: DeepseekV4MoE
+    hc_ffn: HyperConnection
+
+    def __init__(self, args: ModelArgs, layer_id: int) -> None: ...
+    def __call__(
+        self,
+        h: mx.array,
+        cache: Optional[Any],
+        input_ids: mx.array,
+    ) -> mx.array: ...
+
+class DeepseekV4Model(nn.Module):
+    args: ModelArgs
+    vocab_size: int
+    embed_tokens: nn.Embedding
+    layers: list[DeepseekV4Block]
+    norm: nn.RMSNorm
+    hc_head: HyperHead
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[List[Any]] = None,
+    ) -> mx.array: ...
+
+class Model(nn.Module):
+    args: ModelArgs
+    model_type: str
+    model: DeepseekV4Model
+    lm_head: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[List[Any]] = None,
+    ) -> mx.array: ...
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
+    def make_cache(self) -> list[RotatingKVCache | DeepseekV4Cache]: ...
+    @property
+    def layers(self) -> list[DeepseekV4Block]: ...

bench/exo_bench.py

@@ -122,8 +122,48 @@ def load_tokenizer_for_bench(model_id: str) -> Any:
 
         return hf_tokenizer
 
-    # Default: use AutoTokenizer
-    return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+    # TODO: Change back to using only transformers
+    try:
+        return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+    except (AttributeError, ValueError):
+        from huggingface_hub import snapshot_download
+        from transformers import PretrainedConfig
+
+        model_path = Path(
+            snapshot_download(
+                model_id,
+                allow_patterns=[
+                    "*.json",
+                    "*.py",
+                    "tokenizer.model",
+                    "*.tiktoken",
+                    "tiktoken.model",
+                    "*.txt",
+                    "*.jsonl",
+                    "*.jinja",
+                ],
+            )
+        )
+        stub_kwargs: dict[str, Any] = {}
+        config_file = model_path / "config.json"
+        if config_file.exists():
+            with open(config_file) as f:
+                raw = json.load(f)
+            for key in (
+                "model_type",
+                "max_position_embeddings",
+                "vocab_size",
+                "bos_token_id",
+                "eos_token_id",
+                "pad_token_id",
+            ):
+                if key in raw:
+                    stub_kwargs[key] = raw[key]
+        return AutoTokenizer.from_pretrained(
+            str(model_path),
+            config=PretrainedConfig(**stub_kwargs),
+            trust_remote_code=True,
+        )
 
 
 def format_peak_memory(b: float) -> str:
@@ -278,9 +318,19 @@ class PromptSizer:
     def _make_counter(tokenizer: Any) -> Callable[[str], int]:
         def count_fn(user_content: str) -> int:
             messages = [{"role": "user", "content": user_content}]
-            ids = tokenizer.apply_chat_template(
-                messages, tokenize=True, add_generation_prompt=True
-            )
+            try:
+                ids = tokenizer.apply_chat_template(
+                    messages, tokenize=True, add_generation_prompt=True
+                )
+            except ValueError:
+                # Models without a Jinja chat template (e.g. DeepSeek V4 which
+                # ships its own Python encoder). Use the exo-side V4 encoder.
+                from exo.worker.engines.mlx.deepseek_v4_encoding import (
+                    encode_messages as encode_v4,
+                )
+
+                prompt = encode_v4(messages, thinking_mode="thinking")
+                ids = tokenizer.encode(prompt, add_special_tokens=False)
             # Fix for transformers 5.x
             if hasattr(ids, "input_ids"):
                 ids = ids.input_ids

pyproject.toml

@@ -29,11 +29,12 @@ dependencies = [
   "msgspec>=0.19.0",
   "zstandard>=0.23.0",
   "mlx-vlm>=0.3.11",
-  "transformers>=5.0.0,<5.4.0",
+  "transformers>=5.6.2",
 ]
 
 [project.scripts]
 exo = "exo.main:main"
+exo-reasoning-proxy = "exo.reasoning_proxy.main:main"
 
 # dependencies only required for development
 [dependency-groups]
@@ -77,7 +78,7 @@ members = ["rust/exo_pyo3_bindings", "bench"]
 [tool.uv.sources]
 exo-pyo3-bindings = { workspace = true }
 mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
-mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/fix-arrayscache-leak" }
+mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/deepseek-v4" }
 torch = [
   { index = "pytorch-cu130", marker = "sys_platform == 'linux' and extra == 'cuda13' and extra != 'cpu' and extra != 'cuda12'" },
   { index = "pytorch-cu120", marker = "sys_platform == 'linux' and extra == 'cuda12' and extra != 'cpu' and extra != 'cuda13'" },
@@ -154,7 +155,7 @@ required-version = ">=0.8.6"
 prerelease = "allow"
 environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
 conflicts = [[{ extra = "cuda12" }, { extra = "cuda13" }, { extra = "cpu" }]]
-constraint-dependencies = ["transformers>=5.0.0,<5.4.0"]
+constraint-dependencies = ["transformers>=5.6.2"]
 override-dependencies = [
   "mlx==0.31.1; sys_platform=='linux'",
   "mlx; sys_platform=='darwin'",

resources/inference_model_cards/mlx-community--DeepSeek-V4-Flash.toml

@@ -0,0 +1,20 @@
+model_id = "mlx-community/DeepSeek-V4-Flash"
+n_layers = 43
+hidden_size = 4096
+num_key_value_heads = 1
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "deepseek"
+quantization = "8bit"
+base_model = "DeepSeek V4 Flash"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+context_length = 1048576
+
+[storage_size]
+in_bytes = 155095760030
+
+# Source: https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash
+[sampling_defaults]
+temperature = 1.0
+top_p = 1.0

resources/inference_model_cards/mlx-community--DeepSeek-V4-Pro.toml

@@ -0,0 +1,20 @@
+model_id = "mlx-community/DeepSeek-V4-Pro"
+n_layers = 61
+hidden_size = 7168
+num_key_value_heads = 1
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "deepseek"
+quantization = "8bit"
+base_model = "DeepSeek V4 Pro"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+context_length = 1048576
+
+[storage_size]
+in_bytes = 849681803879
+
+# Source: https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro
+[sampling_defaults]
+temperature = 1.0
+top_p = 1.0

resources/inference_model_cards/mlx-community--GLM-4.5-Air-8bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "8bit"
 base_model = "GLM 4.5 Air"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 131072
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.5-Air-bf16.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "bf16"
 base_model = "GLM 4.5 Air"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 131072
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-4bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "4bit"
 base_model = "GLM 4.7"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-6bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "6bit"
 base_model = "GLM 4.7"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-8bit-gs32.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "8bit"
 base_model = "GLM 4.7"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-Flash-4bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "4bit"
 base_model = "GLM 4.7 Flash"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-Flash-5bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "5bit"
 base_model = "GLM 4.7 Flash"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-Flash-6bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "6bit"
 base_model = "GLM 4.7 Flash"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-Flash-8bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "8bit"
 base_model = "GLM 4.7 Flash"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-5-8bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "8bit"
 base_model = "GLM-5"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-5-MXFP4-Q8.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "MXFP4-Q8"
 base_model = "GLM-5"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-5-bf16.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "bf16"
 base_model = "GLM-5"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-5.1-DQ4plus-q8.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "8bit"
 base_model = "GLM-5.1"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-5.1-MXFP4-Q8.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "MXFP4-Q8"
 base_model = "GLM-5.1"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-5.1.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "bf16"
 base_model = "GLM-5.1"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--Kimi-K2-Thinking.toml

@@ -8,7 +8,7 @@ family = "kimi"
 quantization = ""
 base_model = "Kimi K2"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "suffix"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Kimi-K2.5.toml

@@ -8,7 +8,7 @@ family = "kimi"
 quantization = ""
 base_model = "Kimi K2.5"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "suffix"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Kimi-K2.6-mlx-DQ3_K_M-q8.toml

@@ -8,7 +8,7 @@ family = "kimi"
 quantization = "3bit"
 base_model = "Kimi K2.6"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "suffix"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.1-3bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "3bit"
 base_model = "MiniMax M2.1"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.1-8bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "8bit"
 base_model = "MiniMax M2.1"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.5-4bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "4bit"
 base_model = "MiniMax M2.5"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.5-6bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "6bit"
 base_model = "MiniMax M2.5"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.5-8bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "8bit"
 base_model = "MiniMax M2.5"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.7-4bit-mxfp4.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "4bit-mxfp4"
 base_model = "MiniMax M2.7"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.7-4bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "4bit"
 base_model = "MiniMax M2.7"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.7-5bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "5bit"
 base_model = "MiniMax M2.7"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.7-6bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "6bit"
 base_model = "MiniMax M2.7"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.7-8bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "8bit"
 base_model = "MiniMax M2.7"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.7.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "bf16"
 base_model = "MiniMax M2.7"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3-Next-80B-A3B-Thinking-4bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3 Next 80B"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3-Next-80B-A3B-Thinking-8bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3 Next 80B"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-4bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3.5 122B A10B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-6bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "6bit"
 base_model = "Qwen3.5 122B A10B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-8bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3.5 122B A10B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-bf16.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "bf16"
 base_model = "Qwen3.5 122B A10B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.5-27B-4bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3.5 27B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.5-27B-8bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3.5 27B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.5-2B-MLX-8bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3.5 2B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.5-35B-A3B-4bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3.5 35B A3B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.5-35B-A3B-8bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3.5 35B A3B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.5-397B-A17B-4bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3.5 397B A17B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.5-397B-A17B-6bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "6bit"
 base_model = "Qwen3.5 397B A17B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.5-397B-A17B-8bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3.5 397B A17B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.5-9B-4bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3.5 9B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.5-9B-8bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3.5 9B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.6-27B-4bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3.6 27B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.6-27B-8bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3.6 27B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.6-27B-bf16.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "bf16"
 base_model = "Qwen3.6 27B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.6-35B-A3B-4bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3.6 35B A3B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.6-35B-A3B-5bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "5bit"
 base_model = "Qwen3.6 35B A3B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.6-35B-A3B-8bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3.6 35B A3B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.6-35B-A3B-bf16.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "bf16"
 base_model = "Qwen3.6 35B A3B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--gpt-oss-120b-MXFP4-Q8.toml

@@ -8,7 +8,7 @@ family = "gpt-oss"
 quantization = "MXFP4-Q8"
 base_model = "GPT-OSS 120B"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "channel"
 context_length = 131072
 
 [storage_size]

resources/inference_model_cards/mlx-community--gpt-oss-20b-MXFP4-Q8.toml

@@ -8,7 +8,7 @@ family = "gpt-oss"
 quantization = "MXFP4-Q8"
 base_model = "GPT-OSS 20B"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "channel"
 context_length = 131072
 
 [storage_size]

src/exo/api/main.py

@@ -1663,6 +1663,7 @@ class API:
                     quantization=card.quantization,
                     base_model=card.base_model,
                     capabilities=card.capabilities,
+                    reasoning_dialect=card.reasoning_dialect,
                     context_length=card.context_length,
                 )
                 for card in cards

src/exo/api/types/api.py

@@ -8,7 +8,7 @@ from pydantic import BaseModel, Field, field_validator
 from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.memory import Memory
-from exo.shared.types.text_generation import ReasoningEffort
+from exo.shared.types.text_generation import ReasoningDialect, ReasoningEffort
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding, ShardMetadata
 from exo.utils.pydantic_ext import FrozenModel
@@ -48,6 +48,7 @@ class ModelListModel(BaseModel):
     quantization: str = Field(default="")
     base_model: str = Field(default="")
     capabilities: list[str] = Field(default_factory=list)
+    reasoning_dialect: ReasoningDialect = "none"
 
 
 class ModelList(BaseModel):

src/exo/master/placement.py

@@ -133,12 +133,16 @@ def place_instance(
                 f"Requested Tensor sharding but this model does not support tensor parallelism: {command.model_card.model_id}"
             )
         # TODO: the condition here for tensor parallel is not correct, but it works good enough for now.
+        # DeepSeek V4 is MQA (num_key_value_heads=1) but its sharding strategy
+        # head-parallelises wq_b/wo_a and shards MoE experts instead of splitting
+        # KV heads, so the kv-head divisibility check doesn't apply.
+        is_deepseek_v4 = command.model_card.base_model.startswith("DeepSeek V4")
         kv_heads = command.model_card.num_key_value_heads
         cycles_with_sufficient_memory = [
             cycle
             for cycle in cycles_with_sufficient_memory
             if command.model_card.hidden_size % len(cycle) == 0
-            and (kv_heads is None or kv_heads % len(cycle) == 0)
+            and (is_deepseek_v4 or kv_heads is None or kv_heads % len(cycle) == 0)
         ]
         if not cycles_with_sufficient_memory:
             raise ValueError(

src/exo/reasoning_proxy/_helpers.py

@@ -0,0 +1,33 @@
+from typing import cast
+
+
+def as_str(value: object) -> str | None:
+    return value if isinstance(value, str) else None
+
+
+def as_list(value: object) -> list[object] | None:
+    if isinstance(value, list):
+        return cast(list[object], value)
+    return None
+
+
+def as_dict(value: object) -> dict[str, object] | None:
+    if isinstance(value, dict):
+        return cast(dict[str, object], value)
+    return None
+
+
+def as_int(value: object, default: int = 0) -> int:
+    return value if isinstance(value, int) and not isinstance(value, bool) else default
+
+
+def dict_get_str(d: dict[str, object], key: str) -> str | None:
+    return as_str(d.get(key))
+
+
+def dict_get_list(d: dict[str, object], key: str) -> list[object] | None:
+    return as_list(d.get(key))
+
+
+def dict_get_dict(d: dict[str, object], key: str) -> dict[str, object] | None:
+    return as_dict(d.get(key))

src/exo/reasoning_proxy/accumulator.py

@@ -0,0 +1,261 @@
+"""Accumulators for capturing the emitted assistant shape from a streaming response.
+
+Both accumulators are fed raw SSE chunks (bytes) as they pass through. At stream
+end, they expose a canonical assistant-message shape suitable for hashing, plus
+the reasoning text that should be cached against that hash.
+"""
+
+import json
+import logging
+from typing import cast
+
+from exo.reasoning_proxy._helpers import (
+    as_dict,
+    as_str,
+    dict_get_dict,
+    dict_get_list,
+    dict_get_str,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class OpenAIAccumulator:
+    """Captures content, tool_calls, and reasoning_content from OpenAI SSE chunks.
+
+    OpenAI can emit multiple choices per chunk; we only track choice index 0
+    (the common case for chat completions; n>1 is uncommon and re-hash misses
+    there degrade gracefully to no-op cache insert).
+    """
+
+    def __init__(self) -> None:
+        self._content_parts: list[str] = []
+        self._reasoning_parts: list[str] = []
+        self._tool_calls_by_index: dict[int, dict[str, object]] = {}
+        self._buffer = ""
+
+    def feed_bytes(self, chunk: bytes) -> None:
+        self._buffer += chunk.decode("utf-8", errors="replace")
+        while "\n" in self._buffer:
+            line, self._buffer = self._buffer.split("\n", 1)
+            line = line.strip()
+            if not line.startswith("data:"):
+                continue
+            payload = line[len("data:") :].strip()
+            if payload == "[DONE]" or not payload:
+                continue
+            try:
+                parsed = cast(object, json.loads(payload))
+            except json.JSONDecodeError:
+                continue
+            event = as_dict(parsed)
+            if event is None:
+                continue
+            self._consume_event(event)
+
+    def _consume_event(self, event: dict[str, object]) -> None:
+        choices = dict_get_list(event, "choices")
+        if not choices:
+            return
+        for choice_raw in choices:
+            choice = as_dict(choice_raw)
+            if choice is None:
+                continue
+            index_val = choice.get("index", 0)
+            if (
+                not (isinstance(index_val, int) and not isinstance(index_val, bool))
+                or index_val != 0
+            ):
+                continue
+            delta = dict_get_dict(choice, "delta")
+            if delta is None:
+                continue
+            content = dict_get_str(delta, "content")
+            if content is not None:
+                self._content_parts.append(content)
+            reasoning = dict_get_str(delta, "reasoning_content")
+            if reasoning is not None:
+                self._reasoning_parts.append(reasoning)
+            tool_calls = dict_get_list(delta, "tool_calls")
+            if tool_calls is not None:
+                self._merge_tool_calls(tool_calls)
+
+    def _merge_tool_calls(self, deltas: list[object]) -> None:
+        for raw in deltas:
+            d = as_dict(raw)
+            if d is None:
+                continue
+            index_val = d.get("index", 0)
+            if not (isinstance(index_val, int) and not isinstance(index_val, bool)):
+                continue
+            entry = self._tool_calls_by_index.setdefault(
+                index_val,
+                {
+                    "id": "",
+                    "type": "function",
+                    "function": {"name": "", "arguments": ""},
+                },
+            )
+            tc_id = dict_get_str(d, "id")
+            if tc_id is not None:
+                entry["id"] = tc_id
+            tc_type = dict_get_str(d, "type")
+            if tc_type is not None:
+                entry["type"] = tc_type
+            fn = dict_get_dict(d, "function")
+            if fn is not None:
+                entry_fn = entry.get("function")
+                if not isinstance(entry_fn, dict):
+                    entry_fn = {"name": "", "arguments": ""}
+                    entry["function"] = entry_fn
+                entry_fn_typed = cast(dict[str, object], entry_fn)
+                name = dict_get_str(fn, "name")
+                if name is not None:
+                    prev_name = as_str(entry_fn_typed.get("name")) or ""
+                    entry_fn_typed["name"] = prev_name + name
+                args = dict_get_str(fn, "arguments")
+                if args is not None:
+                    prev_args = as_str(entry_fn_typed.get("arguments")) or ""
+                    entry_fn_typed["arguments"] = prev_args + args
+
+    @property
+    def content(self) -> str | None:
+        joined = "".join(self._content_parts)
+        return joined if joined else None
+
+    @property
+    def tool_calls(self) -> list[dict[str, object]] | None:
+        if not self._tool_calls_by_index:
+            return None
+        ordered = [
+            self._tool_calls_by_index[i] for i in sorted(self._tool_calls_by_index)
+        ]
+        return ordered
+
+    @property
+    def reasoning(self) -> str:
+        return "".join(self._reasoning_parts)
+
+
+class ClaudeAccumulator:
+    """Captures Claude streaming content blocks.
+
+    Tracks per-index content blocks. At end, exposes the final `content_blocks`
+    list (excluding thinking blocks — those go into `reasoning` as joined text)
+    in a shape suitable for hashing.
+    """
+
+    def __init__(self) -> None:
+        self._blocks_by_index: dict[int, dict[str, object]] = {}
+        self._buffer = ""
+        self._current_event: str | None = None
+
+    def feed_bytes(self, chunk: bytes) -> None:
+        self._buffer += chunk.decode("utf-8", errors="replace")
+        while "\n" in self._buffer:
+            line, self._buffer = self._buffer.split("\n", 1)
+            line = line.rstrip("\r")
+            if not line:
+                self._current_event = None
+                continue
+            if line.startswith("event:"):
+                self._current_event = line[len("event:") :].strip()
+                continue
+            if line.startswith("data:"):
+                payload = line[len("data:") :].strip()
+                try:
+                    parsed = cast(object, json.loads(payload))
+                except json.JSONDecodeError:
+                    continue
+                event = as_dict(parsed)
+                if event is None:
+                    continue
+                self._consume_event(event)
+
+    def _consume_event(self, event: dict[str, object]) -> None:
+        event_type = dict_get_str(event, "type") or self._current_event
+        if event_type == "content_block_start":
+            index_val = event.get("index", 0)
+            if not (isinstance(index_val, int) and not isinstance(index_val, bool)):
+                return
+            block = dict_get_dict(event, "content_block")
+            if block is None:
+                return
+            btype = dict_get_str(block, "type")
+            if btype == "text":
+                self._blocks_by_index[index_val] = {"type": "text", "text": ""}
+            elif btype == "thinking":
+                self._blocks_by_index[index_val] = {
+                    "type": "thinking",
+                    "thinking": "",
+                }
+            elif btype == "tool_use":
+                self._blocks_by_index[index_val] = {
+                    "type": "tool_use",
+                    "id": dict_get_str(block, "id") or "",
+                    "name": dict_get_str(block, "name") or "",
+                    "input_json": "",
+                }
+        elif event_type == "content_block_delta":
+            index_val = event.get("index", 0)
+            if not (isinstance(index_val, int) and not isinstance(index_val, bool)):
+                return
+            delta = dict_get_dict(event, "delta")
+            if delta is None:
+                return
+            block = self._blocks_by_index.get(index_val)
+            if block is None:
+                return
+            dtype = dict_get_str(delta, "type")
+            if dtype == "text_delta":
+                text = dict_get_str(delta, "text")
+                if text is not None:
+                    prev = as_str(block.get("text")) or ""
+                    block["text"] = prev + text
+            elif dtype == "thinking_delta":
+                thinking = dict_get_str(delta, "thinking")
+                if thinking is not None:
+                    prev = as_str(block.get("thinking")) or ""
+                    block["thinking"] = prev + thinking
+            elif dtype == "input_json_delta":
+                partial = dict_get_str(delta, "partial_json")
+                if partial is not None:
+                    prev = as_str(block.get("input_json")) or ""
+                    block["input_json"] = prev + partial
+
+    @property
+    def content_blocks(self) -> list[dict[str, object]]:
+        """Public blocks (excludes thinking), with tool_use input parsed from JSON."""
+        public: list[dict[str, object]] = []
+        for index in sorted(self._blocks_by_index):
+            block = self._blocks_by_index[index]
+            if block.get("type") == "thinking":
+                continue
+            if block.get("type") == "tool_use":
+                input_json = as_str(block.get("input_json")) or "{}"
+                parsed_input_raw: object
+                try:
+                    parsed_input_raw = cast(object, json.loads(input_json))
+                except json.JSONDecodeError:
+                    parsed_input_raw = {}
+                parsed_input = as_dict(parsed_input_raw) or {}
+                public.append(
+                    {
+                        "type": "tool_use",
+                        "id": as_str(block.get("id")) or "",
+                        "name": as_str(block.get("name")) or "",
+                        "input": parsed_input,
+                    }
+                )
+            else:
+                public.append({k: v for k, v in block.items() if k != "input_json"})
+        return public
+
+    @property
+    def reasoning(self) -> str:
+        parts: list[str] = []
+        for index in sorted(self._blocks_by_index):
+            block = self._blocks_by_index[index]
+            if block.get("type") == "thinking":
+                parts.append(as_str(block.get("thinking")) or "")
+        return "".join(parts)

src/exo/reasoning_proxy/cache.py

@@ -0,0 +1,21 @@
+import threading
+
+
+class ReasoningCache:
+    def __init__(self) -> None:
+        self._store: dict[str, str] = {}
+        self._lock = threading.Lock()
+
+    def get(self, content_hash: str) -> str | None:
+        with self._lock:
+            return self._store.get(content_hash)
+
+    def put(self, content_hash: str, reasoning: str) -> None:
+        if not reasoning:
+            return
+        with self._lock:
+            self._store[content_hash] = reasoning
+
+    def size(self) -> int:
+        with self._lock:
+            return len(self._store)

src/exo/reasoning_proxy/dialects.py

@@ -0,0 +1,170 @@
+"""Dialect strategies for deciding which assistant history indices should receive
+cached reasoning on inbound requests.
+
+Each dialect inspects the message list and returns the set of indices where
+reasoning_content (OpenAI) or a thinking block (Claude) should be reattached if
+the cache has it. The dialect does not mutate messages — the caller does.
+
+Dialect selection is driven by the `reasoning_dialect` field on each model card,
+surfaced through /v1/models.
+"""
+
+from typing import Protocol
+
+from exo.reasoning_proxy._helpers import as_dict, as_list, dict_get_list, dict_get_str
+from exo.shared.types.text_generation import ReasoningDialect
+
+
+class Dialect(Protocol):
+    def select_attach_indices(
+        self, messages: list[dict[str, object]], has_tools: bool
+    ) -> set[int]: ...
+
+
+def _is_assistant(msg: dict[str, object]) -> bool:
+    return msg.get("role") == "assistant"
+
+
+def _is_user(msg: dict[str, object]) -> bool:
+    return msg.get("role") == "user"
+
+
+class NoneDialect:
+    def select_attach_indices(
+        self, messages: list[dict[str, object]], has_tools: bool
+    ) -> set[int]:
+        return set()
+
+
+class PostLastUserDialect:
+    """MiniMax / GLM / Qwen-thinking / V4-with-tools.
+
+    Preserve reasoning on every assistant message appearing after the last
+    non-tool-response user message. Tool-response user messages (role=tool, or
+    role=user with tool_call_id set, or Claude's tool_result block) don't count
+    as "real" user turns — they're part of the assistant's tool-calling chain.
+    """
+
+    def select_attach_indices(
+        self, messages: list[dict[str, object]], has_tools: bool
+    ) -> set[int]:
+        last_user_index = -1
+        for i, msg in enumerate(messages):
+            if _is_user(msg) and not _is_tool_response(msg):
+                last_user_index = i
+        return {
+            i
+            for i, msg in enumerate(messages)
+            if i > last_user_index and _is_assistant(msg)
+        }
+
+
+class SuffixDialect:
+    """Kimi K2 Thinking / K2.6.
+
+    Preserve reasoning only on the tail run of tool-call-carrying assistant
+    messages (the current, unresolved tool-call chain). Walk backward: include
+    every assistant with tool_calls until we hit an assistant without tool_calls
+    or a non-assistant message.
+    """
+
+    def select_attach_indices(
+        self, messages: list[dict[str, object]], has_tools: bool
+    ) -> set[int]:
+        indices: set[int] = set()
+        for i in range(len(messages) - 1, -1, -1):
+            msg = messages[i]
+            if not _is_assistant(msg):
+                if _is_tool_response(msg):
+                    continue
+                break
+            if not _has_tool_calls(msg):
+                break
+            indices.add(i)
+        return indices
+
+
+class ChannelDialect:
+    """GPT-OSS Harmony format.
+
+    Preserve analysis-channel content on assistant turns that follow the most
+    recent assistant message tagged with a "final" channel marker. If no prior
+    final exists, the whole conversation is one unresolved chain.
+    """
+
+    def select_attach_indices(
+        self, messages: list[dict[str, object]], has_tools: bool
+    ) -> set[int]:
+        last_final_index = -1
+        for i, msg in enumerate(messages):
+            if _is_assistant(msg) and _has_final_channel(msg):
+                last_final_index = i
+        return {
+            i
+            for i, msg in enumerate(messages)
+            if i > last_final_index and _is_assistant(msg)
+        }
+
+
+class ToolConditionalDialect:
+    """DeepSeek V4 Flash.
+
+    If the request has tools, behave as PostLastUserDialect; otherwise passthrough.
+    """
+
+    def __init__(self) -> None:
+        self._inner = PostLastUserDialect()
+
+    def select_attach_indices(
+        self, messages: list[dict[str, object]], has_tools: bool
+    ) -> set[int]:
+        if not has_tools:
+            return set()
+        return self._inner.select_attach_indices(messages, has_tools)
+
+
+def _is_tool_response(msg: dict[str, object]) -> bool:
+    if msg.get("role") == "tool":
+        return True
+    if msg.get("role") == "user" and msg.get("tool_call_id"):
+        return True
+    content = as_list(msg.get("content"))
+    if content is not None:
+        for raw in content:
+            block = as_dict(raw)
+            if block is not None and block.get("type") == "tool_result":
+                return True
+    return False
+
+
+def _has_tool_calls(msg: dict[str, object]) -> bool:
+    tc = dict_get_list(msg, "tool_calls")
+    if tc:
+        return True
+    content = as_list(msg.get("content"))
+    if content is not None:
+        for raw in content:
+            block = as_dict(raw)
+            if block is not None and block.get("type") == "tool_use":
+                return True
+    return False
+
+
+def _has_final_channel(msg: dict[str, object]) -> bool:
+    if msg.get("channel") == "final":
+        return True
+    content = dict_get_str(msg, "content")
+    return bool(content and content.strip())
+
+
+_DIALECTS: dict[ReasoningDialect, Dialect] = {
+    "none": NoneDialect(),
+    "post_last_user": PostLastUserDialect(),
+    "suffix": SuffixDialect(),
+    "channel": ChannelDialect(),
+    "tool_conditional": ToolConditionalDialect(),
+}
+
+
+def get_dialect(name: ReasoningDialect) -> Dialect:
+    return _DIALECTS[name]

src/exo/reasoning_proxy/hashing.py

@@ -0,0 +1,75 @@
+import hashlib
+import json
+
+from exo.reasoning_proxy._helpers import as_dict, dict_get_str
+
+
+def _canonical_tool_calls(
+    tool_calls: list[dict[str, object]] | None,
+) -> list[dict[str, object]]:
+    if not tool_calls:
+        return []
+    result: list[dict[str, object]] = []
+    for tc in tool_calls:
+        entry: dict[str, object] = {}
+        if "id" in tc:
+            entry["id"] = tc["id"]
+        fn = as_dict(tc.get("function"))
+        if fn is not None:
+            entry["function"] = {
+                "name": dict_get_str(fn, "name") or "",
+                "arguments": dict_get_str(fn, "arguments") or "",
+            }
+        if "type" in tc:
+            entry["type"] = tc["type"]
+        result.append(entry)
+    return result
+
+
+def hash_openai_assistant(
+    content: str | list[object] | None,
+    tool_calls: list[dict[str, object]] | None,
+) -> str:
+    """Deterministic hash of an OpenAI assistant message's observable surface.
+
+    Canonicalizes None content to "" and tool_calls to a minimal id/function shape
+    so trivial shape differences between client render and our re-emit don't miss.
+    """
+    shape: dict[str, object] = {
+        "content": content if content is not None else "",
+        "tool_calls": _canonical_tool_calls(tool_calls),
+    }
+    payload = json.dumps(
+        shape, sort_keys=True, ensure_ascii=False, separators=(",", ":")
+    )
+    return hashlib.sha256(payload.encode("utf-8")).hexdigest()
+
+
+def hash_claude_assistant(content_blocks: list[dict[str, object]]) -> str:
+    """Deterministic hash of a Claude assistant message's observable surface.
+
+    Skips thinking blocks (we're hashing what the *client sends back*, which typically
+    omits thinking) and normalizes tool_use blocks to id/name/input.
+    """
+    normalized: list[dict[str, object]] = []
+    for block in content_blocks:
+        btype = block.get("type")
+        if btype == "text":
+            normalized.append(
+                {"type": "text", "text": dict_get_str(block, "text") or ""}
+            )
+        elif btype == "tool_use":
+            normalized.append(
+                {
+                    "type": "tool_use",
+                    "id": dict_get_str(block, "id") or "",
+                    "name": dict_get_str(block, "name") or "",
+                    "input": block.get("input")
+                    if block.get("input") is not None
+                    else {},
+                }
+            )
+    payload = json.dumps(
+        normalized, sort_keys=True, ensure_ascii=False, separators=(",", ":")
+    )
+    return hashlib.sha256(payload.encode("utf-8")).hexdigest()

src/exo/reasoning_proxy/main.py

@@ -0,0 +1,70 @@
+import argparse
+import logging
+from contextlib import asynccontextmanager
+from typing import cast
+
+import httpx
+import uvicorn
+from fastapi import FastAPI
+
+from exo.reasoning_proxy.cache import ReasoningCache
+from exo.reasoning_proxy.registry import DialectRegistry
+from exo.reasoning_proxy.routes import register_routes
+
+logger = logging.getLogger("exo.reasoning_proxy")
+
+
+def build_app(upstream: str) -> FastAPI:
+    client = httpx.AsyncClient(timeout=httpx.Timeout(None, connect=10.0))
+    cache = ReasoningCache()
+    registry = DialectRegistry(upstream=upstream, client=client)
+
+    @asynccontextmanager
+    async def lifespan(_: FastAPI):
+        await registry.refresh()
+        try:
+            yield
+        finally:
+            await client.aclose()
+
+    app = FastAPI(lifespan=lifespan, title="exo-reasoning-proxy")
+    register_routes(
+        app=app,
+        client=client,
+        upstream=upstream.rstrip("/"),
+        cache=cache,
+        registry=registry,
+    )
+    return app
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(prog="exo-reasoning-proxy")
+    _ = parser.add_argument("--upstream", default="http://localhost:52415")
+    _ = parser.add_argument("--host", default="127.0.0.1")
+    _ = parser.add_argument("--port", type=int, default=52416)
+    _ = parser.add_argument("-v", "--verbose", action="count", default=0)
+    args = parser.parse_args()
+
+    verbose = cast(int, args.verbose)
+    upstream = cast(str, args.upstream)
+    host = cast(str, args.host)
+    port = cast(int, args.port)
+
+    level = logging.WARNING
+    if verbose == 1:
+        level = logging.INFO
+    elif verbose >= 2:
+        level = logging.DEBUG
+    logging.basicConfig(
+        level=level, format="%(asctime)s %(levelname)s %(name)s: %(message)s"
+    )
+
+    logger.info("Starting exo-reasoning-proxy on %s:%d → %s", host, port, upstream)
+
+    app = build_app(upstream=upstream)
+    uvicorn.run(app, host=host, port=port, log_level=level)
+
+
+if __name__ == "__main__":
+    main()

src/exo/reasoning_proxy/registry.py

@@ -0,0 +1,75 @@
+import asyncio
+import logging
+from typing import cast, get_args
+
+import httpx
+
+from exo.reasoning_proxy._helpers import as_dict, as_list, dict_get_str
+from exo.shared.types.text_generation import ReasoningDialect
+
+logger = logging.getLogger(__name__)
+
+
+class DialectRegistry:
+    def __init__(self, upstream: str, client: httpx.AsyncClient) -> None:
+        self._upstream = upstream.rstrip("/")
+        self._client = client
+        self._by_model: dict[str, ReasoningDialect] = {}
+        self._unknown_logged: set[str] = set()
+        self._lock = asyncio.Lock()
+        self._initialized = False
+
+    async def refresh(self) -> None:
+        await self._fetch()
+
+    async def _fetch(self) -> None:
+        try:
+            resp = await self._client.get(f"{self._upstream}/v1/models", timeout=10.0)
+            resp.raise_for_status()
+            body = as_dict(cast(object, resp.json()))
+            if body is None:
+                return
+            data = as_list(body.get("data")) or []
+            updated: dict[str, ReasoningDialect] = {}
+            for entry_raw in data:
+                entry = as_dict(entry_raw)
+                if entry is None:
+                    continue
+                model_id = dict_get_str(entry, "id")
+                dialect_raw = entry.get("reasoning_dialect", "none")
+                if model_id is not None:
+                    updated[model_id] = _coerce_dialect(dialect_raw)
+            self._by_model = updated
+            self._initialized = True
+            logger.info(
+                "Loaded %d model dialects from %s", len(updated), self._upstream
+            )
+        except Exception as exc:
+            logger.warning(
+                "Failed to fetch /v1/models from %s: %s", self._upstream, exc
+            )
+
+    async def resolve(self, model_id: str) -> ReasoningDialect:
+        async with self._lock:
+            if not self._initialized:
+                await self._fetch()
+            if model_id in self._by_model:
+                return self._by_model[model_id]
+            await self._fetch()
+            if model_id in self._by_model:
+                return self._by_model[model_id]
+            if model_id not in self._unknown_logged:
+                logger.info(
+                    "No dialect declared for model %s; passing through", model_id
+                )
+                self._unknown_logged.add(model_id)
+            return "none"
+
+
+_VALID_DIALECTS: frozenset[str] = frozenset(get_args(ReasoningDialect))
+
+
+def _coerce_dialect(value: object) -> ReasoningDialect:
+    if isinstance(value, str) and value in _VALID_DIALECTS:
+        return cast(ReasoningDialect, value)
+    return "none"

src/exo/reasoning_proxy/routes.py

@@ -0,0 +1,384 @@
+"""FastAPI handlers for the reasoning proxy.
+
+Two handlers, one shape: read body → resolve dialect → reattach cached
+reasoning to designated history indices → forward → tee the response stream
+→ capture emitted reasoning → cache under the emitted assistant's hash.
+"""
+
+import json
+import logging
+from collections.abc import AsyncIterator
+from typing import cast
+
+import httpx
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse, Response, StreamingResponse
+from starlette.datastructures import Headers as StarletteHeaders
+
+from exo.reasoning_proxy._helpers import (
+    as_dict,
+    as_list,
+    dict_get_dict,
+    dict_get_list,
+    dict_get_str,
+)
+from exo.reasoning_proxy.accumulator import ClaudeAccumulator, OpenAIAccumulator
+from exo.reasoning_proxy.cache import ReasoningCache
+from exo.reasoning_proxy.dialects import get_dialect
+from exo.reasoning_proxy.hashing import hash_claude_assistant, hash_openai_assistant
+from exo.reasoning_proxy.registry import DialectRegistry
+
+logger = logging.getLogger(__name__)
+
+
+def _parse_body(raw_body: bytes) -> dict[str, object] | None:
+    try:
+        parsed = cast(object, json.loads(raw_body))
+    except json.JSONDecodeError:
+        return None
+    return as_dict(parsed)
+
+
+def _content_for_hash(value: object) -> str | list[object] | None:
+    if isinstance(value, str):
+        return value
+    as_listed = as_list(value)
+    if as_listed is not None:
+        return as_listed
+    return None
+
+
+def _attach_openai_reasoning(
+    messages: list[dict[str, object]],
+    indices: set[int],
+    cache: ReasoningCache,
+) -> None:
+    for i in indices:
+        msg = messages[i]
+        existing = dict_get_str(msg, "reasoning_content")
+        if existing:
+            continue
+        content = _content_for_hash(msg.get("content"))
+        tool_calls_raw = dict_get_list(msg, "tool_calls") or []
+        tool_calls: list[dict[str, object]] = [
+            d for d in (as_dict(t) for t in tool_calls_raw) if d is not None
+        ]
+        h = hash_openai_assistant(content, tool_calls or None)
+        cached = cache.get(h)
+        if cached is not None:
+            msg["reasoning_content"] = cached
+
+
+def _attach_claude_reasoning(
+    messages: list[dict[str, object]],
+    indices: set[int],
+    cache: ReasoningCache,
+) -> None:
+    for i in indices:
+        msg = messages[i]
+        content = as_list(msg.get("content"))
+        if content is None:
+            continue
+        has_thinking = False
+        normalized: list[dict[str, object]] = []
+        for raw in content:
+            block = as_dict(raw)
+            if block is None:
+                continue
+            if block.get("type") == "thinking":
+                has_thinking = True
+            normalized.append(block)
+        if has_thinking:
+            continue
+        h = hash_claude_assistant(normalized)
+        cached = cache.get(h)
+        if cached is None:
+            continue
+        new_content: list[dict[str, object]] = [
+            {"type": "thinking", "thinking": cached}
+        ]
+        new_content.extend(normalized)
+        msg["content"] = new_content
+
+
+async def _stream_and_capture_openai(
+    upstream_resp: httpx.Response,
+    cache: ReasoningCache,
+) -> AsyncIterator[bytes]:
+    accumulator = OpenAIAccumulator()
+    try:
+        async for chunk in upstream_resp.aiter_raw():
+            accumulator.feed_bytes(chunk)
+            yield chunk
+    finally:
+        await upstream_resp.aclose()
+    reasoning = accumulator.reasoning
+    if not reasoning:
+        return
+    h = hash_openai_assistant(accumulator.content, accumulator.tool_calls)
+    cache.put(h, reasoning)
+
+
+async def _stream_and_capture_claude(
+    upstream_resp: httpx.Response,
+    cache: ReasoningCache,
+) -> AsyncIterator[bytes]:
+    accumulator = ClaudeAccumulator()
+    try:
+        async for chunk in upstream_resp.aiter_raw():
+            accumulator.feed_bytes(chunk)
+            yield chunk
+    finally:
+        await upstream_resp.aclose()
+    reasoning = accumulator.reasoning
+    if not reasoning:
+        return
+    h = hash_claude_assistant(accumulator.content_blocks)
+    cache.put(h, reasoning)
+
+
+def _capture_openai_nonstream(body_text: str, cache: ReasoningCache) -> None:
+    body = _parse_body(body_text.encode("utf-8"))
+    if body is None:
+        return
+    choices = dict_get_list(body, "choices")
+    if not choices:
+        return
+    first = as_dict(choices[0])
+    if first is None:
+        return
+    message = dict_get_dict(first, "message")
+    if message is None:
+        return
+    reasoning = dict_get_str(message, "reasoning_content")
+    if not reasoning:
+        return
+    content = _content_for_hash(message.get("content"))
+    tool_calls_raw = dict_get_list(message, "tool_calls") or []
+    tool_calls: list[dict[str, object]] = [
+        d for d in (as_dict(t) for t in tool_calls_raw) if d is not None
+    ]
+    h = hash_openai_assistant(content, tool_calls or None)
+    cache.put(h, reasoning)
+
+
+def _capture_claude_nonstream(body_text: str, cache: ReasoningCache) -> None:
+    body = _parse_body(body_text.encode("utf-8"))
+    if body is None:
+        return
+    content = as_list(body.get("content"))
+    if content is None:
+        return
+    reasoning_parts: list[str] = []
+    public_blocks: list[dict[str, object]] = []
+    for raw in content:
+        block = as_dict(raw)
+        if block is None:
+            continue
+        if block.get("type") == "thinking":
+            thinking = dict_get_str(block, "thinking")
+            if thinking is not None:
+                reasoning_parts.append(thinking)
+        else:
+            public_blocks.append(block)
+    reasoning = "".join(reasoning_parts)
+    if not reasoning:
+        return
+    h = hash_claude_assistant(public_blocks)
+    cache.put(h, reasoning)
+
+
+def _messages_from_body(body: dict[str, object]) -> list[dict[str, object]] | None:
+    raw = as_list(body.get("messages"))
+    if raw is None:
+        return None
+    result: list[dict[str, object]] = []
+    for item in raw:
+        m = as_dict(item)
+        if m is None:
+            return None
+        result.append(m)
+    return result
+
+
+def register_routes(
+    app: FastAPI,
+    client: httpx.AsyncClient,
+    upstream: str,
+    cache: ReasoningCache,
+    registry: DialectRegistry,
+) -> None:
+    async def handle_chat_completions(request: Request) -> Response:
+        raw_body = await request.body()
+        body = _parse_body(raw_body)
+        if body is None:
+            return _bad_request("invalid JSON body")
+
+        model_id = dict_get_str(body, "model")
+        if model_id is None:
+            return _bad_request("missing or invalid 'model' field")
+
+        dialect_name = await registry.resolve(model_id)
+        dialect = get_dialect(dialect_name)
+
+        messages = _messages_from_body(body)
+        if messages is not None:
+            has_tools = bool(body.get("tools"))
+            indices = dialect.select_attach_indices(messages, has_tools=has_tools)
+            if indices:
+                _attach_openai_reasoning(messages, indices, cache)
+                body["messages"] = messages
+
+        forward_body = json.dumps(body).encode("utf-8")
+        forward_headers = _copy_headers(request.headers)
+        forward_headers["content-length"] = str(len(forward_body))
+
+        is_stream = bool(body.get("stream"))
+
+        try:
+            if is_stream:
+                req = client.build_request(
+                    "POST",
+                    f"{upstream}/v1/chat/completions",
+                    content=forward_body,
+                    headers=forward_headers,
+                )
+                upstream_resp = await client.send(req, stream=True)
+                return StreamingResponse(
+                    _stream_and_capture_openai(upstream_resp, cache),
+                    status_code=upstream_resp.status_code,
+                    media_type=_media_type(upstream_resp.headers, "text/event-stream"),
+                    headers=_response_headers(upstream_resp.headers),
+                )
+            upstream_resp = await client.post(
+                f"{upstream}/v1/chat/completions",
+                content=forward_body,
+                headers=forward_headers,
+            )
+            text = upstream_resp.text
+            if upstream_resp.status_code == 200:
+                _capture_openai_nonstream(text, cache)
+            return Response(
+                content=text,
+                status_code=upstream_resp.status_code,
+                media_type=_media_type(upstream_resp.headers, "application/json"),
+                headers=_response_headers(upstream_resp.headers),
+            )
+        except httpx.RequestError as exc:
+            logger.warning("Upstream request failed: %s", exc)
+            return _bad_gateway(str(exc))
+
+    async def handle_claude_messages(request: Request) -> Response:
+        raw_body = await request.body()
+        body = _parse_body(raw_body)
+        if body is None:
+            return _bad_request("invalid JSON body")
+
+        model_id = dict_get_str(body, "model")
+        if model_id is None:
+            return _bad_request("missing or invalid 'model' field")
+
+        dialect_name = await registry.resolve(model_id)
+        dialect = get_dialect(dialect_name)
+
+        messages = _messages_from_body(body)
+        if messages is not None:
+            has_tools = bool(body.get("tools"))
+            indices = dialect.select_attach_indices(messages, has_tools=has_tools)
+            if indices:
+                _attach_claude_reasoning(messages, indices, cache)
+                body["messages"] = messages
+
+        forward_body = json.dumps(body).encode("utf-8")
+        forward_headers = _copy_headers(request.headers)
+        forward_headers["content-length"] = str(len(forward_body))
+
+        is_stream = bool(body.get("stream"))
+
+        try:
+            if is_stream:
+                req = client.build_request(
+                    "POST",
+                    f"{upstream}/v1/messages",
+                    content=forward_body,
+                    headers=forward_headers,
+                )
+                upstream_resp = await client.send(req, stream=True)
+                return StreamingResponse(
+                    _stream_and_capture_claude(upstream_resp, cache),
+                    status_code=upstream_resp.status_code,
+                    media_type=_media_type(upstream_resp.headers, "text/event-stream"),
+                    headers=_response_headers(upstream_resp.headers),
+                )
+            upstream_resp = await client.post(
+                f"{upstream}/v1/messages",
+                content=forward_body,
+                headers=forward_headers,
+            )
+            text = upstream_resp.text
+            if upstream_resp.status_code == 200:
+                _capture_claude_nonstream(text, cache)
+            return Response(
+                content=text,
+                status_code=upstream_resp.status_code,
+                media_type=_media_type(upstream_resp.headers, "application/json"),
+                headers=_response_headers(upstream_resp.headers),
+            )
+        except httpx.RequestError as exc:
+            logger.warning("Upstream request failed: %s", exc)
+            return _bad_gateway(str(exc))
+
+    async def health() -> dict[str, object]:
+        return {"status": "ok", "cache_entries": cache.size()}
+
+    _ = app.post("/v1/chat/completions")(handle_chat_completions)
+    _ = app.post("/v1/messages")(handle_claude_messages)
+    _ = app.get("/health")(health)
+
+
+_HOP_BY_HOP = {
+    "connection",
+    "keep-alive",
+    "proxy-authenticate",
+    "proxy-authorization",
+    "te",
+    "trailers",
+    "transfer-encoding",
+    "upgrade",
+    "content-length",
+    "host",
+}
+
+
+def _copy_headers(headers: httpx.Headers | StarletteHeaders) -> dict[str, str]:
+    out: dict[str, str] = {}
+    for k, v in headers.items():
+        if k.lower() in _HOP_BY_HOP:
+            continue
+        out[k] = v
+    return out
+
+
+def _response_headers(headers: httpx.Headers) -> dict[str, str]:
+    return _copy_headers(headers)
+
+
+def _media_type(headers: httpx.Headers, default: str) -> str:
+    value = cast(object, headers.get("content-type", default))
+    return value if isinstance(value, str) else default
+
+
+def _bad_request(msg: str) -> JSONResponse:
+    return JSONResponse(
+        status_code=400,
+        content={"error": {"message": msg, "type": "invalid_request_error"}},
+    )
+
+
+def _bad_gateway(msg: str) -> JSONResponse:
+    return JSONResponse(
+        status_code=502,
+        content={
+            "error": {"message": f"upstream unreachable: {msg}", "type": "bad_gateway"}
+        },
+    )

src/exo/shared/models/model_cards.py

@@ -28,6 +28,7 @@ from exo.shared.constants import (
 )
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
+from exo.shared.types.text_generation import ReasoningDialect
 from exo.utils.pydantic_ext import FrozenModel
 
 # kinda ugly...
@@ -145,6 +146,7 @@ class ModelCard(FrozenModel):
     quantization: str = ""
     base_model: str = ""
     capabilities: list[str] = []
+    reasoning_dialect: ReasoningDialect = "none"
     context_length: int = 0
     uses_cfg: bool = False
     trust_remote_code: bool = True
@@ -270,6 +272,7 @@ class ConfigData(BaseModel):
         return self.architectures in [
             ["Glm4MoeLiteForCausalLM"],
             ["GlmMoeDsaForCausalLM"],
+            ["DeepseekV4ForCausalLM"],
             ["DeepseekV32ForCausalLM"],
             ["DeepseekV3ForCausalLM"],
             ["Qwen3NextForCausalLM"],

src/exo/shared/types/text_generation.py

@@ -13,6 +13,9 @@ from exo.shared.types.common import ModelId, TruncatingString
 
 MessageRole = Literal["user", "assistant", "system", "developer", "tool"]
 ReasoningEffort = Literal["none", "minimal", "low", "medium", "high", "xhigh"]
+ReasoningDialect = Literal[
+    "none", "post_last_user", "suffix", "channel", "tool_conditional"
+]
 
 
 def resolve_reasoning_params(

src/exo/worker/engines/mlx/auto_parallel.py

@@ -17,6 +17,8 @@ from mlx_lm.models.base import (
 from mlx_lm.models.cache import ArraysCache, KVCache
 from mlx_lm.models.deepseek_v3 import DeepseekV3MLP
 from mlx_lm.models.deepseek_v3 import Model as DeepseekV3Model
+from mlx_lm.models.deepseek_v4 import DeepseekV4MoE, V4Attention
+from mlx_lm.models.deepseek_v4 import Model as DeepseekV4Model
 from mlx_lm.models.deepseek_v32 import DeepseekV32MLP
 from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model
 from mlx_lm.models.gemma4 import Model as Gemma4Model
@@ -295,6 +297,7 @@ def pipeline_auto_parallel(
     total = len(layers)
     for i, layer in enumerate(layers):
         mx.eval(layer)  # type: ignore
+        mx.clear_cache()
         yield ModelLoadingResponse(layers_loaded=i, total=total)
 
     layers[0] = PipelineFirstLayer(layers[0], device_rank, group=group)
@@ -510,6 +513,14 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
         )
+    elif isinstance(model, DeepseekV4Model):
+        tensor_parallel_sharding_strategy = DeepseekV4ShardingStrategy(
+            group,
+            all_to_sharded_linear,
+            sharded_to_all_linear,
+            all_to_sharded_linear_in_place,
+            sharded_to_all_linear_in_place,
+        )
     elif isinstance(model, MiniMaxModel):
         tensor_parallel_sharding_strategy = MiniMaxShardingStrategy(
             group,
@@ -648,7 +659,14 @@ def _set_layers(model: nn.Module, layers: list[_LayerCallable]) -> None:
 
         # Update DeepSeek V3 specific parameters when layers are shrunk
         if isinstance(
-            model, (DeepseekV3Model, DeepseekV32Model, Glm4MoeModel, KimiK25Model)
+            model,
+            (
+                DeepseekV3Model,
+                DeepseekV32Model,
+                DeepseekV4Model,
+                Glm4MoeModel,
+                KimiK25Model,
+            ),
         ) and hasattr(inner_model_instance, "num_layers"):
             logger.info(
                 f"Setting num_layers to {len(layers)} for model {model.model.__class__.__name__}"
@@ -749,6 +767,178 @@ class ShardedMoE(CustomMlxLayer):
         return y
 
 
+class ShardedMoEV4(CustomMlxLayer):
+    """Same as ShardedMoE but for DeepseekV4MoE which takes (x, input_ids)."""
+
+    def __init__(self, layer: DeepseekV4MoE):
+        super().__init__(cast(_LayerCallable, cast(object, layer)))
+        self._v4_inner = layer
+        self.sharding_group: mx.distributed.Group | None = None
+
+    def __call__(self, x: mx.array, input_ids: mx.array) -> mx.array:
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(x)
+        y = self._v4_inner(x, input_ids)
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+        return y
+
+
+def _shard_quantized_rows(
+    q: nn.QuantizedLinear,
+    head_dim: int,
+    slicer: Callable[[mx.array, int], mx.array],
+) -> None:
+    weight = q["weight"]
+    scales = q["scales"]
+    assert isinstance(weight, mx.array)
+    assert isinstance(scales, mx.array)
+    q.weight = slicer(weight, head_dim)
+    q.scales = slicer(scales, head_dim)
+    biases = q.get("biases")
+    if isinstance(biases, mx.array):
+        q.biases = slicer(biases, head_dim)
+
+
+class _AllSumLinear(nn.Module):
+    """Wraps an unsharded wo_b that takes a head-sharded partial wo_a output.
+
+    Flow per rank:
+      1. all_sum the incoming partial wo_a output (summed across the head
+         input shards → full wo_a_out on every rank)
+      2. apply the unsharded wo_b → full hidden on every rank
+
+    One collective per layer on the smaller of (n_groups * o_lora_rank) vs
+    hidden. wo_b compute is replicated, but at decode B=1 it's only ~30M FLOPs
+    per layer and 61 extra all_gathers/token cost more than running wo_b on
+    every rank.
+    """
+
+    def __init__(self, inner: nn.Module, group: mx.distributed.Group):
+        super().__init__()
+        self.inner = inner
+        self._group = group
+
+    def __call__(self, x: mx.array) -> mx.array:
+        x = mx.distributed.all_sum(x, group=self._group)
+        return cast(Callable[[mx.array], mx.array], self.inner)(x)
+
+
+def _shard_v4_attention_heads(
+    attn: V4Attention,
+    world_size: int,
+    rank: int,
+) -> None:
+    """Interleaved-per-group head sharding for V4Attention.
+
+    V4 uses a grouped low-rank output projection: `_grouped_output_projection`
+    reshapes the flat `n_heads * head_dim` dim into `(o_groups, heads_per_group,
+    head_dim)`, so group g owns heads `[g * heads_per_group : (g+1) * heads_per_group]`.
+
+    A naive contiguous `shard_linear("all-to-sharded")` on wq_b puts whole
+    original groups on each rank — the per-rank "group g" ends up containing
+    heads that don't belong to original group g. That breaks the wo_a grouped
+    weight mapping. We instead slice heads interleaved-by-group: each rank
+    owns `heads_per_group / N` heads *from every original group*, kept in
+    group-major order so SDPA → reshape → wo_a preserves the group mapping.
+
+    Affects `wq_b.weight` / `wq_b.bias`, `attn_sink`. wo_a is sharded via a
+    normal input-dim block split (the default axis-(-1) behavior of
+    shard_inplace), which now correctly aligns with the interleaved head
+    layout because the last dim of out after reshape is `heads_per_group/N *
+    head_dim` per group.
+    """
+    n_heads: int = attn.n_heads
+    head_dim: int = attn.head_dim
+    o_groups: int = attn.n_groups
+    assert n_heads % o_groups == 0, "n_heads must be divisible by o_groups"
+    heads_per_group = n_heads // o_groups
+    assert heads_per_group % world_size == 0, (
+        f"heads_per_group ({heads_per_group}) must be divisible by world_size "
+        f"({world_size}) for interleaved per-group head sharding"
+    )
+    hpg_per_rank = heads_per_group // world_size
+    start = rank * hpg_per_rank
+    end = start + hpg_per_rank
+
+    def _slice_head_major_flat(arr: mx.array, stride: int) -> mx.array:
+        """Slice arr on axis 0 where the flat 0-axis is (o_groups *
+        heads_per_group * stride), returning a fresh contiguous allocation
+        so the full unsharded array can be freed. Without the contiguous
+        copy the slice is a view and the original weight stays resident —
+        OOM on large V4. Quantized packed weights don't round-trip through
+        numpy so we use mx.contiguous directly."""
+        rest = arr.shape[1:]
+        reshaped = arr.reshape(o_groups, heads_per_group, stride, *rest)
+        sliced = reshaped[:, start:end].reshape(o_groups * hpg_per_rank * stride, *rest)
+        detached = mx.contiguous(sliced)
+        mx.eval(detached)
+        return detached
+
+    wq_b: nn.Module = attn.wq_b
+    if isinstance(wq_b, nn.QuantizedLinear):
+        # Packed weight: (n_heads*head_dim, q_lora_rank/el_per_int).
+        # scales/biases: (n_heads*head_dim, q_lora_rank/group_size).
+        # Slice axis 0 interleaved-by-group with head_dim stride.
+        _shard_quantized_rows(wq_b, head_dim, _slice_head_major_flat)
+    else:
+        dense = wq_b
+        assert isinstance(dense, nn.Linear)
+        w = dense.weight
+        q_lora_rank = w.shape[-1]
+        w_sharded = _slice_head_major_flat(w, head_dim)
+        has_bias = "bias" in dense
+        new_wq_b = nn.Linear(q_lora_rank, w_sharded.shape[0], bias=has_bias)
+        new_wq_b.weight = w_sharded
+        if has_bias:
+            b = dense.bias
+            assert b is not None
+            new_wq_b.bias = _slice_head_major_flat(b[:, None], head_dim).reshape(-1)
+        attn.wq_b = new_wq_b
+
+    sink = attn.attn_sink
+    reshaped = sink.reshape(o_groups, heads_per_group)[:, start:end].reshape(-1)
+    detached_sink = mx.contiguous(reshaped)
+    mx.eval(detached_sink)
+    attn.attn_sink = detached_sink
+    attn.n_heads = o_groups * hpg_per_rank
+
+
+class DeepseekV4ShardingStrategy(TensorParallelShardingStrategy):
+    def shard_model(
+        self,
+        model: nn.Module,
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
+        model = cast(DeepseekV4Model, model)
+        total = len(model.layers)
+
+        for i, layer in enumerate(model.layers):
+            mx.eval(layer.parameters())
+
+            # Head-parallel attention with interleaved-per-group sharding.
+            _shard_v4_attention_heads(layer.attn, self.N, self.group.rank())
+            self.sharded_to_all_linear_in_place(layer.attn.wo_a)
+            layer.attn.wo_b = _AllSumLinear(layer.attn.wo_b, self.group)  # type: ignore[assignment]
+
+            ffn = layer.ffn
+            if getattr(ffn, "shared_experts", None) is not None:
+                self.all_to_sharded_linear_in_place(ffn.shared_experts.gate_proj)
+                self.sharded_to_all_linear_in_place(ffn.shared_experts.down_proj)
+                self.all_to_sharded_linear_in_place(ffn.shared_experts.up_proj)
+            self.all_to_sharded_linear_in_place(ffn.switch_mlp.gate_proj)
+            self.sharded_to_all_linear_in_place(ffn.switch_mlp.down_proj)
+            self.all_to_sharded_linear_in_place(ffn.switch_mlp.up_proj)
+            wrapped = ShardedMoEV4(ffn)
+            wrapped.sharding_group = self.group
+            layer.ffn = wrapped  # type: ignore[assignment]
+
+            mx.eval(layer)
+            mx.clear_cache()
+            yield ModelLoadingResponse(layers_loaded=i, total=total)
+
+        return model
+
+
 class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
     def shard_model(
         self,
@@ -804,6 +994,7 @@ class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
                 layer.mlp = ShardedMoE(layer.mlp)  # type: ignore
                 layer.mlp.sharding_group = self.group  # type: ignore
             mx.eval(layer)
+            mx.clear_cache()
 
             yield ModelLoadingResponse(layers_loaded=i, total=total)
 
@@ -921,6 +1112,7 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
             layer.block_sparse_moe = ShardedMoE(layer.block_sparse_moe)  # type: ignore
             layer.block_sparse_moe.sharding_group = self.group
             mx.eval(layer)
+            mx.clear_cache()
 
             yield ModelLoadingResponse(layers_loaded=i, total=total)
         return model
@@ -1085,6 +1277,7 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                 layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
 
             mx.eval(layer)
+            mx.clear_cache()
 
             yield ModelLoadingResponse(layers_loaded=i, total=total)
         return model
@@ -1130,6 +1323,7 @@ class Glm4MoeShardingStrategy(TensorParallelShardingStrategy):
                 layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
 
             mx.eval(layer)
+            mx.clear_cache()
 
             yield ModelLoadingResponse(layers_loaded=i, total=total)
         return model
@@ -1170,6 +1364,7 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
             layer.mlp = ShardedMoE(layer.mlp)  # type: ignore
             layer.mlp.sharding_group = self.group
             mx.eval(layer)
+            mx.clear_cache()
 
             yield ModelLoadingResponse(layers_loaded=i, total=total)
         return model
@@ -1212,6 +1407,7 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy):
                 self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)
 
             mx.eval(layer)
+            mx.clear_cache()
 
             yield ModelLoadingResponse(layers_loaded=i, total=total)
         return model
@@ -1254,6 +1450,7 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
                 layer.mixer = mixer  # pyright: ignore[reportAttributeAccessIssue]
 
             mx.eval(layer)
+            mx.clear_cache()
             yield ModelLoadingResponse(layers_loaded=i, total=total)
         return model
 
@@ -1391,5 +1588,6 @@ class Gemma4ShardingStrategy(TensorParallelShardingStrategy):
                 layer.experts.sharding_group = self.group
 
             mx.eval(layer)
+            mx.clear_cache()
             yield ModelLoadingResponse(layers_loaded=i, total=total)
         return model

src/exo/worker/engines/mlx/cache.py

@@ -46,7 +46,9 @@ class CacheSnapshot:
     """Snapshot of states at a known token position."""
 
     def __init__(
-        self, states: list[RotatingKVCache | ArraysCache | None], token_count: int
+        self,
+        states: list[RotatingKVCache | ArraysCache | CacheList | None],
+        token_count: int,
     ):
         self.states = states
         self.token_count = token_count
@@ -83,13 +85,55 @@ def copy_rotating_kv_cache(cache: RotatingKVCache) -> RotatingKVCache | None:
     return snap
 
 
+def _copy_arrays_cache(ac: ArraysCache) -> ArraysCache:
+    entries: list[mx.array | None] = []
+    for entry in ac.cache:  # type: ignore[reportUnknownMemberType]
+        if entry is None:
+            entries.append(None)
+            continue
+        assert isinstance(entry, mx.array)
+        entries.append(_detached_copy(entry))
+    copy = ArraysCache(len(entries))
+    copy.cache = entries  # type: ignore[reportUnknownMemberType]
+    return copy
+
+
+def _copy_cache_list(cl: CacheList) -> CacheList:
+    inners: list[object] = list(cl)  # type: ignore[reportUnknownArgumentType]
+    copied: list[object] = []
+    for inner in inners:
+        if isinstance(inner, RotatingKVCache):
+            snap = copy_rotating_kv_cache(inner)
+            copied.append(snap if snap is not None else deepcopy(inner))
+        elif isinstance(inner, ArraysCache):
+            copied.append(_copy_arrays_cache(inner))
+        else:
+            copied.append(deepcopy(inner))
+    return CacheList(*copied)
+
+
+def restore_snapshot_entry(
+    entry: ArraysCache | RotatingKVCache | CacheList | None,
+) -> ArraysCache | RotatingKVCache | CacheList | None:
+    if entry is None:
+        return None
+    if isinstance(entry, RotatingKVCache):
+        snap = copy_rotating_kv_cache(entry)
+        return snap if snap is not None else deepcopy(entry)
+    if isinstance(entry, ArraysCache):
+        return _copy_arrays_cache(entry)
+    return _copy_cache_list(entry)
+
+
 def snapshot_ssm_states(cache: KVCacheType) -> CacheSnapshot:
-    states: list[ArraysCache | RotatingKVCache | None] = []
+    states: list[ArraysCache | RotatingKVCache | CacheList | None] = []
     for c in cache:
         if isinstance(c, ArraysCache):
-            states.append(deepcopy(c))
+            states.append(_copy_arrays_cache(c))
         elif isinstance(c, RotatingKVCache):
             states.append(copy_rotating_kv_cache(c))
+        elif isinstance(c, CacheList) and not bool(c.is_trimmable()):  # type: ignore[reportUnknownMemberType]
+            states.append(_copy_cache_list(c))
         else:
             states.append(None)
     token_count = cache_length(cache)
@@ -111,7 +155,12 @@ def _find_nearest_snapshot(
 
 def has_non_kv_caches(cache: KVCacheType) -> bool:
     """Check if a cache contains any ArraysCache (SSM) entries."""
-    return any(isinstance(c, (ArraysCache, RotatingKVCache)) for c in cache)
+    for c in cache:
+        if isinstance(c, CacheList):
+            return any(isinstance(_c, (ArraysCache, RotatingKVCache)) for _c in c)  # type: ignore[reportUnknownVariableType]
+        elif isinstance(c, (ArraysCache, RotatingKVCache)):
+            return True
+    return False
 
 
 class KVPrefixCache:
@@ -357,11 +406,21 @@ def trim_cache(
     snapshot: CacheSnapshot | None = None,
 ) -> None:
     for i, c in enumerate(cache):
-        if isinstance(c, (ArraysCache, RotatingKVCache)):
+        non_trimmable = isinstance(c, (ArraysCache, RotatingKVCache)) or (
+            isinstance(c, CacheList) and not bool(c.is_trimmable())  # type: ignore[reportUnknownMemberType]
+        )
+        if non_trimmable:
             if snapshot is not None and snapshot.states[i] is not None:
-                cache[i] = deepcopy(snapshot.states[i])  # type: ignore
-            else:
+                restored = restore_snapshot_entry(snapshot.states[i])
+                if restored is not None:
+                    cache[i] = restored  # type: ignore
+            elif isinstance(c, (ArraysCache, RotatingKVCache)):
                 c.state = [None] * len(c.state)
+            else:
+                # CacheList without a snapshot — zero each inner cache's state
+                for inner in c:  # type: ignore[reportUnknownVariableType]
+                    if isinstance(inner, (ArraysCache, RotatingKVCache)):
+                        inner.state = [None] * len(inner.state)
         else:
             c.trim(num_tokens)
 

src/exo/worker/engines/mlx/deepseek_v4_encoding.py

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

src/exo/worker/engines/mlx/generator/generate.py

@@ -2,7 +2,6 @@ import contextlib
 import functools
 import math
 import time
-from copy import deepcopy
 from typing import Callable, Generator, cast, get_args
 
 import mlx.core as mx
@@ -10,7 +9,7 @@ from mlx_lm.generate import (
     maybe_quantize_kv_cache,
     stream_generate,
 )
-from mlx_lm.models.cache import ArraysCache, RotatingKVCache
+from mlx_lm.models.cache import ArraysCache, CacheList, RotatingKVCache
 from mlx_lm.sample_utils import make_logits_processors, make_sampler
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
@@ -47,6 +46,7 @@ from exo.worker.engines.mlx.cache import (
     encode_prompt,
     has_non_kv_caches,
     make_kv_cache,
+    restore_snapshot_entry,
     snapshot_ssm_states,
 )
 from exo.worker.engines.mlx.constants import (
@@ -370,14 +370,18 @@ def prefill(
 
     # stream_generate added 1 extra generated token to the cache, so we should trim it.
     # Because of needing to roll back arrays cache, we will generate on 2 tokens so trim 1 more.
-    pre_gen = deepcopy(snapshots[-2]) if has_ssm else None
+    pre_gen = snapshots[-2] if has_ssm else None
     for i, c in enumerate(cache):
-        if has_ssm and isinstance(c, (ArraysCache, RotatingKVCache)):
+        non_trimmable = isinstance(c, (ArraysCache, RotatingKVCache)) or (
+            isinstance(c, CacheList) and not bool(c.is_trimmable())  # type: ignore[reportUnknownMemberType]
+        )
+        if has_ssm and non_trimmable:
             assert pre_gen is not None
-            if pre_gen.states[i] is not None:
-                cache[i] = deepcopy(pre_gen.states[i])  # type: ignore
+            restored = restore_snapshot_entry(pre_gen.states[i])
+            if restored is not None:
+                cache[i] = restored  # type: ignore
         else:
-            assert not isinstance(c, (ArraysCache, RotatingKVCache))
+            assert not non_trimmable
             c.trim(2)
 
     elapsed = time.perf_counter() - start_time

src/exo/worker/engines/mlx/patches/__init__.py

@@ -1,5 +1,6 @@
 from exo.worker.engines.mlx.patches.opt_batch_gen import apply_batch_gen_patch
 from exo.worker.engines.mlx.patches.standard_yarn_rope import patch_yarn_rope
+from exo.worker.engines.mlx.patches.v4_offset_sync import apply as apply_v4_offset_sync
 
 _applied = False
 
@@ -11,3 +12,4 @@ def apply_mlx_patches() -> None:
     _applied = True
     patch_yarn_rope()
     apply_batch_gen_patch()
+    apply_v4_offset_sync()

src/exo/worker/engines/mlx/patches/opt_batch_gen.py

@@ -58,6 +58,7 @@ def _patched_step(self: GenerationBatch) -> tuple[list[int], list[mx.array]]:
     self._current_tokens = self._next_tokens
     self._current_logprobs = self._next_logprobs
     inputs = self._current_tokens
+    assert inputs is not None, "_step requires initialized _next_tokens"
 
     buf = _get_buffer(self)
     buf.ready = buf.pending
@@ -87,7 +88,7 @@ def _patched_step(self: GenerationBatch) -> tuple[list[int], list[mx.array]]:
         sampled = self.fallback_sampler(logprobs)
 
     self._next_tokens = sampled
-    self._next_logprobs = list(logprobs)
+    self._next_logprobs = logprobs
 
     if buf.needs_topk:
         batch_size = len(self.uids)
@@ -106,19 +107,29 @@ def _patched_step(self: GenerationBatch) -> tuple[list[int], list[mx.array]]:
         )
         mx.async_eval(
             self._next_tokens,
-            *self._next_logprobs,
+            self._next_logprobs,
             pending_indices,
             pending_values,
             pending_selected,
         )
     else:
-        mx.async_eval(self._next_tokens, *self._next_logprobs)
+        mx.async_eval(self._next_tokens, self._next_logprobs)
+
+    current_lp = self._current_logprobs
+    if isinstance(current_lp, mx.array):
+        mx.eval(inputs, current_lp)
+    elif current_lp:
+        mx.eval(inputs, *current_lp)
+    else:
+        mx.eval(inputs)
 
-    mx.eval(inputs, *self._current_logprobs)
     token_list = cast(list[int], inputs.tolist())
     for sti, ti in zip(self.tokens, token_list, strict=True):
         sti.append(ti)
-    return token_list, self._current_logprobs
+
+    if isinstance(current_lp, mx.array):
+        current_lp = list(current_lp)
+    return token_list, current_lp
 
 
 def apply_batch_gen_patch() -> None:

src/exo/worker/engines/mlx/patches/v4_offset_sync.py

@@ -0,0 +1,74 @@
+from typing import Callable, cast
+
+import mlx.core as mx
+from mlx_lm.models.deepseek_v4 import Compressor, DeepseekV4Model
+
+_current_int_offset: int | None = None
+_applied: bool = False
+
+
+def _extract_int_offset(cache: object | None) -> int | None:
+    if cache is None:
+        return None
+    for entry in cast(list[object], cache):
+        inner_caches = getattr(entry, "caches", None)
+        win = inner_caches[0] if inner_caches is not None else entry
+        int_off = getattr(win, "_offset", None)
+        if isinstance(int_off, int):
+            return int_off
+        maybe_int = getattr(win, "offset", None)
+        if isinstance(maybe_int, int):
+            return maybe_int
+    return None
+
+
+_ModelCall = Callable[[DeepseekV4Model, mx.array, list[object] | None], mx.array]
+_CompressorCall = Callable[
+    [Compressor, mx.array, object, object, int, int, int], mx.array | None
+]
+
+
+def apply() -> None:
+    global _applied
+    if _applied:
+        return
+    _applied = True
+
+    original_model_call = cast(_ModelCall, DeepseekV4Model.__call__)
+
+    def patched_model_call(
+        self: DeepseekV4Model,
+        inputs: mx.array,
+        cache: list[object] | None = None,
+    ) -> mx.array:
+        global _current_int_offset
+        prev = _current_int_offset
+        _current_int_offset = _extract_int_offset(cache)
+        try:
+            return original_model_call(self, inputs, cache)
+        finally:
+            _current_int_offset = prev
+
+    DeepseekV4Model.__call__ = patched_model_call
+
+    original_compressor_call = cast(_CompressorCall, Compressor.__call__)
+
+    def patched_compressor_call(
+        self: Compressor,
+        x: mx.array,
+        state: object,
+        offset: object,
+        slot_compressed: int,
+        slot_kv_state: int,
+        slot_score_state: int,
+    ) -> mx.array | None:
+        if isinstance(offset, mx.array) and _current_int_offset is not None:
+            offset = _current_int_offset
+        return original_compressor_call(
+            self, x, state, offset, slot_compressed, slot_kv_state, slot_score_state
+        )
+
+    Compressor.__call__ = patched_compressor_call
+
+
+apply()

src/exo/worker/engines/mlx/utils_mlx.py

@@ -165,6 +165,8 @@ def load_mlx_items(
 ) -> Generator[
     ModelLoadingResponse, None, tuple[Model, TokenizerWrapper, "VisionProcessor | None"]
 ]:
+    set_wired_limit_for_model(get_weights_size(bound_instance.bound_shard))
+
     if group is None:
         logger.info(f"Single device used for {bound_instance.instance}")
         model_path = build_model_path(bound_instance.bound_shard.model_card.model_id)
@@ -199,8 +201,6 @@ def load_mlx_items(
             f"Time taken to shard and load model: {(end_time - start_time):.2f}s"
         )
 
-    set_wired_limit_for_model(get_weights_size(bound_instance.bound_shard))
-
     mx.clear_cache()
 
     vision_config = bound_instance.bound_shard.model_card.vision
@@ -486,6 +486,34 @@ def _needs_dsml_encoding(task_params: TextGenerationTaskParams) -> bool:
     return "deepseek-v3.2" in task_params.model.lower()
 
 
+def _needs_v4_encoding(task_params: TextGenerationTaskParams) -> bool:
+    return "deepseek-v4" in task_params.model.lower()
+
+
+def _v4_reasoning_effort(task_params: TextGenerationTaskParams) -> str | None:
+    effort = task_params.reasoning_effort
+    if effort == "xhigh":
+        return "max"
+    if effort == "high":
+        return "high"
+    return None
+
+
+_V4_THINK_BLOCK_RE = re.compile(r"<think>.*?</think>", re.DOTALL)
+
+
+def _strip_v4_thinking_markers(content: str) -> str:
+    """Remove `<think>…</think>` blocks and any stray `<think>`/`</think>` tags
+    from prior-turn assistant content.
+
+    The V4 encoder drops `reasoning_content` for older turns when
+    `drop_thinking=True`"""
+    if not content:
+        return content
+    cleaned = _V4_THINK_BLOCK_RE.sub("", content)
+    return cleaned.replace("<think>", "").replace("</think>", "")
+
+
 def consolidate_system_messages(
     messages: list[dict[str, Any]],
 ) -> list[dict[str, Any]]:
@@ -549,6 +577,38 @@ def render_chat_template(
             prompt += partial_assistant_content
         return prompt
 
+    if _needs_v4_encoding(task_params):
+        from exo.worker.engines.mlx.deepseek_v4_encoding import (
+            encode_messages as encode_messages_v4,
+        )
+
+        v4_messages = [dict(m) for m in formatted_messages]
+        for msg in v4_messages:
+            if msg.get("role") == "assistant":
+                content = msg.get("content")
+                if isinstance(content, str):
+                    msg["content"] = _strip_v4_thinking_markers(content)
+        if task_params.tools:
+            for msg in v4_messages:
+                if msg.get("role") in ("system", "developer"):
+                    msg["tools"] = task_params.tools
+                    break
+            else:
+                v4_messages.insert(
+                    0, {"role": "system", "content": "", "tools": task_params.tools}
+                )
+
+        prompt = encode_messages_v4(
+            messages=v4_messages,
+            thinking_mode="chat"
+            if task_params.enable_thinking is False
+            else "thinking",
+            reasoning_effort=_v4_reasoning_effort(task_params),
+        )
+        if partial_assistant_content:
+            prompt += partial_assistant_content
+        return prompt
+
     for msg in formatted_messages:
         _normalize_tool_calls(msg)
 

src/exo/worker/runner/llm_inference/model_output_parsers.py

@@ -1,4 +1,4 @@
-from collections.abc import Generator, Iterator
+from collections.abc import Callable, Generator, Iterator
 from functools import cache
 from typing import Any
 
@@ -23,6 +23,7 @@ from exo.shared.types.chunks import (
 from exo.shared.types.common import ModelId
 from exo.shared.types.mlx import Model
 from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse
+from exo.worker.engines.mlx.dsml_encoding import parse_dsml_output
 from exo.worker.engines.mlx.utils_mlx import (
     detect_thinking_prompt_suffix,
 )
@@ -73,12 +74,10 @@ def apply_all_parsers(
 ) -> Iterator[GenerationChunk | None]:
     generator = receiver
 
+    normalized_id = model_id.normalize().lower()
     if issubclass(model_type, GptOssModel):
         generator = parse_gpt_oss(generator)
-    elif (
-        issubclass(model_type, DeepseekV32Model)
-        and "deepseek" in model_id.normalize().lower()
-    ):
+    elif issubclass(model_type, DeepseekV32Model) and "deepseek" in normalized_id:
         if tokenizer.has_thinking:
             generator = parse_thinking_models(
                 generator,
@@ -87,6 +86,15 @@ def apply_all_parsers(
                 starts_in_thinking=detect_thinking_prompt_suffix(prompt, tokenizer),
             )
         generator = parse_deepseek_v32(generator)
+    elif "deepseek-v4" in normalized_id:
+        if tokenizer.has_thinking:
+            generator = parse_thinking_models(
+                generator,
+                tokenizer.think_start,
+                tokenizer.think_end,
+                starts_in_thinking=detect_thinking_prompt_suffix(prompt, tokenizer),
+            )
+        generator = parse_deepseek_v4(generator)
     else:
         if tokenizer.has_thinking:
             generator = parse_thinking_models(
@@ -226,6 +234,26 @@ def parse_deepseek_v32(
         parse_dsml_output,
     )
 
+    return _parse_dsml_stream(
+        responses, TOOL_CALLS_START, TOOL_CALLS_END, parse_dsml_output
+    )
+
+
+def parse_deepseek_v4(
+    responses: Generator[GenerationResponse | None],
+) -> Generator[GenerationResponse | ToolCallResponse | None]:
+    dsml_token = "｜DSML｜"
+    start = f"<{dsml_token}tool_calls>"
+    end = f"</{dsml_token}tool_calls>"
+    return _parse_dsml_stream(responses, start, end, parse_dsml_output)
+
+
+def _parse_dsml_stream(
+    responses: Generator[GenerationResponse | None],
+    tool_calls_start: str,
+    tool_calls_end: str,
+    parse_body: Callable[[str], list[ToolCallItem] | None],
+) -> Generator[GenerationResponse | ToolCallResponse | None]:
     accumulated = ""
     in_tool_call = False
     # Tokens buffered while we detect the start of a DSML block
@@ -236,7 +264,7 @@ def parse_deepseek_v32(
     def _try_parse_tool_call(
         text: str, response: GenerationResponse
     ) -> ToolCallResponse | GenerationResponse:
-        parsed = parse_dsml_output(text)
+        parsed = parse_body(text)
         if parsed is not None:
             return ToolCallResponse(
                 tool_calls=parsed, usage=response.usage, stats=response.stats
@@ -256,11 +284,11 @@ def parse_deepseek_v32(
                 tool_call_text += response.text
                 yield (
                     _try_parse_tool_call(tool_call_text, response)
-                    if TOOL_CALLS_END in tool_call_text
+                    if tool_calls_end in tool_call_text
                     else response.model_copy(update={"text": tool_call_text})
                 )
-            elif TOOL_CALLS_START in response.text and TOOL_CALLS_END in response.text:
-                dsml_start = response.text.index(TOOL_CALLS_START)
+            elif tool_calls_start in response.text and tool_calls_end in response.text:
+                dsml_start = response.text.index(tool_calls_start)
                 before = response.text[:dsml_start]
                 if before:
                     yield response.model_copy(update={"text": before})
@@ -269,25 +297,21 @@ def parse_deepseek_v32(
                 yield response
             break
 
-        # ── Handle tool call accumulation ──
         if in_tool_call:
             tool_call_text += response.text
-            if TOOL_CALLS_END in tool_call_text:
+            if tool_calls_end in tool_call_text:
                 yield _try_parse_tool_call(tool_call_text, response)
                 in_tool_call = False
                 tool_call_text = ""
             continue
 
-        # ── Detect start of tool call block ──
         accumulated += response.text
 
-        if TOOL_CALLS_START in accumulated:
-            # The start marker might be split across pending_buffer + current token
-            start_idx = accumulated.index(TOOL_CALLS_START)
-            # Yield any pending tokens that are purely before the marker
+        if tool_calls_start in accumulated:
+            start_idx = accumulated.index(tool_calls_start)
             pre_text = accumulated[:start_idx]
+            # Flush pending buffer tokens that contributed text before the marker
             if pre_text:
-                # Flush pending buffer tokens that contributed text before the marker
                 for buf_resp in pending_buffer:
                     if not pre_text:
                         break
@@ -302,17 +326,14 @@ def parse_deepseek_v32(
             tool_call_text = accumulated[start_idx:]
             accumulated = ""
 
-            # Check if the end marker is already present (entire tool call in one token)
-            if TOOL_CALLS_END in tool_call_text:
+            if tool_calls_end in tool_call_text:
                 yield _try_parse_tool_call(tool_call_text, response)
                 tool_call_text = ""
             else:
                 in_tool_call = True
             continue
 
-        # Check if accumulated text might be the start of a DSML marker
-        # Buffer tokens if we see a partial match at the end
-        if _could_be_dsml_prefix(accumulated):
+        if _could_be_marker_prefix(accumulated, tool_calls_start):
             pending_buffer.append(response)
             continue
 
@@ -326,22 +347,12 @@ def parse_deepseek_v32(
     yield from pending_buffer
 
 
-def _could_be_dsml_prefix(text: str) -> bool:
-    """Check if the end of text could be the start of a DSML function_calls marker.
-
-    We look for suffixes of text that are prefixes of the TOOL_CALLS_START pattern.
-    This allows us to buffer tokens until we can determine if a tool call is starting.
-    """
-    from exo.worker.engines.mlx.dsml_encoding import TOOL_CALLS_START
-
-    # Only check the last portion of text that could overlap with the marker
-    max_check = len(TOOL_CALLS_START)
+def _could_be_marker_prefix(text: str, marker: str) -> bool:
+    max_check = len(marker)
     tail = text[-max_check:] if len(text) > max_check else text
-
-    # Check if any suffix of tail is a prefix of TOOL_CALLS_START
     for i in range(len(tail)):
         suffix = tail[i:]
-        if TOOL_CALLS_START.startswith(suffix):
+        if marker.startswith(suffix):
             return True
     return False
 

src/exo/worker/tests/unittests/test_mlx/test_tp_bit_exact.py

@@ -234,6 +234,42 @@ MODEL_CONFIGS = {
             rope_theta=10000.0,
         ),
     ),
+    "deepseek_v4": dict(
+        module="mlx_lm.models.deepseek_v4",
+        args=dict(
+            model_type="deepseek_v4",
+            vocab_size=256,
+            hidden_size=64,
+            num_hidden_layers=4,
+            num_attention_heads=4,
+            num_key_value_heads=1,
+            q_lora_rank=32,
+            o_lora_rank=32,
+            o_groups=1,
+            head_dim=16,
+            qk_rope_head_dim=8,
+            sliding_window=32,
+            compress_ratios=[0, 0, 4, 0, 0],
+            index_n_heads=4,
+            index_head_dim=16,
+            index_topk=16,
+            moe_intermediate_size=32,
+            n_routed_experts=4,
+            n_shared_experts=1,
+            num_experts_per_tok=2,
+            num_hash_layers=1,
+            hc_mult=1,
+            num_nextn_predict_layers=0,
+            max_position_embeddings=2048,
+            rope_scaling={
+                "beta_fast": 32,
+                "beta_slow": 1,
+                "factor": 2,
+                "original_max_position_embeddings": 1024,
+                "type": "yarn",
+            },
+        ),
+    ),
     "gemma4": dict(
         module="mlx_lm.models.gemma4",
         args=dict(

uv.lock

@@ -24,7 +24,7 @@ members = [
     "exo-bench",
     "exo-pyo3-bindings",
 ]
-constraints = [{ name = "transformers", specifier = ">=5.0.0,<5.4.0" }]
+constraints = [{ name = "transformers", specifier = ">=5.6.2" }]
 overrides = [
     { name = "mlx", marker = "sys_platform == 'darwin'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks" },
     { name = "mlx", marker = "sys_platform == 'linux'", specifier = "==0.31.1" },
@@ -394,8 +394,8 @@ dependencies = [
     { name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "mflux", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "mlx", version = "0.31.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
-    { name = "mlx", version = "0.31.2.dev20260422+ec49d18e", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#ec49d18ec4cfba0e0c7a37f20d1cf4d75fe56731" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
-    { name = "mlx-lm", version = "0.31.3", source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-arrayscache-leak#c7010341e1f41ac15815feb5dc55134f44e3b044" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
+    { name = "mlx", version = "0.32.0.dev20260425+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
+    { name = "mlx-lm", version = "0.31.3", source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Fdeepseek-v4#0f230c9d108592cea4639a4cac528ce9e1aa0017" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "mlx-vlm", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "openai-harmony", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
@@ -416,21 +416,21 @@ build = [
 ]
 cpu = [
     { name = "mlx", version = "0.31.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
-    { name = "mlx", version = "0.31.2.dev20260422+ec49d18e", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#ec49d18ec4cfba0e0c7a37f20d1cf4d75fe56731" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-cpu') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
+    { name = "mlx", version = "0.32.0.dev20260425+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-cpu') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "mlx-cpu", marker = "sys_platform == 'linux'" },
-    { name = "mlx-lm", version = "0.31.3", source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-arrayscache-leak#c7010341e1f41ac15815feb5dc55134f44e3b044" }, marker = "sys_platform == 'linux'" },
+    { name = "mlx-lm", version = "0.31.3", source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Fdeepseek-v4#0f230c9d108592cea4639a4cac528ce9e1aa0017" }, marker = "sys_platform == 'linux'" },
 ]
 cuda12 = [
     { name = "mlx", version = "0.31.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13')" },
-    { name = "mlx", version = "0.31.2.dev20260422+ec49d18e", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#ec49d18ec4cfba0e0c7a37f20d1cf4d75fe56731" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13')" },
+    { name = "mlx", version = "0.32.0.dev20260425+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13')" },
     { name = "mlx-cuda-12", marker = "sys_platform == 'linux'" },
-    { name = "mlx-lm", version = "0.31.3", source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-arrayscache-leak#c7010341e1f41ac15815feb5dc55134f44e3b044" }, marker = "sys_platform == 'linux'" },
+    { name = "mlx-lm", version = "0.31.3", source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Fdeepseek-v4#0f230c9d108592cea4639a4cac528ce9e1aa0017" }, marker = "sys_platform == 'linux'" },
 ]
 cuda13 = [
     { name = "mlx", version = "0.31.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13')" },
-    { name = "mlx", version = "0.31.2.dev20260422+ec49d18e", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#ec49d18ec4cfba0e0c7a37f20d1cf4d75fe56731" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13')" },
+    { name = "mlx", version = "0.32.0.dev20260425+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13')" },
     { name = "mlx-cuda-13", marker = "sys_platform == 'linux'" },
-    { name = "mlx-lm", version = "0.31.3", source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-arrayscache-leak#c7010341e1f41ac15815feb5dc55134f44e3b044" }, marker = "sys_platform == 'linux'" },
+    { name = "mlx-lm", version = "0.31.3", source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Fdeepseek-v4#0f230c9d108592cea4639a4cac528ce9e1aa0017" }, marker = "sys_platform == 'linux'" },
 ]
 
 [package.dev-dependencies]
@@ -463,10 +463,10 @@ requires-dist = [
     { name = "mlx-cpu", marker = "sys_platform == 'linux' and extra == 'cpu'", specifier = "==0.31.1" },
     { name = "mlx-cuda-12", marker = "sys_platform == 'linux' and extra == 'cuda12'", specifier = "==0.31.1" },
     { name = "mlx-cuda-13", marker = "sys_platform == 'linux' and extra == 'cuda13'", specifier = "==0.31.1" },
-    { name = "mlx-lm", marker = "sys_platform == 'darwin'", git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-arrayscache-leak" },
-    { name = "mlx-lm", marker = "sys_platform == 'linux' and extra == 'cpu'", git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-arrayscache-leak" },
-    { name = "mlx-lm", marker = "sys_platform == 'linux' and extra == 'cuda12'", git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-arrayscache-leak" },
-    { name = "mlx-lm", marker = "sys_platform == 'linux' and extra == 'cuda13'", git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-arrayscache-leak" },
+    { name = "mlx-lm", marker = "sys_platform == 'darwin'", git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Fdeepseek-v4" },
+    { name = "mlx-lm", marker = "sys_platform == 'linux' and extra == 'cpu'", git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Fdeepseek-v4" },
+    { name = "mlx-lm", marker = "sys_platform == 'linux' and extra == 'cuda12'", git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Fdeepseek-v4" },
+    { name = "mlx-lm", marker = "sys_platform == 'linux' and extra == 'cuda13'", git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Fdeepseek-v4" },
     { name = "mlx-vlm", specifier = ">=0.3.11" },
     { name = "msgspec", specifier = ">=0.19.0" },
     { name = "nanobind", marker = "extra == 'build'" },
@@ -483,7 +483,7 @@ requires-dist = [
     { name = "torch", marker = "(sys_platform == 'linux' and extra == 'cpu' and extra == 'cuda12') or (sys_platform == 'linux' and extra == 'cpu' and extra == 'cuda13')", specifier = ">=2.10.0" },
     { name = "torch", marker = "sys_platform == 'linux' and extra == 'cpu' and extra != 'cuda12' and extra != 'cuda13'", specifier = ">=2.10.0", index = "https://download.pytorch.org/whl/cpu" },
     { name = "torch", marker = "(sys_platform == 'linux' and extra == 'cpu' and extra == 'cuda13') or (sys_platform == 'linux' and extra == 'cuda12' and extra == 'cuda13')", specifier = ">=2.10.0" },
-    { name = "transformers", specifier = ">=5.0.0,<5.4.0" },
+    { name = "transformers", specifier = ">=5.6.2" },
     { name = "types-aiofiles", specifier = ">=24.1.0.20250708" },
     { name = "zstandard", specifier = ">=0.23.0" },
 ]
@@ -1207,7 +1207,7 @@ dependencies = [
     { name = "hf-transfer", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "huggingface-hub", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "matplotlib", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
-    { name = "mlx", version = "0.31.2.dev20260422+ec49d18e", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#ec49d18ec4cfba0e0c7a37f20d1cf4d75fe56731" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
+    { name = "mlx", version = "0.32.0.dev20260425+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "numpy", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "opencv-python", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "piexif", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
@@ -1257,8 +1257,8 @@ wheels = [
 
 [[package]]
 name = "mlx"
-version = "0.31.2.dev20260422+ec49d18e"
-source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#ec49d18ec4cfba0e0c7a37f20d1cf4d75fe56731" }
+version = "0.32.0.dev20260425+cc3f3e60"
+source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }
 resolution-markers = [
     "sys_platform == 'darwin'",
 ]
@@ -1326,7 +1326,7 @@ wheels = [
 [[package]]
 name = "mlx-lm"
 version = "0.31.3"
-source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-arrayscache-leak#c7010341e1f41ac15815feb5dc55134f44e3b044" }
+source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Fdeepseek-v4#0f230c9d108592cea4639a4cac528ce9e1aa0017" }
 resolution-markers = [
     "sys_platform == 'darwin'",
     "sys_platform == 'linux'",
@@ -1334,7 +1334,7 @@ resolution-markers = [
 dependencies = [
     { name = "jinja2", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "mlx", version = "0.31.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
-    { name = "mlx", version = "0.31.2.dev20260422+ec49d18e", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#ec49d18ec4cfba0e0c7a37f20d1cf4d75fe56731" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
+    { name = "mlx", version = "0.32.0.dev20260425+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "numpy", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "protobuf", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "pyyaml", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
@@ -1351,9 +1351,9 @@ dependencies = [
     { name = "fastapi", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "miniaudio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "mlx", version = "0.31.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
-    { name = "mlx", version = "0.31.2.dev20260422+ec49d18e", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#ec49d18ec4cfba0e0c7a37f20d1cf4d75fe56731" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
+    { name = "mlx", version = "0.32.0.dev20260425+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "mlx-lm", version = "0.31.3", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-3-exo-cpu' and extra != 'extra-3-exo-cuda12' and extra != 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
-    { name = "mlx-lm", version = "0.31.3", source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-arrayscache-leak#c7010341e1f41ac15815feb5dc55134f44e3b044" }, marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
+    { name = "mlx-lm", version = "0.31.3", source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Fdeepseek-v4#0f230c9d108592cea4639a4cac528ce9e1aa0017" }, marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "pillow", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
@@ -2634,7 +2634,7 @@ wheels = [
 
 [[package]]
 name = "transformers"
-version = "5.2.0"
+version = "5.6.2"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
@@ -2645,11 +2645,11 @@ dependencies = [
     { name = "safetensors", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "tokenizers", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
     { name = "tqdm", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
-    { name = "typer-slim", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
+    { name = "typer", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/bd/7e/8a0c57d562015e5b16c97c1f0b8e0e92ead2c7c20513225dc12c2043ba9f/transformers-5.2.0.tar.gz", hash = "sha256:0088b8b46ccc9eff1a1dca72b5d618a5ee3b1befc3e418c9512b35dea9f9a650", size = 8618176, upload-time = "2026-02-16T18:54:02.867Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/a4/e9/c6c80a07690142a7d05444271f47b9f3c8aac7dea01d52e1137ee480ad78/transformers-5.6.2.tar.gz", hash = "sha256:e657134c3e5a6bc00a3c35f4e2674bb51adfcd89898495b788a18552bac2b91a", size = 8311867, upload-time = "2026-04-23T18:33:29.332Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/4e/93/79754b0ca486e556c2b95d4f5afc66aaf4b260694f3d6e1b51da2d036691/transformers-5.2.0-py3-none-any.whl", hash = "sha256:9ecaf243dc45bee11a7d93f8caf03746accc0cb069181bbf4ad8566c53e854b4", size = 10403304, upload-time = "2026-02-16T18:53:59.699Z" },
+    { url = "https://files.pythonhosted.org/packages/5d/95/0b0218149b0d6f14df35f5b8f676fa83df4f19ed253c3cc447107ef86eca/transformers-5.6.2-py3-none-any.whl", hash = "sha256:f8d3a1bb96778fed9b8aabfd0dd6e19843e4b0f2bb6b59f32b8a92051b0f348f", size = 10364898, upload-time = "2026-04-23T18:33:26.081Z" },
 ]
 
 [[package]]
@@ -2706,18 +2706,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/4a/91/48db081e7a63bb37284f9fbcefda7c44c277b18b0e13fbc36ea2335b71e6/typer-0.24.1-py3-none-any.whl", hash = "sha256:112c1f0ce578bfb4cab9ffdabc68f031416ebcc216536611ba21f04e9aa84c9e", size = 56085, upload-time = "2026-02-21T16:54:41.616Z" },
 ]
 
-[[package]]
-name = "typer-slim"
-version = "0.24.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "typer", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/a7/a7/e6aecc4b4eb59598829a3b5076a93aff291b4fdaa2ded25efc4e1f4d219c/typer_slim-0.24.0.tar.gz", hash = "sha256:f0ed36127183f52ae6ced2ecb2521789995992c521a46083bfcdbb652d22ad34", size = 4776, upload-time = "2026-02-16T22:08:51.2Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/a7/24/5480c20380dfd18cf33d14784096dca45a24eae6102e91d49a718d3b6855/typer_slim-0.24.0-py3-none-any.whl", hash = "sha256:d5d7ee1ee2834d5020c7c616ed5e0d0f29b9a4b1dd283bdebae198ec09778d0e", size = 3394, upload-time = "2026-02-16T22:08:49.92Z" },
-]
-
 [[package]]
 name = "types-aiofiles"
 version = "25.1.0.20251011"