Temporary reasoning proxy just to test

Update mlx-lm
cleanup
2026-09-08 19:41:32 -04:00 · 2026-04-25 16:11:45 +01:00 · 2026-04-25 16:06:08 +01:00 · 2026-04-25 14:32:21 +01:00 · 2026-04-25 04:45:55 +01:00 · 2026-04-25 02:44:26 +01:00
127 changed files with 4989 additions and 1553 deletions
@@ -32,7 +32,6 @@ jobs:
      SPARKLE_ED25519_PRIVATE: ${{ secrets.SPARKLE_ED25519_PRIVATE }}
      SPARKLE_S3_BUCKET: ${{ secrets.SPARKLE_S3_BUCKET }}
      SPARKLE_S3_PREFIX: ${{ secrets.SPARKLE_S3_PREFIX }}
-      EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT: ${{ secrets.EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT }}
      AWS_REGION: ${{ secrets.AWS_REGION }}
      EXO_BUILD_NUMBER: ${{ github.run_number }}
      EXO_LIBP2P_NAMESPACE: ${{ github.ref_name }}
@@ -347,7 +346,6 @@ jobs:
            EXO_BUILD_COMMIT="$GITHUB_SHA" \
            SPARKLE_FEED_URL="$SPARKLE_FEED_URL" \
            SPARKLE_ED25519_PUBLIC="$SPARKLE_ED25519_PUBLIC" \
-            EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT="$EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT" \
            CODE_SIGNING_IDENTITY="$SIGNING_IDENTITY" \
            CODE_SIGN_INJECT_BASE_ENTITLEMENTS=YES
          mkdir -p ../../output
