fix: validate num_key_value_heads in tensor sharding placement (#1669)

## Problem Models with fewer KV heads than nodes crash during tensor parallelism. For example, Qwen3.5 MoE models have only 2 KV heads — trying to shard across 4 nodes produces empty tensors and a reshape error at runtime. The placement system already validates `hidden_size % num_nodes == 0` but doesn't check KV heads, so it creates configurations that look valid but blow up when the worker tries to split the attention heads. Affected models include Qwen3.5-35B-A3B, Qwen3.5-122B-A10B, Qwen3.5-397B-A17B, Qwen3-Next-80B-A3B, and Qwen3-Coder-Next (all have 2 KV heads). ## Changes **Placement validation** (`src/exo/master/placement.py`): - Combined KV heads divisibility check with the existing hidden_size filter in a single pass - Cycles where `num_key_value_heads % len(cycle) != 0` are now excluded for tensor sharding - Error message includes both constraints when no valid cycle is found **Model card schema** (`src/exo/shared/models/model_cards.py`): - Added optional `num_key_value_heads` field to `ModelCard` and `ConfigData` - Extracted from HuggingFace `config.json` (handles both top-level and `text_config` nesting) - Passed through in `fetch_from_hf()` for dynamically fetched cards **All 68 inference model cards** (`resources/inference_model_cards/*.toml`): - Populated `num_key_value_heads` from each model's HuggingFace config **Utility script** (`scripts/fetch_kv_heads.py`): - Fetches `num_key_value_heads` from HuggingFace and updates TOML cards - `--missing`: only fills in cards that don't have the field yet - `--all`: re-fetches and overwrites everything - Uses tomlkit for safe TOML editing and ThreadPoolExecutor for parallel fetches ## Behavior - Instance previews no longer show tensor options for models that can't split their KV heads across the cluster size - `place_instance()` rejects with a clear error instead of crash-looping - Pipeline parallelism is unaffected - 2-node tensor still works for 2-KV-head models (2 ÷ 2 = 1) - Field is optional — existing custom cards without it continue to work (validation is skipped when `None`)
2026-03-11 16:46:33 +03:00
parent 38f0c09175
commit 2994b41089
71 changed files with 211 additions and 1 deletions
@@ -1,6 +1,7 @@
 model_id = "mlx-community/DeepSeek-V3.1-4bit"
 n_layers = 61
 hidden_size = 7168
+num_key_value_heads = 128
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "deepseek"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/DeepSeek-V3.1-8bit"
 n_layers = 61
 hidden_size = 7168
+num_key_value_heads = 128
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "deepseek"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-4.5-Air-8bit"
 n_layers = 46
 hidden_size = 4096
+num_key_value_heads = 8
 supports_tensor = false
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-4.5-Air-bf16"
 n_layers = 46
 hidden_size = 4096
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-4.7-4bit"
 n_layers = 91
 hidden_size = 5120
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-4.7-6bit"
 n_layers = 91
 hidden_size = 5120
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-4.7-8bit-gs32"
 n_layers = 91
 hidden_size = 5120
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-4.7-Flash-4bit"
 n_layers = 47
 hidden_size = 2048
+num_key_value_heads = 20
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-4.7-Flash-5bit"
 n_layers = 47
 hidden_size = 2048
+num_key_value_heads = 20
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-4.7-Flash-6bit"
 n_layers = 47
 hidden_size = 2048
+num_key_value_heads = 20
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-4.7-Flash-8bit"
 n_layers = 47
 hidden_size = 2048
+num_key_value_heads = 20
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-5-8bit-MXFP8"
 n_layers = 78
 hidden_size = 6144
+num_key_value_heads = 64
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-5-MXFP4-Q8"
 n_layers = 78
 hidden_size = 6144
+num_key_value_heads = 64
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-5"
 n_layers = 78
 hidden_size = 6144
+num_key_value_heads = 64
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Kimi-K2-Instruct-4bit"
 n_layers = 61
 hidden_size = 7168
+num_key_value_heads = 64
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "kimi"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Kimi-K2-Thinking"
 n_layers = 61
 hidden_size = 7168
+num_key_value_heads = 64
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "kimi"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Kimi-K2.5"
 n_layers = 61
 hidden_size = 7168
+num_key_value_heads = 64
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "kimi"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Llama-3.2-1B-Instruct-4bit"
 n_layers = 16
 hidden_size = 2048
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "llama"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Llama-3.2-3B-Instruct-4bit"
 n_layers = 28
 hidden_size = 3072
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "llama"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Llama-3.2-3B-Instruct-8bit"
 n_layers = 28
 hidden_size = 3072
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "llama"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Llama-3.3-70B-Instruct-4bit"
 n_layers = 80
 hidden_size = 8192
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "llama"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Llama-3.3-70B-Instruct-8bit"
 n_layers = 80
 hidden_size = 8192
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "llama"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Meta-Llama-3.1-70B-Instruct-4bit"
 n_layers = 80
 hidden_size = 8192
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "llama"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Meta-Llama-3.1-8B-Instruct-4bit"
 n_layers = 32
 hidden_size = 4096
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "llama"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Meta-Llama-3.1-8B-Instruct-8bit"
 n_layers = 32
 hidden_size = 4096
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "llama"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Meta-Llama-3.1-8B-Instruct-bf16"
 n_layers = 32
 hidden_size = 4096
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "llama"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/MiniMax-M2.1-3bit"
 n_layers = 61
 hidden_size = 3072
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "minimax"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/MiniMax-M2.1-8bit"
 n_layers = 61
 hidden_size = 3072
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "minimax"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/MiniMax-M2.5-4bit"
 n_layers = 62
 hidden_size = 3072
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "minimax"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/MiniMax-M2.5-6bit"
 n_layers = 62
 hidden_size = 3072
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "minimax"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/MiniMax-M2.5-8bit"
 n_layers = 62
 hidden_size = 3072
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "minimax"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-0.6B-4bit"
 n_layers = 28
 hidden_size = 1024
+num_key_value_heads = 8
 supports_tensor = false
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-0.6B-8bit"
 n_layers = 28
 hidden_size = 1024
+num_key_value_heads = 8
 supports_tensor = false
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-235B-A22B-Instruct-2507-4bit"
 n_layers = 94
 hidden_size = 4096
+num_key_value_heads = 4
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-235B-A22B-Instruct-2507-8bit"
 n_layers = 94
 hidden_size = 4096
+num_key_value_heads = 4
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-30B-A3B-4bit"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 4
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-30B-A3B-8bit"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 4
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Coder-480B-A35B-Instruct-4bit"
 n_layers = 62
 hidden_size = 6144
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Coder-480B-A35B-Instruct-8bit"
 n_layers = 62
 hidden_size = 6144
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Coder-Next-4bit"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Coder-Next-5bit"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Coder-Next-6bit"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Coder-Next-8bit"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Coder-Next-bf16"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Next-80B-A3B-Instruct-4bit"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Next-80B-A3B-Instruct-8bit"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Next-80B-A3B-Thinking-4bit"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Next-80B-A3B-Thinking-8bit"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-122B-A10B-4bit"
 n_layers = 48
 hidden_size = 3072
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-122B-A10B-6bit"
 n_layers = 48
 hidden_size = 3072
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-122B-A10B-8bit"
 n_layers = 48
 hidden_size = 3072
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-122B-A10B-bf16"
 n_layers = 48
 hidden_size = 3072
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-27B-4bit"
 n_layers = 64
 hidden_size = 5120
+num_key_value_heads = 4
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-27B-8bit"
 n_layers = 64
 hidden_size = 5120
+num_key_value_heads = 4
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-2B-MLX-8bit"
 n_layers = 24
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-35B-A3B-4bit"
 n_layers = 40
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-35B-A3B-8bit"
 n_layers = 40
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-397B-A17B-4bit"
 n_layers = 60
 hidden_size = 4096
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-397B-A17B-6bit"
 n_layers = 60
 hidden_size = 4096
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-397B-A17B-8bit"
 n_layers = 60
 hidden_size = 4096
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-9B-4bit"
 n_layers = 32
 hidden_size = 4096
+num_key_value_heads = 4
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-9B-8bit"
 n_layers = 32
 hidden_size = 4096
+num_key_value_heads = 4
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Step-3.5-Flash-4bit"
 n_layers = 45
 hidden_size = 4096
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "step"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Step-3.5-Flash-6bit"
 n_layers = 45
 hidden_size = 4096
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "step"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Step-3.5-Flash-8Bit"
 n_layers = 45
 hidden_size = 4096
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "step"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/gpt-oss-120b-MXFP4-Q8"
 n_layers = 36
 hidden_size = 2880
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "gpt-oss"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/gpt-oss-20b-MXFP4-Q8"
 n_layers = 24
 hidden_size = 2880
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "gpt-oss"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/llama-3.3-70b-instruct-fp16"
 n_layers = 80
 hidden_size = 8192
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "llama"
@@ -0,0 +1,133 @@
+#!/usr/bin/env python3
+"""Fetch num_key_value_heads from HuggingFace config.json and update TOML model cards.
+
+Usage:
+    # Update only cards missing num_key_value_heads
+    uv run python scripts/fetch_kv_heads.py --missing
+
+    # Update all cards (overwrite existing values)
+    uv run python scripts/fetch_kv_heads.py --all
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import urllib.request
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+
+import tomlkit
+
+CARDS_DIR = (
+    Path(__file__).resolve().parent.parent / "resources" / "inference_model_cards"
+)
+MAX_WORKERS = 5
+
+
+def fetch_kv_heads(model_id: str) -> int | None:
+    """Fetch num_key_value_heads from HuggingFace config.json."""
+    url = f"https://huggingface.co/{model_id}/raw/main/config.json"
+    try:
+        with urllib.request.urlopen(url, timeout=15) as resp:
+            config = json.loads(resp.read())
+    except Exception as e:
+        print(f"  ERROR fetching {url}: {e}", file=sys.stderr)
+        return None
+
+    for source in [config, config.get("text_config", {})]:
+        if "num_key_value_heads" in source:
+            return int(source["num_key_value_heads"])
+
+    return None
+
+
+def update_toml(path: Path, kv_heads: int) -> bool:
+    """Insert or update num_key_value_heads in a TOML file. Returns True if changed."""
+    content = path.read_text()
+    doc = tomlkit.parse(content)
+
+    if doc.get("num_key_value_heads") == kv_heads:
+        return False
+
+    # Insert after hidden_size if adding for the first time
+    if "num_key_value_heads" not in doc:
+        new_doc = tomlkit.document()
+        for key, value in doc.items():
+            new_doc[key] = value
+            if key == "hidden_size":
+                new_doc["num_key_value_heads"] = kv_heads
+        path.write_text(tomlkit.dumps(new_doc))
+    else:
+        doc["num_key_value_heads"] = kv_heads
+        path.write_text(tomlkit.dumps(doc))
+
+    return True
+
+
+def process_card(path: Path) -> tuple[str, str]:
+    """Fetch and update a single card. Returns (filename, status)."""
+    content = path.read_text()
+    doc = tomlkit.parse(content)
+    model_id = doc.get("model_id")
+    if not model_id:
+        return path.name, "SKIP (no model_id)"
+
+    kv_heads = fetch_kv_heads(str(model_id))
+    if kv_heads is None:
+        return path.name, "FAILED"
+
+    changed = update_toml(path, kv_heads)
+    return path.name, f"{kv_heads} ({'UPDATED' if changed else 'UNCHANGED'})"
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Fetch num_key_value_heads from HuggingFace and update TOML cards."
+    )
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument(
+        "--all",
+        action="store_true",
+        help="Update all model cards (overwrite existing values)",
+    )
+    group.add_argument(
+        "--missing",
+        action="store_true",
+        help="Only update cards missing num_key_value_heads",
+    )
+    args = parser.parse_args()
+
+    toml_files = sorted(CARDS_DIR.glob("*.toml"))
+    if not toml_files:
+        print(f"No TOML files found in {CARDS_DIR}", file=sys.stderr)
+        sys.exit(1)
+
+    to_process = []
+    skipped = 0
+
+    for path in toml_files:
+        if args.missing and "num_key_value_heads" in path.read_text():
+            skipped += 1
+            continue
+        to_process.append(path)
+
+    updated = 0
+    failed = 0
+
+    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool:
+        futures = {pool.submit(process_card, path): path for path in to_process}
+        for future in as_completed(futures):
+            name, status = future.result()
+            print(f"  {name}: {status}")
+            if "UPDATED" in status:
+                updated += 1
+            elif "FAILED" in status:
+                failed += 1
+
+    print(f"\nDone: {updated} updated, {skipped} skipped, {failed} failed")
+
+
+if __name__ == "__main__":
+    main()
@@ -90,14 +90,19 @@ 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.
+        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)
        ]
        if not cycles_with_sufficient_memory:
            raise ValueError(
-                f"No tensor sharding found for model with hidden_size {command.model_card.hidden_size} candidate cycles"
+                f"No tensor sharding found for model with "
+                f"hidden_size={command.model_card.hidden_size}"
+                f"{f', num_key_value_heads={kv_heads}' if kv_heads is not None else ''}"
+                f" across candidate cycles"
            )
    if command.sharding == Sharding.Pipeline and command.model_card.model_id == ModelId(
        "mlx-community/DeepSeek-V3.1-8bit"
@@ -83,6 +83,7 @@ class ModelCard(CamelCaseModel):
    n_layers: PositiveInt
    hidden_size: PositiveInt
    supports_tensor: bool
+    num_key_value_heads: PositiveInt | None = None
    tasks: list[ModelTask]
    components: list[ComponentInfo] | None = None
    family: str = ""
@@ -137,6 +138,7 @@ class ModelCard(CamelCaseModel):
            n_layers=num_layers,
            hidden_size=config_data.hidden_size or 0,
            supports_tensor=config_data.supports_tensor,
+            num_key_value_heads=config_data.num_key_value_heads,
            tasks=[ModelTask.TextGeneration],
            trust_remote_code=False,
        )
@@ -170,6 +172,7 @@ class ConfigData(BaseModel):

    architectures: list[str] | None = None
    hidden_size: Annotated[int, Field(ge=0)] | None = None
+    num_key_value_heads: PositiveInt | None = None
    layer_count: int = Field(
        validation_alias=AliasChoices(
            "num_hidden_layers",
@@ -209,6 +212,7 @@ class ConfigData(BaseModel):
        for field in [
            "architectures",
            "hidden_size",
+            "num_key_value_heads",
            "num_hidden_layers",
            "num_layers",
            "n_layer",