@@ -1767,12 +1767,12 @@ def clip(
        array: The clipped array.
    """

-def compile(
-    fun: Callable,
+def compile[F: Callable[..., object]](
+    fun: F,
    inputs: object | None = ...,
    outputs: object | None = ...,
    shapeless: bool = ...,
-) -> Callable:
+) -> F:
    """
    Returns a compiled function which produces the same output as ``fun``.

@@ -2915,8 +2915,8 @@ def gather_mm(
    a: array,
    b: array,
    /,
-    lhs_indices: array,
-    rhs_indices: array,
+    lhs_indices: array | None = ...,
+    rhs_indices: array | None = ...,
    *,
    sorted_indices: bool = ...,
    stream: Stream | Device | None = ...,
@@ -4707,6 +4707,7 @@ def softmax(
    /,
    axis: int | Sequence[int] | None = ...,
    *,
+    precise: bool = ...,
    stream: Stream | Device | None = ...,
 ) -> array:
    """
@@ -57,6 +57,10 @@ class Module(dict):
    def __init__(self) -> None:
        """Should be called by the subclasses of ``Module``."""

+    def __getitem__(self, key: str) -> mx.array | Module: ...
+    def get(
+        self, key: str, default: mx.array | Module | None = ...
+    ) -> mx.array | Module | None: ...
    @property
    def training(self):  # -> bool:
        """Boolean indicating if the model is in training mode."""
@@ -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,
@@ -3,7 +3,7 @@ This type stub file was generated by pyright.
 """

 from dataclasses import dataclass
-from typing import Optional
+from typing import Any, Optional

 import mlx.core as mx

@@ -37,10 +37,10 @@ def quantized_scaled_dot_product_attention(
    bits: int = ...,
 ) -> mx.array: ...
 def scaled_dot_product_attention(
-    queries,
-    keys,
-    values,
-    cache,
+    queries: mx.array,
+    keys: mx.array,
+    values: mx.array,
+    cache: Optional[Any],
    scale: float,
    mask: Optional[mx.array],
    sinks: Optional[mx.array] = ...,
@@ -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]: ...
@@ -0,0 +1,103 @@
+"""Type stubs for mlx_lm.models.gpt_oss"""
+
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .cache import KVCache
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    num_local_experts: int
+    num_experts_per_tok: int
+    vocab_size: int
+    rms_norm_eps: float
+    sliding_window: int
+    layer_types: Optional[List[str]]
+
+def mlx_topk(a: mx.array, k: int, axis: int = -1) -> tuple[mx.array, mx.array]: ...
+
+class AttentionBlock(nn.Module):
+    head_dim: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    num_key_value_groups: int
+    sinks: mx.array
+    q_proj: nn.Linear
+    k_proj: nn.Linear
+    v_proj: nn.Linear
+    o_proj: nn.Linear
+    sm_scale: float
+    rope: nn.Module
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class TransformerBlock(nn.Module):
+    self_attn: AttentionBlock
+    mlp: MLPBlock
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class MLPBlock(nn.Module):
+    hidden_size: int
+    num_local_experts: int
+    num_experts_per_tok: int
+    experts: SwitchGLU
+    router: nn.Linear
+    sharding_group: Optional[mx.distributed.Group]
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class GptOssMoeModel(nn.Module):
+    embed_tokens: nn.Embedding
+    norm: nn.RMSNorm
+    layer_types: List[str]
+    layers: list[TransformerBlock]
+    window_size: int
+    swa_idx: int
+    ga_idx: int
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class Model(nn.Module):
+    model_type: str
+    model: GptOssMoeModel
+    lm_head: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+    @property
+    def layers(self) -> list[nn.Module]: ...
+    def make_cache(self) -> list[KVCache]: ...
@@ -0,0 +1,94 @@
+"""Type stubs for mlx_lm.models.minimax"""
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    num_local_experts: int
+    num_experts_per_tok: int
+    max_position_embeddings: int
+
+class MiniMaxAttention(nn.Module):
+    num_heads: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    head_dim: int
+    scale: float
+    q_proj: nn.Linear
+    k_proj: nn.Linear
+    v_proj: nn.Linear
+    o_proj: nn.Linear
+    q_norm: nn.Module
+    k_norm: nn.Module
+    rope: nn.Module
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class MiniMaxSparseMoeBlock(nn.Module):
+    num_experts_per_tok: int
+    gate: nn.Linear
+    switch_mlp: SwitchGLU
+    e_score_correction_bias: mx.array
+    sharding_group: Optional[mx.distributed.Group]
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class MiniMaxDecoderLayer(nn.Module):
+    self_attn: MiniMaxAttention
+    block_sparse_moe: MiniMaxSparseMoeBlock
+    input_layernorm: nn.RMSNorm
+    post_attention_layernorm: nn.RMSNorm
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class MiniMaxModel(nn.Module):
+    embed_tokens: nn.Embedding
+    layers: list[MiniMaxDecoderLayer]
+    norm: nn.RMSNorm
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class Model(nn.Module):
+    model_type: str
+    model: MiniMaxModel
+    lm_head: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+    @property
+    def layers(self) -> list[MiniMaxDecoderLayer]: ...
@@ -92,6 +92,15 @@ class NemotronHAttention(nn.Module):
        cache: Optional[KVCache] = None,
    ) -> mx.array: ...

+class MoEGate(nn.Module):
+    config: ModelArgs
+    top_k: int
+    norm_topk_prob: bool
+    weight: mx.array
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...
+
 class NemotronHMLP(nn.Module):
    up_proj: nn.Linear
    down_proj: nn.Linear
@@ -102,9 +111,14 @@ class NemotronHMLP(nn.Module):
    def __call__(self, x: mx.array) -> mx.array: ...

 class NemotronHMoE(nn.Module):
+    config: ModelArgs
    num_experts_per_tok: int
+    moe_latent_size: Optional[int]
    switch_mlp: SwitchMLP
+    gate: MoEGate
    shared_experts: NemotronHMLP
+    fc1_latent_proj: nn.Linear
+    fc2_latent_proj: nn.Linear

    def __init__(self, config: ModelArgs) -> None: ...
    def __call__(self, x: mx.array) -> mx.array: ...
@@ -71,6 +71,7 @@ class Qwen3NextAttention(nn.Module):
 class Qwen3NextSparseMoeBlock(nn.Module):
    norm_topk_prob: bool
    num_experts: int
+    num_experts_per_tok: int
    top_k: int
    gate: nn.Linear
    switch_mlp: SwitchGLU
@@ -584,9 +584,18 @@ struct ContentView: View {

            case .prompting:
                VStack(alignment: .leading, spacing: 6) {
-                    Text("What's the issue? (optional)")
+                    VStack(alignment: .leading, spacing: 2) {
+                        Text("Tell us what went wrong (optional)")
+                            .font(.caption2)
+                            .foregroundColor(.secondary)
+                        Text(
+                            "A quick description of what you were doing and what happened helps us track down the bug for you."
+                        )
                        .font(.caption2)
                        .foregroundColor(.secondary)
+                        .opacity(0.8)
+                        .fixedSize(horizontal: false, vertical: true)
+                    }
                    TextEditor(text: $bugReportUserDescription)
                        .font(.caption2)
                        .frame(height: 60)
@@ -9,7 +9,7 @@
 	<key>EXOBuildCommit</key>
 	<string>$(EXO_BUILD_COMMIT)</string>
 	<key>EXOBugReportPresignedUrlEndpoint</key>
-	<string>$(EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT)</string>
+	<string>https://reports.exolabs.net/presigned-urls</string>
 	<key>NSLocalNetworkUsageDescription</key>
 	<string>EXO needs local network access to discover and connect to other devices in your cluster for distributed AI inference.</string>
 	<key>NSBonjourServices</key>
@@ -3,11 +3,13 @@ from __future__ import annotations

 import argparse
 import contextlib
+import io
 import json
 import os
 import sys
 import time
 import tomllib
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Literal
@@ -209,7 +211,7 @@ def _openai_build_request(
        "model": model,
        "messages": messages,
        "tools": tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
        "temperature": 0.0,
    }
    return "/v1/chat/completions", body
@@ -276,7 +278,7 @@ def _openai_build_followup(
        "model": model,
        "messages": followup_messages,
        "tools": tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
        "temperature": 0.0,
    }
    return "/v1/chat/completions", body
@@ -379,7 +381,7 @@ def _claude_build_request(
        "model": model,
        "messages": claude_messages,
        "tools": claude_tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
        "temperature": 0.0,
    }
    if system_content is not None:
@@ -489,7 +491,7 @@ def _claude_build_followup(
        "model": model,
        "messages": claude_messages,
        "tools": claude_tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
        "temperature": 0.0,
    }
    if system_content is not None:
@@ -913,6 +915,12 @@ Examples:
        default=1,
        help="Repeat each scenario N times (default: 1)",
    )
+    parser.add_argument(
+        "--concurrency",
+        type=int,
+        default=1,
+        help="Run up to N scenarios in parallel against the same instance (default: 1)",
+    )
    parser.add_argument(
        "--scenarios",
        nargs="*",
@@ -935,6 +943,13 @@ Examples:
    )
    args = parser.parse_args()

+    if args.concurrency < 1:
+        print(
+            f"--concurrency must be >= 1 (got {args.concurrency})",
+            file=sys.stderr,
+        )
+        sys.exit(2)
+
    all_scenarios = load_scenarios(SCENARIOS_PATH)
    if args.scenarios:
        scenarios = [s for s in all_scenarios if s.name in args.scenarios]
@@ -1010,42 +1025,72 @@ Examples:
    cluster_snapshot = capture_cluster_snapshot(exo)
    all_results: list[ScenarioResult] = []

+    tasks: list[tuple[int, Scenario, ApiName]] = [
+        (run_idx, scenario, api_name)
+        for run_idx in range(args.repeat)
+        for scenario in scenarios
+        for api_name in api_names
+    ]
+
+    def _run_one(
+        http_client: httpx.Client,
+        task: tuple[int, Scenario, ApiName],
+    ) -> tuple[tuple[int, Scenario, ApiName], list[ScenarioResult], str]:
+        run_idx, scenario, api_name = task
+        buf = io.StringIO()
+        run_tag = f"[run {run_idx + 1}/{args.repeat}]" if args.repeat > 1 else ""
+        print(
+            f"\n  {run_tag}[{api_name:>9}] {scenario.name}: {scenario.description}",
+            file=buf,
+        )
+        scenario_results = run_scenario(
+            http_client,
+            args.host,
+            args.port,
+            full_model_id,
+            scenario,
+            api_name,
+            args.timeout,
+            args.verbose,
+        )
+        for r in scenario_results:
+            status = "PASS" if r.passed else "FAIL"
+            print(
+                f"    [{r.phase:>10}] {status}  ({r.latency_ms:.0f}ms)",
+                file=buf,
+            )
+            for check_name, check_ok in r.checks.items():
+                mark = "+" if check_ok else "-"
+                print(f"      {mark} {check_name}", file=buf)
+            if r.error:
+                print(f"      ! {r.error}", file=buf)
+        return task, scenario_results, buf.getvalue()
+
    try:
        with httpx.Client() as http_client:
-            for run_idx in range(args.repeat):
-                if args.repeat > 1:
-                    print(f"\n--- Run {run_idx + 1}/{args.repeat} ---", file=log)
-
-                for scenario in scenarios:
-                    for api_name in api_names:
-                        print(
-                            f"\n  [{api_name:>9}] {scenario.name}: {scenario.description}",
-                            file=log,
-                        )
-
-                        scenario_results = run_scenario(
-                            http_client,
-                            args.host,
-                            args.port,
-                            full_model_id,
-                            scenario,
-                            api_name,
-                            args.timeout,
-                            args.verbose,
-                        )
+            if args.concurrency == 1:
+                current_run = -1
+                for task in tasks:
+                    run_idx = task[0]
+                    if args.repeat > 1 and run_idx != current_run:
+                        print(f"\n--- Run {run_idx + 1}/{args.repeat} ---", file=log)
+                        current_run = run_idx
+                    _, scenario_results, buffered = _run_one(http_client, task)
+                    all_results.extend(scenario_results)
+                    log.write(buffered)
+                    log.flush()
+            else:
+                print(
+                    f"Running {len(tasks)} tasks with concurrency={args.concurrency}",
+                    file=log,
+                )
+                with ThreadPoolExecutor(max_workers=args.concurrency) as pool:
+                    futures = [pool.submit(_run_one, http_client, t) for t in tasks]
+                    for fut in as_completed(futures):
+                        _, scenario_results, buffered = fut.result()
                        all_results.extend(scenario_results)
-
-                        for r in scenario_results:
-                            status = "PASS" if r.passed else "FAIL"
-                            print(
-                                f"    [{r.phase:>10}] {status}  ({r.latency_ms:.0f}ms)",
-                                file=log,
-                            )
-                            for check_name, check_ok in r.checks.items():
-                                mark = "+" if check_ok else "-"
-                                print(f"      {mark} {check_name}", file=log)
-                            if r.error:
-                                print(f"      ! {r.error}", file=log)
+                        log.write(buffered)
+                        log.flush()
    finally:
        try:
            exo.request_json("DELETE", f"/instance/{instance_id}")
@@ -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
@@ -564,7 +564,7 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
    ap.add_argument(
        "--settle-timeout",
        type=float,
-        default=0,
+        default=60.0,
        help="Max seconds to wait for the cluster to produce valid placements (0 = try once).",
    )
    ap.add_argument(
@@ -88,10 +88,12 @@
  let codexModel = $state("");
  let codexMcpPath = $state("/Users/username");
  let openClawModel = $state("");
+  let piModel = $state("");
  $effect(() => {
    const def = modelsBySize.length > 0 ? modelsBySize[0] : "your-model-id";
    codexModel = def;
    openClawModel = def;
+    piModel = def;
  });

  const claudeShellCommand = $derived(
@@ -218,6 +220,55 @@
    ),
  );

+  const piModelsJson = $derived.by(() => {
+    const models: Record<string, unknown>[] = [];
+    for (const modelId of runningModels) {
+      const caps = modelCapabilities[modelId] || [];
+      const ctxLen = modelContextLengths[modelId] || 0;
+      const entry: Record<string, unknown> = { id: modelId };
+      if (caps.includes("vision")) {
+        entry.input = ["text", "image"];
+      }
+      // Mark thinking-capable models so pi surfaces its thinking-level selector
+      // for them. exo capability strings: "thinking" (model emits reasoning
+      // content) and "thinking_toggle" (user can turn it on/off).
+      if (caps.includes("thinking") || caps.includes("thinking_toggle")) {
+        entry.reasoning = true;
+      }
+      if (ctxLen > 0) {
+        entry.contextWindow = ctxLen;
+      }
+      models.push(entry);
+    }
+    if (models.length === 0) {
+      models.push({ id: "your-model-id" });
+    }
+    return JSON.stringify(
+      {
+        providers: {
+          exo: {
+            baseUrl: `${apiUrl}/v1`,
+            api: "openai-completions",
+            apiKey: "exo",
+            compat: {
+              supportsDeveloperRole: false,
+              // exo's OpenAI surface takes a boolean `enable_thinking` toggle,
+              // not graded effort levels, so disable pi's `reasoning_effort`
+              // parameter and use the matching top-level-boolean format.
+              supportsReasoningEffort: false,
+              thinkingFormat: "qwen",
+            },
+            models,
+          },
+        },
+      },
+      null,
+      2,
+    );
+  });
+
+  const piShellCommand = $derived(`pi --provider exo --model ${piModel}`);
+
  const ollamaCommand = $derived(
    `OLLAMA_HOST=${apiUrl}/ollama ollama run ${modelsBySize.length > 0 ? modelsBySize[0] : "your-model-id"}`,
  );
@@ -277,6 +328,7 @@
    "OpenCode",
    "Codex",
    "OpenClaw",
+    "Pi",
    "Open WebUI",
    "n8n",
    "Firefox",
@@ -515,6 +567,33 @@
          config={`openclaw doctor --fix${(modelCapabilities[openClawModel] || []).includes("vision") ? `\nopenclaw models set-image exo/${openClawModel}` : ""}\nopenclaw gateway &\nopenclaw dashboard`}
          language="bash"
        />
+      {:else if activeTab === "Pi"}
+        {#if runningModels.length > 1}
+          <div class="text-xs">
+            <span
+              class="text-exo-light-gray/50 text-[10px] uppercase tracking-wider block mb-1"
+              >Model</span
+            >
+            <select bind:value={piModel} class={selectClass}>
+              {#each runningModels as model}
+                <option value={model}>{model.split("/").pop()}</option>
+              {/each}
+            </select>
+          </div>
+        {/if}
+        <IntegrationCard
+          title="Models Config"
+          subtitle="~/.pi/agent/models.json"
+          description="Register exo as a custom provider in pi. Create or edit this file, then run pi and pick an exo model via /model. Install pi with: npm install -g @mariozechner/pi-coding-agent"
+          config={piModelsJson}
+        />
+        <IntegrationCard
+          title="Shell Command"
+          subtitle="Run in terminal"
+          description="Launch pi directly with the exo provider and model selected."
+          config={piShellCommand}
+          language="bash"
+        />
      {:else if activeTab === "Open WebUI"}
        <IntegrationCard
          title="1. Start Open WebUI"
@@ -0,0 +1,6 @@
+{
+  "name": "exo",
+  "lockfileVersion": 3,
+  "requires": true,
+  "packages": {}
+}
@@ -3,7 +3,7 @@ name = "exo"
 version = "0.3.70"
 description = "Exo"
 readme = "README.md"
-requires-python = ">=3.13"
+requires-python = "==3.13.*"
 dependencies = [
  "aiofiles>=24.1.0",
  "aiohttp>=3.12.14",
@@ -17,8 +17,8 @@ dependencies = [
  "loguru>=0.7.3",
  "exo-pyo3-bindings",                       # rust bindings
  "anyio==4.11.0",
-  "mlx==0.31.1; sys_platform == 'darwin'",
-  "mlx-lm",
+  "mlx==0.31.2; sys_platform == 'darwin'",
+  "mlx-lm; sys_platform=='darwin'",
  "tiktoken>=0.12.0",                        # required for kimi k2 tokenizer
  "hypercorn>=0.18.0",
  "openai-harmony>=0.0.8",
@@ -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]
@@ -49,15 +50,21 @@ dev = [
 [project.optional-dependencies]
 build = ["nanobind"]
 cpu = [
+  "mlx==0.31.1; sys_platform == 'linux'",
  "mlx-cpu==0.31.1; sys_platform == 'linux'",
+  "mlx-lm; sys_platform == 'linux'",
  "torch>=2.10.0; sys_platform == 'linux'",
 ]
 cuda12 = [
+  "mlx==0.31.1; sys_platform == 'linux'",
  "mlx-cuda-12==0.31.1; sys_platform == 'linux'",
+  "mlx-lm; sys_platform == 'linux'",
  "torch>=2.10.0; sys_platform == 'linux'",
 ]
 cuda13 = [
+  "mlx==0.31.1; sys_platform == 'linux'",
  "mlx-cuda-13==0.31.1; sys_platform == 'linux'",
+  "mlx-lm; sys_platform == 'linux'",
  "torch>=2.10.0; sys_platform == 'linux'",
 ]

@@ -71,11 +78,11 @@ 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'" },
-  { index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'cpu' and extra != 'cuda12' and extra != 'cuda13'" },
+  { index = "pytorch-cpu", marker = "(extra != 'cuda12' and extra != 'cuda13' and sys_platform == 'linux') or sys_platform == 'darwin'" },
 ]
 vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }

@@ -148,7 +155,11 @@ 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'",
+]

 [tool.uv.extra-build-dependencies]
 miniaudio = ["setuptools", "cffi", "pycparser"]
@@ -5,14 +5,13 @@ let
    workspaceRoot = ../.;
  };

-  mkPythonSet = { pkgs, lib, self' }:
+  mkPythonSet = { pkgs, lib, self', members }:
    let
      inherit (pkgs.stdenv.hostPlatform) isLinux isDarwin isx86_64;
      inherit (pkgs.config) cudaSupport;
      inherit (pkgs) cudaPackages;
      cuda13Support = cudaSupport && cudaPackages.cudaMajorVersion == "13";
      libmlx_source = if cuda13Support then "mlx-cuda-13" else if cudaSupport then "mlx-cuda-12" else "mlx-cpu";
-      uv_extra = if cuda13Support then "cuda13" else if cudaSupport then "cuda12" else "cpu";
      python = pkgs.python313;
      cudaLibs = with cudaPackages; [
        cuda_cudart
@@ -51,7 +50,7 @@ let
          '';
        };
      };
-      buildSystemsOverlay = final: prev: { } //
+      buildSystemsOverlay = final: prev:
        lib.optionalAttrs isDarwin
          {
            mlx = prev.mlx.overrideAttrs (old:
@@ -81,7 +80,7 @@ let
                nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.cmake self'.packages.metal-toolchain ];
                # TODO: non-sdk_26 support
                buildInputs = (old.buildInputs or [ ])
-                  ++ [ gguf-tools pkgs.fmt pkgs.nlohmann_json pkgs.apple-sdk_26 ];
+                ++ [ gguf-tools pkgs.fmt pkgs.nlohmann_json pkgs.apple-sdk_26 ];
                patches = [
                  (pkgs.replaceVars ../nix/darwin-build-fixes.patch {
                    sdkVersion = pkgs.apple-sdk_26.version;
@@ -113,42 +112,42 @@ let
                MACOSX_DEPLOYMENT_TARGET = pkgs.apple-sdk_26.version;
              });
          } // lib.optionalAttrs isLinux {
-        mlx = prev.mlx.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ lib.optionals cudaSupport cudaLibs;
-          autoPatchelfIgnoreMissingDeps = lib.optionals cudaSupport [ "libcuda.so.1" ];
-          postInstall = (old.postInstall or "") + ''
-            cp -r "${final.${libmlx_source}}/${final.python.sitePackages}/mlx" "$out/${final.python.sitePackages}/mlx/"
-          '';
-        });
-      } // lib.optionalAttrs cudaSupport {
-        "${libmlx_source}" = prev."${libmlx_source}".overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ cudaLibs;
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-        nvidia-cufile = prev.nvidia-cufile.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ [ pkgs.rdma-core ];
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-        nvidia-cusolver = prev.nvidia-cusolver.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ cudaLibs;
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-        nvidia-nvshmem-cu13 = prev.nvidia-nvshmem-cu13.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ [ pkgs.rdma-core pkgs.pmix pkgs.libfabric pkgs.ucx pkgs.openmpi ];
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-        nvidia-cusparse = prev.nvidia-cusparse.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ [ cudaLibs ];
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-        torch = prev.torch.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ cudaLibs;
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-      };
+          mlx = prev.mlx.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ lib.optionals cudaSupport cudaLibs;
+            autoPatchelfIgnoreMissingDeps = lib.optionals cudaSupport [ "libcuda.so.1" ];
+            postInstall = ''
+              cp -r "${final.${libmlx_source}}/${final.python.sitePackages}/mlx" "$out/${final.python.sitePackages}/mlx/"
+            '';
+          });
+        } // lib.optionalAttrs cudaSupport {
+          "${libmlx_source}" = prev."${libmlx_source}".overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ cudaLibs;
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          nvidia-cufile = prev.nvidia-cufile.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ [ pkgs.rdma-core ];
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          nvidia-cusolver = prev.nvidia-cusolver.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ cudaLibs;
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          nvidia-nvshmem-cu13 = prev.nvidia-nvshmem-cu13.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ [ pkgs.rdma-core pkgs.pmix pkgs.libfabric pkgs.ucx pkgs.openmpi ];
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          nvidia-cusparse = prev.nvidia-cusparse.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ [ cudaLibs ];
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          torch = prev.torch.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ cudaLibs;
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+        };
      pyprojectOverlay = workspace.mkPyprojectOverlay {
        sourcePreference = "wheel";
-        dependencies = { exo = [ uv_extra ]; exo-bench = [ ]; };
+        dependencies = members;
      };
      editableOverlay = workspace.mkEditablePyprojectOverlay {
        # Use environment variable pointing to editable root directory
@@ -165,8 +164,8 @@ let
          buildSystemsOverlay
        ]
      );
-
-      mkApp = cmd: name: members: pkgs.writeShellApplication {
+      venv = name: (pythonSet.mkVirtualEnv "${name}-env" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" ]; });
+      mkApp = cmd: name: pkgs.writeShellApplication {
        inherit name;
        runtimeEnv = {
          EXO_DASHBOARD_DIR = self'.packages.dashboard;
@@ -174,17 +173,17 @@ let
        };
        runtimeInputs = [
          # mlx and mlx-cuda ship clashing cmake files - we dont need them at runtime anyway
-          ((pythonSet.mkVirtualEnv "${name}-env" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" ]; }))
+          (venv name)
        ]
        ++ lib.optionals isDarwin [ pkgs.macmon ];
        text = "exec " + lib.optionalString cudaSupport "${lib.getExe pkgs.nix-gl-host} " + cmd;
      };
    in
    {
-      inherit pythonSet;
+      inherit venv;
      editablePythonSet = pythonSet.overrideScope editableOverlay;
-      mkPythonScript = members: name: path: mkApp ''python ${path} "$@"'' name members;
-      mkExo = name: members: mkApp ''exo "$@"'' name members;
+      mkPythonScript = path: mkApp ''python ${path} "$@"'';
+      mkExo = mkApp ''exo "$@"'';
    };
 in
 {
@@ -192,16 +191,21 @@ in
    { self', pkgs, unfreePkgs, lib, ... }:
    let
      inherit (pkgs.stdenv.hostPlatform) isLinux;
-      inherit (mkPythonSet { inherit self' pkgs lib; }) pythonSet editablePythonSet mkPythonScript mkExo;
-
-      exoVenv = pythonSet.mkVirtualEnv "exo-env" { exo = lib.optionals isLinux [ "cpu" ]; };
+      inherit (mkPythonSet { inherit self' pkgs lib; members = { exo = [ "cpu" ]; }; }) editablePythonSet mkExo;

      # Virtual environment with dev dependencies for testing
-      testVenv = pythonSet.mkVirtualEnv "exo-test-env" {
-        exo = [ "dev" ] ++ lib.optionals isLinux [ "cpu" ]; # Include pytest, pytest-asyncio, pytest-env
+      testVenv = (mkPythonSet {
+        inherit self' pkgs lib; members = {
+        exo = [ "dev" "cpu" ]; # Include pytest, pytest-asyncio, pytest-env
      };
+      }).venv "exo-test";

-      mkBenchScript = mkPythonScript { exo-bench = [ ]; };
+      mkBenchScript = (mkPythonSet {
+        inherit self' pkgs lib; members = {
+        exo = [ "cpu" ];
+        exo-bench = [ ]; # Include pytest, pytest-asyncio, pytest-env
+      };
+      }).mkPythonScript;

      mkSimplePythonScript = name: path: pkgs.writeShellApplication {
        inherit name;
@@ -212,9 +216,7 @@ in
    in
    {
      packages = {
-        exo = mkExo "exo" { exo = lib.optionals isLinux [ "cpu" ]; };
-        # for devShell
-        exo-venv = exoVenv;
+        exo = mkExo "exo";
        editableVenv = editablePythonSet.mkVirtualEnv "exo-dev-env" { exo = [ "dev" ]; };
        # for running tests in ci
        exo-test-env = testVenv;
@@ -224,8 +226,8 @@ in
        # used by ./tests/run_exo_on.sh
        exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
      } // lib.optionalAttrs isLinux {
-        exo-cuda-12 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; }).mkExo "exo-cuda-12" { exo = [ "cuda12" ]; };
-        exo-cuda-13 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; }).mkExo "exo-cuda-13" { exo = [ "cuda13" ]; };
+        exo-cuda-12 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; members = { exo = [ "cuda12" ]; }; }).mkExo "exo-cuda-12";
+        exo-cuda-13 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; members = { exo = [ "cuda13" ]; }; }).mkExo "exo-cuda-13";
      };

      checks = {
@@ -235,19 +237,11 @@ in
          touch $out
        '';

-        typecheck = pkgs.runCommand "typecheck"
-          {
-            nativeBuildInputs = [
-              testVenv
-              pkgs.basedpyright
-            ];
-          }
-          ''
-            cd ${inputs.self}
-            export HOME=$TMPDIR
-            basedpyright --pythonpath ${testVenv}/bin/python --project ${inputs.self}/pyproject.toml
-            touch $out
-          '';
+        typecheck = pkgs.runCommand "typecheck" { nativeBuildInputs = [ testVenv ]; } ''
+          cd ${inputs.self}
+          basedpyright
+          touch $out
+        '';
      };
    };
 }
@@ -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
@@ -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
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -0,0 +1,21 @@
+model_id = "mlx-community/GLM-5.1-DQ4plus-q8"
+n_layers = 78
+hidden_size = 6144
+num_key_value_heads = 64
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "glm"
+quantization = "8bit"
+base_model = "GLM-5.1"
+capabilities = ["text", "thinking"]
+reasoning_dialect = "post_last_user"
+context_length = 202752
+
+[storage_size]
+in_bytes = 465173655552
+
+# Source: https://huggingface.co/zai-org/GLM-5.1
+# Source: https://docs.z.ai/api-reference/llm/chat-completion
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
@@ -0,0 +1,21 @@
+model_id = "mlx-community/GLM-5.1-MXFP4-Q8"
+n_layers = 78
+hidden_size = 6144
+num_key_value_heads = 64
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "glm"
+quantization = "MXFP4-Q8"
+base_model = "GLM-5.1"
+capabilities = ["text", "thinking"]
+reasoning_dialect = "post_last_user"
+context_length = 202752
+
+[storage_size]
+in_bytes = 405480321024
+
+# Source: https://huggingface.co/zai-org/GLM-5.1
+# Source: https://docs.z.ai/api-reference/llm/chat-completion
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
@@ -0,0 +1,21 @@
+model_id = "mlx-community/GLM-5.1"
+n_layers = 78
+hidden_size = 6144
+num_key_value_heads = 64
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "glm"
+quantization = "bf16"
+base_model = "GLM-5.1"
+capabilities = ["text", "thinking"]
+reasoning_dialect = "post_last_user"
+context_length = 202752
+
+[storage_size]
+in_bytes = 1487822475264
+
+# Source: https://huggingface.co/zai-org/GLM-5.1
+# Source: https://docs.z.ai/api-reference/llm/chat-completion
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
@@ -8,7 +8,7 @@ family = "kimi"
 quantization = ""
 base_model = "Kimi K2"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "suffix"
 context_length = 262144

 [storage_size]
@@ -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]
@@ -0,0 +1,33 @@
+model_id = "mlx-community/Kimi-K2.6-mlx-DQ3_K_M-q8"
+n_layers = 61
+hidden_size = 7168
+num_key_value_heads = 64
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "kimi"
+quantization = "3bit"
+base_model = "Kimi K2.6"
+capabilities = ["text", "thinking", "thinking_toggle", "vision"]
+reasoning_dialect = "suffix"
+context_length = 262144
+
+[storage_size]
+in_bytes = 470628683776
+
+[vision]
+image_token_id = 163605
+model_type = "kimi_vl"
+weights_repo = "exolabs/Kimi-K2.6-vision"
+processor_repo = "moonshotai/Kimi-K2.6"
+
+# Source: https://huggingface.co/moonshotai/Kimi-K2.6
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
+min_p = 0.01
+
+# Source: https://huggingface.co/moonshotai/Kimi-K2.6
+[sampling_defaults.non_thinking]
+temperature = 0.6
+top_p = 0.95
+min_p = 0.01
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -0,0 +1,35 @@
+model_id = "mlx-community/Qwen3.6-27B-4bit"
+n_layers = 64
+hidden_size = 5120
+num_key_value_heads = 4
+supports_tensor = true
+tasks = ["TextGeneration"]
+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]
+in_bytes = 16054262240
+
+# Source: https://huggingface.co/Qwen/Qwen3.6-27B#best-practices
+# Source: https://unsloth.ai/docs/models/qwen3.5
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
+top_k = 20
+min_p = 0.0
+repetition_penalty = 1.0
+presence_penalty = 1.5
+
+# Source: https://huggingface.co/Qwen/Qwen3.6-27B#best-practices
+# Source: https://unsloth.ai/docs/models/qwen3.5
+[sampling_defaults.non_thinking]
+temperature = 0.7
+top_p = 0.8
+top_k = 20
+min_p = 0.0
+repetition_penalty = 1.0
+presence_penalty = 1.5
@@ -0,0 +1,35 @@
+model_id = "mlx-community/Qwen3.6-27B-8bit"
+n_layers = 64
+hidden_size = 5120
+num_key_value_heads = 4
+supports_tensor = true
+tasks = ["TextGeneration"]
+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]
+in_bytes = 29500938720
+
+# Source: https://huggingface.co/Qwen/Qwen3.6-27B#best-practices
+# Source: https://unsloth.ai/docs/models/qwen3.5
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
+top_k = 20
+min_p = 0.0
+repetition_penalty = 1.0
+presence_penalty = 1.5
+
+# Source: https://huggingface.co/Qwen/Qwen3.6-27B#best-practices
+# Source: https://unsloth.ai/docs/models/qwen3.5
+[sampling_defaults.non_thinking]
+temperature = 0.7
+top_p = 0.8
+top_k = 20
+min_p = 0.0
+repetition_penalty = 1.0
+presence_penalty = 1.5
@@ -0,0 +1,35 @@
+model_id = "mlx-community/Qwen3.6-27B-bf16"
+n_layers = 64
+hidden_size = 5120
+num_key_value_heads = 4
+supports_tensor = true
+tasks = ["TextGeneration"]
+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]
+in_bytes = 54713457120
+
+# Source: https://huggingface.co/Qwen/Qwen3.6-27B#best-practices
+# Source: https://unsloth.ai/docs/models/qwen3.5
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
+top_k = 20
+min_p = 0.0
+repetition_penalty = 1.0
+presence_penalty = 1.5
+
+# Source: https://huggingface.co/Qwen/Qwen3.6-27B#best-practices
+# Source: https://unsloth.ai/docs/models/qwen3.5
+[sampling_defaults.non_thinking]
+temperature = 0.7
+top_p = 0.8
+top_k = 20
+min_p = 0.0
+repetition_penalty = 1.0
+presence_penalty = 1.5
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -0,0 +1,33 @@
+model_id = "moonshotai/Kimi-K2.6"
+n_layers = 61
+hidden_size = 7168
+num_key_value_heads = 64
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "kimi"
+quantization = ""
+base_model = "Kimi K2.6"
+capabilities = ["text", "thinking", "thinking_toggle", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 595148192736
+
+[vision]
+image_token_id = 163605
+model_type = "kimi_vl"
+weights_repo = "exolabs/Kimi-K2.6-vision"
+processor_repo = "moonshotai/Kimi-K2.6"
+
+# Source: https://huggingface.co/moonshotai/Kimi-K2.6
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
+min_p = 0.01
+
+# Source: https://huggingface.co/moonshotai/Kimi-K2.6
+[sampling_defaults.non_thinking]
+temperature = 0.6
+top_p = 0.95
+min_p = 0.01
@@ -113,6 +113,23 @@ def _extract_content(content: str | list[ResponseContentPart]) -> str:
    )


+def _append_tool_call(
+    chat_template_messages: list[dict[str, Any]], tool_call: dict[str, Any]
+) -> None:
+    if chat_template_messages:
+        prev = chat_template_messages[-1]
+        if prev.get("role") == "assistant" and isinstance(prev.get("content"), str):
+            existing: list[dict[str, Any]] | None = prev.get("tool_calls")
+            if existing is None:
+                prev["tool_calls"] = [tool_call]
+            else:
+                existing.append(tool_call)
+            return
+    chat_template_messages.append(
+        {"role": "assistant", "content": "", "tool_calls": [tool_call]}
+    )
+
+
 async def responses_request_to_text_generation(
    request: ResponsesRequest,
 ) -> TextGenerationTaskParams:
@@ -182,59 +199,44 @@ async def responses_request_to_text_generation(
                    | McpCallInputItem()
                    | CustomToolCallInputItem()
                ):
-                    chat_template_messages.append(
+                    _append_tool_call(
+                        chat_template_messages,
                        {
-                            "role": "assistant",
-                            "content": "",
-                            "tool_calls": [
-                                {
-                                    "id": item.call_id,
-                                    "type": "function",
-                                    "function": {
-                                        "name": item.name,
-                                        "arguments": item.arguments,
-                                    },
-                                }
-                            ],
-                        }
+                            "id": item.call_id,
+                            "type": "function",
+                            "function": {
+                                "name": item.name,
+                                "arguments": item.arguments,
+                            },
+                        },
                    )
                case (
                    LocalShellCallInputItem()
                    | ShellCallInputItem()
                    | ComputerCallInputItem()
                ):
-                    chat_template_messages.append(
+                    _append_tool_call(
+                        chat_template_messages,
                        {
-                            "role": "assistant",
-                            "content": "",
-                            "tool_calls": [
-                                {
-                                    "id": item.call_id,
-                                    "type": "function",
-                                    "function": {
-                                        "name": item.type,
-                                        "arguments": json.dumps(item.action),
-                                    },
-                                }
-                            ],
-                        }
+                            "id": item.call_id,
+                            "type": "function",
+                            "function": {
+                                "name": item.type,
+                                "arguments": json.dumps(item.action),
+                            },
+                        },
                    )
                case ApplyPatchCallInputItem():
-                    chat_template_messages.append(
+                    _append_tool_call(
+                        chat_template_messages,
                        {
-                            "role": "assistant",
-                            "content": "",
-                            "tool_calls": [
-                                {
-                                    "id": item.call_id,
-                                    "type": "function",
-                                    "function": {
-                                        "name": "apply_patch",
-                                        "arguments": json.dumps({"patch": item.patch}),
-                                    },
-                                }
-                            ],
-                        }
+                            "id": item.call_id,
+                            "type": "function",
+                            "function": {
+                                "name": "apply_patch",
+                                "arguments": json.dumps({"patch": item.patch}),
+                            },
+                        },
                    )
                case (
                    WebSearchCallInputItem()
@@ -254,21 +256,16 @@ async def responses_request_to_text_generation(
                        args = {"prompt": item.prompt}
                    else:
                        args = {"query": item.query}
-                    chat_template_messages.append(
+                    _append_tool_call(
+                        chat_template_messages,
                        {
-                            "role": "assistant",
-                            "content": "",
-                            "tool_calls": [
-                                {
-                                    "id": item.call_id,
-                                    "type": "function",
-                                    "function": {
-                                        "name": item.type,
-                                        "arguments": json.dumps(args),
-                                    },
-                                }
-                            ],
-                        }
+                            "id": item.call_id,
+                            "type": "function",
+                            "function": {
+                                "name": item.type,
+                                "arguments": json.dumps(args),
+                            },
+                        },
                    )
                case (
                    FunctionCallOutputInputItem()
@@ -320,21 +317,16 @@ async def responses_request_to_text_generation(
                            }
                        )
                case McpApprovalRequestInputItem():
-                    chat_template_messages.append(
+                    _append_tool_call(
+                        chat_template_messages,
                        {
-                            "role": "assistant",
-                            "content": "",
-                            "tool_calls": [
-                                {
-                                    "id": item.call_id,
-                                    "type": "function",
-                                    "function": {
-                                        "name": item.name,
-                                        "arguments": item.arguments,
-                                    },
-                                }
-                            ],
-                        }
+                            "id": item.call_id,
+                            "type": "function",
+                            "function": {
+                                "name": item.name,
+                                "arguments": item.arguments,
+                            },
+                        },
                    )
                case McpApprovalResponseInputItem():
                    chat_template_messages.append(
@@ -185,7 +185,10 @@ from exo.shared.types.tasks import (
 from exo.shared.types.tasks import (
    TextGeneration as TextGenerationTask,
 )
-from exo.shared.types.text_generation import Base64Image, TextGenerationTaskParams
+from exo.shared.types.text_generation import (
+    Base64ImageHash,
+    TextGenerationTaskParams,
+)
 from exo.shared.types.worker.downloads import DownloadCompleted
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding
@@ -234,6 +237,7 @@ class API:
        self.node_id: NodeId = node_id
        self.last_completed_election: int = 0
        self.port = port
+        self._sent_image_hashes: set[str] = set()

        self.paused: bool = False
        self.paused_ev: anyio.Event = anyio.Event()
@@ -283,6 +287,7 @@ class API:
        self.event_receiver.close()
        self.event_receiver = event_receiver
        self._tg.start_soon(self._apply_state)
+        self._sent_image_hashes = set()

    def unpause(self, result_clock: int):
        logger.info("Unpausing API")
@@ -737,8 +742,6 @@ class API:
            "TODO: we should send a notification to the user to download the model"
        )

-    _sent_image_hashes: set[str] = set()
-
    async def _send_text_generation_with_images(
        self, task_params: TextGenerationTaskParams
    ) -> TextGeneration:
@@ -750,23 +753,19 @@ class API:
            return command

        hashes = [hashlib.sha256(img.encode("ascii")).hexdigest() for img in images]
+        all_hashes = {idx: Base64ImageHash(h) for idx, h in enumerate(hashes)}
+        task_params = task_params.model_copy(
+            update={"images": [], "image_hashes": all_hashes}
+        )
+        command = TextGeneration(task_params=task_params)

-        cached_hashes: dict[int, str] = {}
        new_images: list[tuple[int, str]] = []
        for idx, (img, h) in enumerate(zip(images, hashes, strict=True)):
-            if h in self._sent_image_hashes:
-                cached_hashes[idx] = h
-            else:
+            if h not in self._sent_image_hashes:
                self._sent_image_hashes.add(h)
                new_images.append((idx, img))

-        wrapped_hashes = {idx: Base64Image(h) for idx, h in cached_hashes.items()}
-
        if not new_images:
-            task_params = task_params.model_copy(
-                update={"images": [], "image_hashes": wrapped_hashes}
-            )
-            command = TextGeneration(task_params=task_params)
            await self._send(command)
            return command

@@ -775,16 +774,6 @@ class API:
            for i in range(0, len(img_data), EXO_MAX_CHUNK_SIZE):
                all_chunks.append((img_idx, img_data[i : i + EXO_MAX_CHUNK_SIZE]))

-        task_params = task_params.model_copy(
-            update={
-                "images": [],
-                "image_hashes": wrapped_hashes,
-                "total_input_chunks": len(all_chunks),
-                "image_count": len(new_images),
-            }
-        )
-        command = TextGeneration(task_params=task_params)
-
        for global_idx, (img_idx, chunk_data) in enumerate(all_chunks):
            await self._send(
                SendInputChunk(
@@ -1674,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
@@ -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):
@@ -88,7 +88,9 @@ class DownloadCoordinator:

        try:
            if progress.status == "complete":
-                found = await to_thread.run_sync(resolve_existing_model, model_id)
+                found = await to_thread.run_sync(
+                    resolve_existing_model, model_id, callback_shard.model_card
+                )
                if found is not None:
                    completed = self._completed_from_path(
                        callback_shard, found, progress.total
@@ -193,7 +195,9 @@ class DownloadCoordinator:
                return

        # Check all model directories for pre-existing complete models
-        found_path = await to_thread.run_sync(resolve_existing_model, model_id)
+        found_path = await to_thread.run_sync(
+            resolve_existing_model, model_id, shard.model_card
+        )
        if found_path is not None:
            logger.info(f"DownloadCoordinator: Model {model_id} found at {found_path}")
            completed = self._completed_from_path(
@@ -220,7 +224,9 @@ class DownloadCoordinator:
        )

        if initial_progress.status == "complete":
-            found = await to_thread.run_sync(resolve_existing_model, model_id)
+            found = await to_thread.run_sync(
+                resolve_existing_model, model_id, shard.model_card
+            )
            if found is not None:
                completed = self._completed_from_path(
                    shard, found, initial_progress.total
@@ -351,7 +357,9 @@ class DownloadCoordinator:

                    if progress.status == "complete":
                        found = await to_thread.run_sync(
-                            resolve_existing_model, model_id
+                            resolve_existing_model,
+                            model_id,
+                            progress.shard.model_card,
                        )
                        if found is not None:
                            status: DownloadProgress = self._completed_from_path(
@@ -380,7 +388,9 @@ class DownloadCoordinator:
                        # (is_model_directory_complete) which validates that all
                        # safetensors weight files are present.
                        found = await to_thread.run_sync(
-                            resolve_existing_model, model_id
+                            resolve_existing_model,
+                            model_id,
+                            progress.shard.model_card,
                        )
                        if found is not None:
                            status = self._completed_from_path(
@@ -421,7 +431,9 @@ class DownloadCoordinator:
                            (DownloadCompleted, DownloadOngoing, DownloadFailed),
                        ):
                            continue
-                        found = await to_thread.run_sync(resolve_existing_model, mid)
+                        found = await to_thread.run_sync(
+                            resolve_existing_model, mid, card
+                        )
                        if found is not None and is_read_only_model_dir(found):
                            path_shard = PipelineShardMetadata(
                                model_card=card,
@@ -35,7 +35,7 @@ from exo.shared.constants import (
    EXO_MODELS_DIRS,
    EXO_MODELS_READ_ONLY_DIRS,
 )
-from exo.shared.models.model_cards import ModelTask
+from exo.shared.models.model_cards import ModelCard, ModelTask
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
 from exo.shared.types.worker.downloads import (
@@ -118,7 +118,9 @@ class InsufficientDiskSpaceError(Exception):
    """Raised when no writable model directory has enough free space."""


-def resolve_existing_model(model_id: ModelId) -> Path | None:
+def resolve_existing_model(
+    model_id: ModelId, card: ModelCard | None = None
+) -> Path | None:
    """Search all model directories for a complete, pre-existing model.

    Checks read-only directories first, then writable directories.
@@ -128,7 +130,7 @@ def resolve_existing_model(model_id: ModelId) -> Path | None:
    normalized = model_id.normalize()
    for search_dir in (*EXO_MODELS_READ_ONLY_DIRS, *EXO_MODELS_DIRS):
        candidate = search_dir / normalized
-        if candidate.is_dir() and is_model_directory_complete(candidate):
+        if candidate.is_dir() and is_model_directory_complete(candidate, card):
            return candidate
    return None

@@ -165,6 +167,29 @@ def select_download_dir(required_bytes: int) -> Path:
    )


+async def select_download_dir_for_shard(
+    model_id: ModelId,
+    filtered_file_list: list[FileListEntry],
+    total_size: int,
+) -> Path:
+    for candidate_dir in EXO_MODELS_DIRS:
+        if not candidate_dir.exists():
+            continue
+        sub = candidate_dir / model_id.normalize()
+        if not await aios.path.isdir(sub):
+            continue
+        existing_bytes = 0
+        for file_entry in filtered_file_list:
+            existing_bytes += await get_downloaded_size(sub / file_entry.path)
+        remaining = max(total_size - existing_bytes, 0)
+        try:
+            if shutil.disk_usage(candidate_dir).free >= remaining:
+                return candidate_dir
+        except OSError:
+            continue
+    return select_download_dir(total_size)
+
+
 async def resolve_model_dir(model_id: ModelId) -> Path:
    """Return the directory for a model's files, creating it if needed.

@@ -279,10 +304,26 @@ def _scan_model_directory(
    return list(entries_by_path.values())


-def is_model_directory_complete(model_dir: Path) -> bool:
-    """Check if a model directory contains all required weight files."""
+def is_model_directory_complete(model_dir: Path, card: ModelCard | None = None) -> bool:
+    """Check if a model directory contains all required weight files.
+    Also checks for sibling weights repo.
+    """
    file_list = _scan_model_directory(model_dir, recursive=True)
-    return file_list is not None and all(f.size is not None for f in file_list)
+    if file_list is None or not all(f.size is not None for f in file_list):
+        return False
+    if (
+        card is not None
+        and card.vision is not None
+        and card.vision.weights_repo != str(card.model_id)
+    ):
+        vision_id = ModelId(card.vision.weights_repo)
+        normalized = vision_id.normalize()
+        for search_dir in (*EXO_MODELS_READ_ONLY_DIRS, *EXO_MODELS_DIRS):
+            candidate = search_dir / normalized
+            if candidate.is_dir() and is_model_directory_complete(candidate):
+                return True
+        return False
+    return True


 async def _build_file_list_from_local_directory(
@@ -834,7 +875,9 @@ async def download_shard(
            else EXO_DEFAULT_MODELS_DIR / model_id.normalize()
        )
    else:
-        models_dir = select_download_dir(total_size)
+        models_dir = await select_download_dir_for_shard(
+            model_id, filtered_file_list, total_size
+        )
        target_dir = models_dir / model_id.normalize()
        await aios.makedirs(target_dir, exist_ok=True)
    file_progress: dict[str, RepoFileDownloadProgress] = {}
@@ -117,40 +117,39 @@ class ResumableShardDownloader(ShardDownloader):
    ) -> Path:
        allow_patterns = ["config.json"] if config_only else None

+        has_vision_sibling = (
+            not config_only
+            and not self.offline
+            and shard.model_card.vision is not None
+            and shard.model_card.vision.weights_repo != str(shard.model_card.model_id)
+        )
+
+        async def main_progress(
+            cb_shard: ShardMetadata, progress: RepoDownloadProgress
+        ) -> None:
+            if has_vision_sibling and progress.status == "complete":
+                return
+            await self.on_progress_wrapper(cb_shard, progress)
+
        target_dir, _ = await download_shard(
            shard,
-            self.on_progress_wrapper,
+            main_progress,
            max_parallel_downloads=self.max_parallel_downloads,
            allow_patterns=allow_patterns,
            skip_internet=self.offline,
        )

-        if (
-            not config_only
-            and not self.offline
-            and shard.model_card.vision
-            and shard.model_card.vision.weights_repo != str(shard.model_card.model_id)
-        ):
-            vision_repo = shard.model_card.vision.weights_repo
-            vision_card = ModelCard(
-                model_id=ModelId(vision_repo),
-                storage_size=Memory.from_bytes(0),
-                n_layers=1,
-                hidden_size=1,
-                supports_tensor=False,
-                tasks=[ModelTask.TextGeneration],
-            )
-            vision_shard = PipelineShardMetadata(
-                model_card=vision_card,
-                device_rank=0,
-                world_size=1,
-                start_layer=0,
-                end_layer=1,
-                n_layers=1,
-            )
+        if has_vision_sibling:
+            vision_shard = self._build_vision_shard(shard)
+
+            async def vision_progress(
+                _cb_shard: ShardMetadata, progress: RepoDownloadProgress
+            ) -> None:
+                await self.on_progress_wrapper(shard, progress)
+
            await download_shard(
                vision_shard,
-                self.on_progress_wrapper,
+                vision_progress,
                max_parallel_downloads=self.max_parallel_downloads,
                allow_patterns=["*.safetensors", "config.json"],
                skip_internet=self.offline,
@@ -158,6 +157,87 @@ class ResumableShardDownloader(ShardDownloader):

        return target_dir

+    async def _status_for_shard(
+        self, shard: ShardMetadata
+    ) -> tuple[Path, RepoDownloadProgress]:
+        async def _noop(
+            _cb_shard: ShardMetadata, _progress: RepoDownloadProgress
+        ) -> None:
+            return
+
+        path, main_progress = await download_shard(
+            shard,
+            _noop,
+            skip_download=True,
+            skip_internet=self.offline,
+        )
+
+        has_vision_sibling = (
+            shard.model_card.vision is not None
+            and shard.model_card.vision.weights_repo != str(shard.model_card.model_id)
+        )
+        if not has_vision_sibling:
+            return path, main_progress
+
+        vision_shard = self._build_vision_shard(shard)
+        _, vision_progress = await download_shard(
+            vision_shard,
+            _noop,
+            skip_download=True,
+            skip_internet=self.offline,
+        )
+        combined = self._combine_progress(shard, main_progress, vision_progress)
+        return path, combined
+
+    @staticmethod
+    def _build_vision_shard(shard: ShardMetadata) -> PipelineShardMetadata:
+        assert shard.model_card.vision is not None
+        vision_card = ModelCard(
+            model_id=ModelId(shard.model_card.vision.weights_repo),
+            storage_size=Memory.from_bytes(0),
+            n_layers=1,
+            hidden_size=1,
+            supports_tensor=False,
+            tasks=[ModelTask.TextGeneration],
+        )
+        return PipelineShardMetadata(
+            model_card=vision_card,
+            device_rank=0,
+            world_size=1,
+            start_layer=0,
+            end_layer=1,
+            n_layers=1,
+        )
+
+    @staticmethod
+    def _combine_progress(
+        shard: ShardMetadata,
+        main: RepoDownloadProgress,
+        vision: RepoDownloadProgress,
+    ) -> RepoDownloadProgress:
+        status_rank = {"not_started": 0, "in_progress": 1, "complete": 2}
+        combined_status = min(
+            (main.status, vision.status), key=lambda s: status_rank[s]
+        )
+        file_progress = dict(main.file_progress)
+        for file_path, fp in vision.file_progress.items():
+            file_progress[f"{vision.repo_id}/{file_path}"] = fp
+        return RepoDownloadProgress(
+            repo_id=main.repo_id,
+            repo_revision=main.repo_revision,
+            shard=shard,
+            completed_files=main.completed_files + vision.completed_files,
+            total_files=main.total_files + vision.total_files,
+            downloaded=main.downloaded + vision.downloaded,
+            downloaded_this_session=main.downloaded_this_session
+            + vision.downloaded_this_session,
+            total=main.total + vision.total,
+            overall_speed=main.overall_speed + vision.overall_speed,
+            overall_eta=max(main.overall_eta, vision.overall_eta),
+            status=combined_status,
+            file_progress=file_progress,
+        )
+
    async def get_shard_download_status(
        self,
    ) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]:
@@ -166,12 +246,7 @@ class ResumableShardDownloader(ShardDownloader):
        ) -> tuple[Path, RepoDownloadProgress]:
            """Helper coroutine that builds the shard for a model and gets its download status."""
            shard = await build_full_shard(model_id)
-            return await download_shard(
-                shard,
-                self.on_progress_wrapper,
-                skip_download=True,
-                skip_internet=self.offline,
-            )
+            return await self._status_for_shard(shard)

        semaphore = asyncio.Semaphore(self.max_parallel_downloads)

@@ -195,10 +270,5 @@ class ResumableShardDownloader(ShardDownloader):
    async def get_shard_download_status_for_shard(
        self, shard: ShardMetadata
    ) -> RepoDownloadProgress:
-        _, progress = await download_shard(
-            shard,
-            self.on_progress_wrapper,
-            skip_download=True,
-            skip_internet=self.offline,
-        )
+        _, progress = await self._status_for_shard(shard)
        return progress
@@ -410,8 +410,6 @@ class Master:
                        continue

                    logger.debug(f"Master indexing event: {str(event)[:100]}")
-                    indexed = IndexedEvent(event=event, idx=len(self._event_log))
-                    self.state = apply(self.state, indexed)

                    event = event.model_copy(
                        update={"_master_time_stamp": datetime.now(tz=timezone.utc)}
@@ -421,6 +419,9 @@ class Master:
                            update={"when": str(datetime.now(tz=timezone.utc))}
                        )

+                    indexed = IndexedEvent(event=event, idx=len(self._event_log))
+                    self.state = apply(self.state, indexed)
+
                    self._event_log.append(event)
                    await self._send_event(indexed)

@@ -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(
@@ -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))
@@ -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)
@@ -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)
@@ -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]
@@ -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()
@@ -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()
@@ -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"
@@ -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"}
+        },
+    )
@@ -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"],
@@ -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(
@@ -114,8 +117,6 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
    frequency_penalty: float | None = None
    images: list[Base64Image] = Field(default_factory=list)
    image_hashes: dict[int, Base64ImageHash] = Field(default_factory=dict)
-    total_input_chunks: int = 0
-    image_count: int = 0

    def with_card_sampling_defaults(self) -> "TextGenerationTaskParams":
        from exo.shared.models.model_cards import get_card
@@ -70,6 +70,11 @@ class FinishedResponse(BaseRunnerResponse):
    pass


+class ModelLoadingResponse(BaseRunnerResponse):
+    layers_loaded: int
+    total: int
+
+
 class PrefillProgressResponse(BaseRunnerResponse):
    processed_tokens: int
    total_tokens: int
@@ -8,6 +8,7 @@ from PIL import Image

 from exo.api.types import AdvancedImageParams
 from exo.download.download_utils import build_model_path
+from exo.shared.types.common import ModelId
 from exo.shared.types.worker.instances import BoundInstance
 from exo.shared.types.worker.shards import CfgShardMetadata, PipelineShardMetadata
 from exo.worker.engines.image.config import ImageModelConfig
@@ -22,13 +23,14 @@ from exo.worker.runner.bootstrap import logger


 class DistributedImageModel:
+    model_id: ModelId
    _config: ImageModelConfig
    _adapter: ModelAdapter[Any, Any]
    _runner: DiffusionRunner

    def __init__(
        self,
-        model_id: str,
+        model_id: ModelId,
        local_path: Path,
        shard_metadata: PipelineShardMetadata | CfgShardMetadata,
        group: Optional[mx.distributed.Group] = None,
@@ -68,6 +70,7 @@ class DistributedImageModel:
        else:
            logger.info("Single-node initialization")

+        self.model_id = model_id
        self._config = config
        self._adapter = adapter
        self._runner = runner
@@ -3,7 +3,7 @@ import io
 import random
 import tempfile
 import time
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
 from pathlib import Path
 from typing import Generator, Literal

@@ -17,11 +17,10 @@ from exo.api.types import (
    ImageGenerationTaskParams,
    ImageSize,
 )
+from exo.shared.constants import EXO_MAX_CHUNK_SIZE
+from exo.shared.types.chunks import ImageChunk
+from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
-from exo.shared.types.worker.runner_response import (
-    ImageGenerationResponse,
-    PartialImageResponse,
-)
 from exo.worker.engines.image.distributed_model import DistributedImageModel


@@ -71,16 +70,8 @@ def generate_image(
    model: DistributedImageModel,
    task: ImageGenerationTaskParams | ImageEditsTaskParams,
    cancel_checker: Callable[[], bool] | None = None,
-) -> Generator[ImageGenerationResponse | PartialImageResponse, None, None]:
-    """Generate image(s), optionally yielding partial results.
-
-    When partial_images > 0 or stream=True, yields PartialImageResponse for
-    intermediate images, then ImageGenerationResponse for the final image.
-
-    Yields:
-        PartialImageResponse for intermediate images (if partial_images > 0, first image only)
-        ImageGenerationResponse for final complete images
-    """
+) -> Generator[ImageChunk, None, None]:
+    """Generate image(s), optionally yielding partial results."""
    width, height = parse_size(task.size)
    quality: Literal["low", "medium", "high"] = task.quality or "medium"

@@ -142,12 +133,14 @@ def generate_image(
                        image = image.convert("RGB")
                    image.save(buffer, format=image_format)

-                    yield PartialImageResponse(
+                    yield from _process_image_response(
                        image_data=buffer.getvalue(),
-                        format=task.output_format,
+                        image_format=task.output_format,
                        partial_index=partial_idx,
                        total_partials=total_partials,
                        image_index=image_num,
+                        model_id=model.model_id,
+                        stats=None,
                    )
                else:
                    image = result
@@ -189,9 +182,54 @@ def generate_image(
                        image = image.convert("RGB")
                    image.save(buffer, format=image_format)

-                    yield ImageGenerationResponse(
+                    yield from _process_image_response(
                        image_data=buffer.getvalue(),
-                        format=task.output_format,
+                        image_format=task.output_format,
                        stats=stats,
                        image_index=image_num,
+                        model_id=model.model_id,
+                        partial_index=None,
+                        total_partials=None,
                    )
+
+
+def _process_image_response(
+    image_data: bytes,
+    image_index: int,
+    image_format: Literal["png", "jpeg", "webp"],
+    partial_index: int | None,
+    total_partials: int | None,
+    stats: ImageGenerationStats | None,
+    model_id: ModelId,
+) -> Iterator[ImageChunk]:
+    """Process a single image response and send chunks."""
+    is_partial = partial_index is not None
+    encoded_data = base64.b64encode(image_data).decode("utf-8")
+    # Extract stats from final ImageGenerationResponse if available
+    data_chunks = [
+        encoded_data[i : i + EXO_MAX_CHUNK_SIZE]
+        for i in range(0, len(encoded_data), EXO_MAX_CHUNK_SIZE)
+    ]
+    total_chunks = len(data_chunks)
+
+    def _data_to_chunk(item: tuple[int, str]) -> ImageChunk:
+        chunk_index, chunk_data = item
+        # Only include stats on the last chunk of the final image
+        chunk_stats = (
+            stats if chunk_index == total_chunks - 1 and not is_partial else None
+        )
+
+        return ImageChunk(
+            model=model_id,
+            data=chunk_data,
+            chunk_index=chunk_index,
+            total_chunks=total_chunks,
+            image_index=image_index,
+            is_partial=is_partial,
+            partial_index=partial_index,
+            total_partials=total_partials,
+            stats=chunk_stats,
+            format=image_format,
+        )
+
+    return map(_data_to_chunk, enumerate(data_chunks))
@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from collections.abc import Callable
+from collections.abc import Callable, Generator
 from functools import partial
 from inspect import signature
 from typing import TYPE_CHECKING, Literal, Protocol, cast
@@ -12,11 +12,13 @@ from mlx.nn.layers.distributed import (
    sum_gradients,
 )
 from mlx_lm.models.base import (
-    scaled_dot_product_attention,  # pyright: ignore[reportUnknownVariableType]
+    scaled_dot_product_attention,
 )
 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
@@ -59,14 +61,13 @@ from mlx_lm.models.step3p5 import Model as Step35Model
 from mlx_lm.models.step3p5 import Step3p5MLP as Step35MLP
 from mlx_lm.models.step3p5 import Step3p5Model as Step35InnerModel

+from exo.shared.types.worker.runner_response import ModelLoadingResponse
 from exo.shared.types.worker.shards import PipelineShardMetadata
 from exo.worker.runner.bootstrap import logger

 if TYPE_CHECKING:
    from mlx_lm.models.cache import Cache

-LayerLoadedCallback = Callable[[int, int], None]  # (layers_loaded, total_layers)
-

 _pending_prefill_sends: list[tuple[mx.array, int, mx.distributed.Group]] = []

@@ -276,8 +277,7 @@ def pipeline_auto_parallel(
    model: nn.Module,
    group: mx.distributed.Group,
    model_shard_meta: PipelineShardMetadata,
-    on_layer_loaded: LayerLoadedCallback | None,
-) -> nn.Module:
+) -> Generator[ModelLoadingResponse, None, nn.Module]:
    """
    Automatically parallelize a model across multiple devices.
    Args:
@@ -297,8 +297,8 @@ def pipeline_auto_parallel(
    total = len(layers)
    for i, layer in enumerate(layers):
        mx.eval(layer)  # type: ignore
-        if on_layer_loaded is not None:
-            on_layer_loaded(i, total)
+        mx.clear_cache()
+        yield ModelLoadingResponse(layers_loaded=i, total=total)

    layers[0] = PipelineFirstLayer(layers[0], device_rank, group=group)
    layers[-1] = PipelineLastLayer(
@@ -309,24 +309,20 @@ def pipeline_auto_parallel(
    )

    if isinstance(inner_model_instance, GptOssMoeModel):
-        inner_model_instance.layer_types = inner_model_instance.layer_types[  # type: ignore
+        inner_model_instance.layer_types = inner_model_instance.layer_types[
            start_layer:end_layer
        ]
        # We can assume the model has at least one layer thanks to placement.
        # If a layer type doesn't exist, we can set it to 0.
        inner_model_instance.swa_idx = (
            0
-            if "sliding_attention" not in inner_model_instance.layer_types  # type: ignore
-            else inner_model_instance.layer_types.index(  # type: ignore
-                "sliding_attention"
-            )
+            if "sliding_attention" not in inner_model_instance.layer_types
+            else inner_model_instance.layer_types.index("sliding_attention")
        )
        inner_model_instance.ga_idx = (
            0
-            if "full_attention" not in inner_model_instance.layer_types  # type: ignore
-            else inner_model_instance.layer_types.index(  # type: ignore
-                "full_attention"
-            )
+            if "full_attention" not in inner_model_instance.layer_types
+            else inner_model_instance.layer_types.index("full_attention")
        )

    if isinstance(inner_model_instance, Step35InnerModel):
@@ -460,8 +456,7 @@ def patch_tensor_model[T](model: T) -> T:
 def tensor_auto_parallel(
    model: nn.Module,
    group: mx.distributed.Group,
-    on_layer_loaded: LayerLoadedCallback | None,
-) -> nn.Module:
+) -> Generator[ModelLoadingResponse, None, nn.Module]:
    all_to_sharded_linear = partial(
        shard_linear,
        sharding="all-to-sharded",
@@ -518,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,
@@ -595,7 +598,7 @@ def tensor_auto_parallel(
    else:
        raise ValueError(f"Unsupported model type: {type(model)}")

-    model = tensor_parallel_sharding_strategy.shard_model(model, on_layer_loaded)
+    model = yield from tensor_parallel_sharding_strategy.shard_model(model)
    return patch_tensor_model(model)


@@ -619,16 +622,14 @@ class TensorParallelShardingStrategy(ABC):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module: ...
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]: ...


 class LlamaShardingStrategy(TensorParallelShardingStrategy):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module:
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
        model = cast(LlamaModel, model)
        total = len(model.layers)
        for i, layer in enumerate(model.layers):
@@ -646,8 +647,8 @@ class LlamaShardingStrategy(TensorParallelShardingStrategy):
            layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
            layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
            mx.eval(layer)
-            if on_layer_loaded is not None:
-                on_layer_loaded(i, total)
+
+            yield ModelLoadingResponse(layers_loaded=i, total=total)
        return model


@@ -658,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__}"
@@ -681,8 +689,7 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module:
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
        model = cast(DeepseekV3Model, model)
        total = len(model.layers)

@@ -738,8 +745,8 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
                layer.mlp.sharding_group = self.group

            mx.eval(layer)
-            if on_layer_loaded is not None:
-                on_layer_loaded(i, total)
+
+            yield ModelLoadingResponse(layers_loaded=i, total=total)

        return model

@@ -760,12 +767,183 @@ 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,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module:
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
        model = cast(GLM4MoeLiteModel, model)
        total = len(model.layers)  # type: ignore
        for i, layer in enumerate(model.layers):  # type: ignore
@@ -816,8 +994,9 @@ class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
                layer.mlp = ShardedMoE(layer.mlp)  # type: ignore
                layer.mlp.sharding_group = self.group  # type: ignore
            mx.eval(layer)
-            if on_layer_loaded is not None:
-                on_layer_loaded(i, total)
+            mx.clear_cache()
+
+            yield ModelLoadingResponse(layers_loaded=i, total=total)

        return model

@@ -891,7 +1070,7 @@ class WrappedMiniMaxAttention(CustomMlxLayer):
            keys,
            values,
            cache=cache,
-            scale=self._original_layer.scale,  # type: ignore
+            scale=self._original_layer.scale,
            mask=mask,
        )

@@ -904,8 +1083,7 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module:
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
        model = cast(MiniMaxModel, model)
        total = len(model.layers)
        for i, layer in enumerate(model.layers):
@@ -931,11 +1109,12 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
            self.all_to_sharded_linear_in_place(
                layer.block_sparse_moe.switch_mlp.up_proj
            )
-            layer.block_sparse_moe = ShardedMoE(layer.block_sparse_moe)  # pyright: ignore[reportAttributeAccessIssue, reportArgumentType]
-            layer.block_sparse_moe.sharding_group = self.group  # pyright: ignore[reportAttributeAccessIssue]
+            layer.block_sparse_moe = ShardedMoE(layer.block_sparse_moe)  # type: ignore
+            layer.block_sparse_moe.sharding_group = self.group
            mx.eval(layer)
-            if on_layer_loaded is not None:
-                on_layer_loaded(i, total)
+            mx.clear_cache()
+
+            yield ModelLoadingResponse(layers_loaded=i, total=total)
        return model


@@ -943,8 +1122,7 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module:
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
        model = cast(
            Qwen3Model
            | Qwen3MoeModel
@@ -1099,8 +1277,9 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)

            mx.eval(layer)
-            if on_layer_loaded is not None:
-                on_layer_loaded(i, total)
+            mx.clear_cache()
+
+            yield ModelLoadingResponse(layers_loaded=i, total=total)
        return model


@@ -1108,8 +1287,7 @@ class Glm4MoeShardingStrategy(TensorParallelShardingStrategy):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module:
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
        model = cast(Glm4MoeModel, model)
        total = len(model.layers)
        for i, layer in enumerate(model.layers):
@@ -1145,8 +1323,9 @@ class Glm4MoeShardingStrategy(TensorParallelShardingStrategy):
                layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)

            mx.eval(layer)
-            if on_layer_loaded is not None:
-                on_layer_loaded(i, total)
+            mx.clear_cache()
+
+            yield ModelLoadingResponse(layers_loaded=i, total=total)
        return model


@@ -1154,8 +1333,7 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module:
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
        model = cast(GptOssMoeModel, model)
        total = len(model.layers)

@@ -1184,10 +1362,11 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
            self.all_to_sharded_linear_in_place(layer.mlp.experts.up_proj)

            layer.mlp = ShardedMoE(layer.mlp)  # type: ignore
-            layer.mlp.sharding_group = self.group  # pyright: ignore[reportAttributeAccessIssue]
+            layer.mlp.sharding_group = self.group
            mx.eval(layer)
-            if on_layer_loaded is not None:
-                on_layer_loaded(i, total)
+            mx.clear_cache()
+
+            yield ModelLoadingResponse(layers_loaded=i, total=total)
        return model


@@ -1195,8 +1374,7 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module:
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
        model = cast(Step35Model, model)
        total = len(model.layers)

@@ -1229,8 +1407,9 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy):
                self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)

            mx.eval(layer)
-            if on_layer_loaded is not None:
-                on_layer_loaded(i, total)
+            mx.clear_cache()
+
+            yield ModelLoadingResponse(layers_loaded=i, total=total)
        return model


@@ -1238,8 +1417,7 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module:
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
        model = cast(NemotronHModel, model)
        rank = self.group.rank()
        total = len(model.layers)
@@ -1272,8 +1450,8 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
                layer.mixer = mixer  # pyright: ignore[reportAttributeAccessIssue]

            mx.eval(layer)
-            if on_layer_loaded is not None:
-                on_layer_loaded(i, total)
+            mx.clear_cache()
+            yield ModelLoadingResponse(layers_loaded=i, total=total)
        return model

    def _shard_mamba2_mixer(self, mixer: NemotronHMamba2Mixer, rank: int) -> None:
@@ -1380,8 +1558,7 @@ class Gemma4ShardingStrategy(TensorParallelShardingStrategy):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module:
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
        model = cast(Gemma4Model, model)
        layers = model.language_model.model.layers
        total = len(layers)
@@ -1390,9 +1567,11 @@ class Gemma4ShardingStrategy(TensorParallelShardingStrategy):

            attn = layer.self_attn
            attn.q_proj = self.all_to_sharded_linear(attn.q_proj)
-            attn.k_proj = self.all_to_sharded_linear(attn.k_proj)
-            if not attn.use_k_eq_v:
-                attn.v_proj = self.all_to_sharded_linear(attn.v_proj)
+            has_kv: bool = cast(bool, attn.has_kv)
+            if has_kv:
+                attn.k_proj = self.all_to_sharded_linear(attn.k_proj)
+                if not attn.use_k_eq_v:
+                    attn.v_proj = self.all_to_sharded_linear(attn.v_proj)
            attn.o_proj = self.sharded_to_all_linear(attn.o_proj)
            attn.n_heads //= self.N
            attn.n_kv_heads //= self.N
@@ -1409,6 +1588,6 @@ class Gemma4ShardingStrategy(TensorParallelShardingStrategy):
                layer.experts.sharding_group = self.group

            mx.eval(layer)
-            if on_layer_loaded is not None:
-                on_layer_loaded(i, total)
+            mx.clear_cache()
+            yield ModelLoadingResponse(layers_loaded=i, total=total)
        return model
@@ -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:
@@ -249,7 +298,12 @@ class KVPrefixCache:
        # For partial match: trim to best_length, remaining has suffix to prefill
        # This ensures stream_generate always has at least one token to start with
        has_ssm = has_non_kv_caches(self.caches[best_index])
-        target = (max_length - 1) if is_exact and not has_ssm else best_length
+        cached_length = cache_length(self.caches[best_index])
+        if has_ssm:
+            target = best_length
+        else:
+            desired = (max_length - 1) if is_exact else best_length
+            target = min(cached_length, desired)
        restore_pos, restore_snap = self._get_snapshot(best_index, target)

        # No usable snapshot — need fresh cache
@@ -257,7 +311,6 @@ class KVPrefixCache:
            return make_kv_cache(model), prompt_tokens, None, False

        prompt_cache = deepcopy(self.caches[best_index])
-        cached_length = cache_length(self.caches[best_index])
        tokens_to_trim = cached_length - restore_pos
        if tokens_to_trim > 0:
            trim_cache(prompt_cache, tokens_to_trim, restore_snap)
@@ -353,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)
Author	SHA1	Message	Date
Ryuichi Leo Takashige	8bc74d2cea	Temporary reasoning proxy just to test	2026-04-25 16:11:45 +01:00
Ryuichi Leo Takashige	5ece607264	Update mlx-lm	2026-04-25 16:06:08 +01:00
Ryuichi Leo Takashige	8106d8a7e3	cleanup	2026-04-25 14:32:21 +01:00
Ryuichi Leo Takashige	2ef3a4c707	maybe optimize	2026-04-25 04:45:55 +01:00
Ryuichi Leo Takashige	bba012f15b	Update	2026-04-25 02:44:26 +01:00
Ryuichi Leo Takashige	3babf9d070	Set up dsml encodings	2026-04-24 14:19:38 +01:00
Ryuichi Leo Takashige	5d7ea4c6c0	Add temporary mlx lm fork	2026-04-24 14:07:43 +01:00
Ryuichi Leo Takashige	e116097f64	Update transformers to v5.6	2026-04-24 13:03:55 +01:00
Ryuichi Leo Takashige	c6467094b1	Add DSV4 model cards	2026-04-24 12:16:57 +01:00
rltakashige	fd707de30b	Add more model cards (#1970 )	2026-04-23 15:28:40 +01:00
Alex CheemaandClaude Opus 4.7	45248c5c85	chore(app): hardcode bug report presigned-URL endpoint (#1971 ) ## Motivation The bug-report presigned-URL endpoint (`https://reports.exolabs.net/presigned-urls`) was injected at build time from the `EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT` GitHub Actions secret into `Info.plist`, then read at runtime by `BugReportService`. It isn't actually a secret — the POST body is just `{"keys":[...]}` with no credential (see `app/EXO/EXO/Services/BugReportService.swift:136-142`), abuse prevention lives server-side on the lambda, and the URL is already visible in every publicly-distributed DMG's `Info.plist`. Treating it as a repo secret added plumbing with no security benefit and broke local dev builds — hitting Send Bug Report on an uncustomised `just build-app` raised "Bug report endpoint is invalid". ## Changes - `app/EXO/EXO/Info.plist`: replace `$(EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT)` with the literal URL. - `.github/workflows/build-app.yml`: drop the `EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT` job-level env var and the xcodebuild build-setting passthrough. No other workflow changes. Swift code is unchanged — `BugReportService` still reads from `Info.plist`, which leaves an escape hatch if anyone ever needs to override via `xcodebuild EXOBugReportPresignedUrlEndpoint=...` without recompiling. Follow-up: the `EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT` repo secret can now be deleted in the GitHub Actions settings UI. ## Why It Works `Info.plist` variable substitution turns `$(FOO)` into whatever build setting `FOO` resolves to. CI was setting `FOO` via xcodebuild; local dev wasn't, so the key resolved to an empty string, which `BugReportService.fetchPresignedUploadUrls` rejects via the `!trimmedEndpointString.isEmpty` guard at `BugReportService.swift:131`. Hardcoding the literal string removes the substitution entirely, so every build — local or CI — gets the right value. ## Test Plan ### Manual Testing <!-- Hardware: MacBook Pro (macOS app build via Xcode) --> - `just build-app` with no extra env vars (reproduces the failure path on `main`). - `/usr/libexec/PlistBuddy -c "Print :EXOBugReportPresignedUrlEndpoint" app/EXO/build/Build/Products/Debug/EXO.app/Contents/Info.plist` → returns `https://reports.exolabs.net/presigned-urls` (was empty before this change). - `open app/EXO/build/Build/Products/Debug/EXO.app` → menubar → Debug Info → Send Bug Report → type a description → Send → upload succeeds and the Create GitHub Issue button appears (was failing with "Bug report endpoint is invalid" before). - Cross-check on the Slack side that the uploaded `report.json` lands under `reports/YYYY/MM/DD/<ts>/` as before. ### Automated Testing <!-- Describe changes to automated tests, or how existing tests cover this change --> - No new tests. This is a single-string change to `Info.plist` plus a workflow cleanup. `nix flake check` in CI verifies formatting/lint for the rest of the tree. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>	2026-04-23 14:08:16 +00:00
rltakashige	290e3fd927	Keep image cache fresh (temporary fix) (#1961 ) ## Motivation When a new node joins, it might not have the cache. Caveat: This is potentially fallible if a new node joins and updates real topology, but the API topology hasn't caught up with this fact and the user queues up a new text generation. In practice, there is only a split second where this is the case, and this is only for users of the dashboard interface. We should fix this properly after the release.	2026-04-23 11:39:36 +01:00
rltakashige	3894cf134e	Fix Gemma 4 E2B TP + DeepSeek V32 thinking parsing (#1967 )	2026-04-23 01:50:39 +00:00
Alex CheemaandClaude Opus 4.7	8993ccaf09	feat(app): add friendly context message to bug report prompt (#1959 ) ## Motivation When a user clicks Send Bug Report in the macOS app, we already give them the option to add more context via an optional text field. But the current prompt is just a terse label — `"What's the issue? (optional)"` — which doesn't tell the user why bothering to fill it in matters. A friendly one-line explanation increases the chance they'll describe what went wrong, which is the single most useful signal when we triage the resulting diagnostic bundle. ## Changes - `app/EXO/EXO/ContentView.swift`: In the `.prompting` phase of `sendBugReportButton`, replace the single label with a two-line hierarchy: - Primary: `Tell us what went wrong (optional)` - Helper: `A quick description of what you were doing and what happened helps us track down the bug for you.` - The helper uses `.caption2` + `.secondary` + `.opacity(0.8)` + `.fixedSize(horizontal: false, vertical: true)` so it stays visually subordinate and wraps cleanly inside the 340pt popover. No changes to `BugReportService`, the `user_description` payload, or any other flow. ## Why It Works The optional description is already plumbed end-to-end (text editor → `bugReportUserDescription` state → `BugReportService.sendReport(..., userDescription:)` → `report.json`'s `user_description` field → GitHub issue pre-fill). The only gap was user-facing motivation, so this is purely a copy/layout tweak inside the existing `.prompting` case — no new state, bindings, or service changes. ## Test Plan ### Manual Testing <!-- Hardware: MacBook Pro (macOS app build via Xcode) --> - Build the macOS app in Xcode (`app/EXO/EXO.xcodeproj`) and launch it. - Open the menubar popover → expand Debug Info → click Send Bug Report. - Verify the new primary label and helper sentence both appear above the text editor and wrap cleanly within the popover width. - Leave the field empty → click Send → upload should succeed (no `user_description` in payload, same as before). - Fill in a description → click Send → upload succeeds and the success card with Create GitHub Issue appears; clicking it opens GitHub with the description pre-filled. - Click Cancel from the prompting state → returns to idle. ### Automated Testing <!-- Describe changes to automated tests, or how existing tests cover this change --> - No new automated tests. This is a SwiftUI copy/layout change; existing `EXOTests` are smoke-level and don't cover `ContentView` view bodies, and UI snapshot tests aren't worth adding for a two-line copy tweak. - `nix fmt` reports 0 files changed after the edit; `nix flake check` in CI will verify formatting/lint for the rest of the tree. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>	2026-04-22 17:39:36 +00:00
Nadeem Hilal Wani	4939fbe995	feat(dashboard): add Pi integration tab (#1925 ) ## Summary Adds a new Pi tab to the Integrations page (`/#/integrations`) alongside the existing Claude Code, OpenCode, Codex, OpenClaw, Open WebUI, n8n, and Firefox tabs. [pi](https://pi.dev) (`@mariozechner/pi-coding-agent`) is a terminal coding agent that supports custom OpenAI-compatible providers via `~/.pi/agent/models.json`. This tab gives users a copy-pasteable config to wire pi up to their exo cluster. ## What's in the tab - Model selector (shown when multiple models are running) — picks the default model for the generated shell command. - Models Config card — generates `~/.pi/agent/models.json` registering `exo` as a custom provider: - `baseUrl` → `<apiUrl>/v1` - `api` → `openai-completions` - `apiKey` → `"exo"` (placeholder; exo ignores it) - `compat.supportsDeveloperRole: false` and `compat.supportsReasoningEffort: false`, per pi docs recommendation for local OpenAI-compatible servers - Auto-populates every running model with `id`, `contextWindow` (from `/v1/models`), and `input: ["text", "image"]` for vision-capable models - Shell Command card — `pi --provider exo --model <model>` for quick launch. The tab gracefully falls back to `your-model-id` when no models are running, matching the behavior of the other tabs. ## Usage 1. `npm install -g @mariozechner/pi-coding-agent` 2. Paste the generated config into `~/.pi/agent/models.json` 3. Run `pi` and pick an exo model via `/model` — or run the shell command directly ## Changes - `dashboard/src/routes/integrations/+page.svelte` — adds `"Pi"` to the `tabs` tuple, `piModel` state, `piModelsJson` + `piShellCommand` derivations, and the tab content block. Single-file, scoped change — no backend or type changes. ## Testing - `cd dashboard && npm run build` — ✅ builds cleanly - `svelte-check` on the edited file — no new errors - Manually verified the tab renders, the model selector updates the generated JSON, and the config reflects `/v1/models` capabilities (vision → `input: ["text","image"]`, `context_length` → `contextWindow`). ## Screenshots <img width="1545" height="1236" alt="pi-tab" src="https://github.com/user-attachments/assets/38aa179f-4ed9-4a1e-9783-d3baa7738263" />	2026-04-22 17:29:47 +00:00
rltakashige	73782ecc65	Fix event mutation causing indexed vs event mismatch (#1964 ) Fixes small issue with #1957	2026-04-22 16:12:24 +00:00
rltakashige	f6e418ed23	Cleanup on #1952 (#1960 )	2026-04-22 17:05:49 +01:00
rltakashigeandEvan	7a312a177b	Misc fixes: upstream JACCL all_sum, API, etc. + Add Kimi K2.6 (#1952 ) ## Motivation This fixes a bunch of observed model quality issues introduced upstream in JACCL, as well as API issues and prefix cache calculation. ## Test Plan ### Manual Testing Tested a bunch ### Automated Testing Added a test, automated eval tool calls on Kimi K2.6, Minimax M2.7, GPT OSS and Qwen3.6 models. --------- Co-authored-by: Evan <evanev7@gmail.com>	2026-04-22 15:43:27 +00:00
Evan Quiney	0a549f8846	remove layer loading callback (#1890 ) first part of modularising the backend is simplifying some of the control flow. more tbd.	2026-04-22 14:03:31 +01:00