[Model Integration] Adding Kimi K2.7 Code

.gitignore

@@ -29,3 +29,8 @@ mstar/worker/ASYNC_REDESIGN.md
 # local AI-assistant context (kept local, not published — cf. vllm-omni)
 CLAUDE.md
 AGENTS.md
+mstar_traces/
+.claude/skills/
+
+# local golden-extraction harness (Kimi-K2.7 port dev tooling; not for main)
+tools/kimi_goldens/

configs/kimi_k2_7.yaml

@@ -0,0 +1,8 @@
+model: "kimi_k2_7"
+# Single-rank Kimi-K2.7 text config; use TP8 configs for the real 1T checkpoint.
+max_seq_len: 262144
+node_groups:
+  - node_names: [LLM]
+    ranks: [0]
+    tp_size: 1
+    graph_walks: [prefill, decode]

configs/kimi_k2_7_code_tp8.yaml

@@ -0,0 +1,21 @@
+model: "kimi_k2_7"
+# Fallback TP8 config for local Kimi-K2.7-Code snapshot; prefer the /dev/shm variant.
+max_seq_len: 8192
+# See kimi_k2_7_code_tp8_shm.yaml — a 1T INT4 load at TP8 outruns PyTorch's
+# default process-group timeout; this path is slower still (disk, not tmpfs).
+dist_timeout_s: 7200
+# Point at your checkpoint at launch — this config hardcodes no path:
+#   mstar-serve --config configs/kimi_k2_7_code_tp8.yaml \
+#               --model-path /path/to/Kimi-K2.7-Code
+# Omit --model-path to pull moonshotai/Kimi-K2.7-Code from HuggingFace.
+model_kwargs:
+  config_variant: k27_code
+  tokenizer_mode: hf
+kv_cache:
+  max_num_pages: 512
+  page_size: 128
+node_groups:
+  - node_names: [LLM]
+    ranks: [0, 1, 2, 3, 4, 5, 6, 7]
+    tp_size: 8
+    graph_walks: [prefill, decode]

configs/kimi_k2_7_code_tp8_shm.yaml

@@ -0,0 +1,21 @@
+model: "kimi_k2_7"
+# TP8 Kimi-K2.7-Code config using a RAM-backed /dev/shm checkpoint copy.
+max_seq_len: 8192
+# Loading ~600GB of INT4 experts across 8 ranks
+# MSTAR_DIST_TIMEOUT_S overrides this.
+dist_timeout_s: 7200
+# Copy the checkpoint into /dev/shm first, then point at it at launch:
+#   mstar serve kimi_k2_7 --model-path /dev/shm/kimi_k2_7_code
+#   mstar-serve --config configs/kimi_k2_7_code_tp8_shm.yaml \
+#               --model-path /dev/shm/kimi_k2_7_code
+model_kwargs:
+  config_variant: k27_code
+  tokenizer_mode: hf
+kv_cache:
+  max_num_pages: 512
+  page_size: 128
+node_groups:
+  - node_names: [LLM]
+    ranks: [0, 1, 2, 3, 4, 5, 6, 7]
+    tp_size: 8
+    graph_walks: [prefill, decode]

configs/synthetic/kimi_k2_7_repro.yaml

@@ -0,0 +1,17 @@
+model: "kimi_k2_7"
+# Reduced synthetic serve config — random weights, NOT a production deployment.
+# Generate the checkpoint first, then point at it:
+#   mstar-serve --config configs/synthetic/kimi_k2_7_repro.yaml \
+#               --model-path tools/kimi_goldens/repro/checkpoint
+max_seq_len: 512
+model_kwargs:
+  config_variant: reduced
+  tokenizer_mode: byte
+kv_cache:
+  max_num_pages: 256
+  page_size: 128
+node_groups:
+  - node_names: [LLM]
+    ranks: [0]
+    tp_size: 1
+    graph_walks: [prefill, decode]

configs/synthetic/kimi_k2_7_tp2.yaml

@@ -0,0 +1,17 @@
+model: "kimi_k2_7"
+# Reduced synthetic TP=2 config — random weights, NOT a production deployment.
+# Generate the checkpoint first, then point at it:
+#   mstar-serve --config configs/synthetic/kimi_k2_7_tp2.yaml \
+#               --model-path tools/kimi_goldens/repro/checkpoint
+max_seq_len: 512
+model_kwargs:
+  config_variant: reduced
+  tokenizer_mode: byte
+kv_cache:
+  max_num_pages: 256
+  page_size: 128
+node_groups:
+  - node_names: [LLM]
+    ranks: [0, 1]
+    tp_size: 2
+    graph_walks: [prefill, decode]

docs/environment_variables.rst

@@ -35,6 +35,15 @@ Communication
      - ``19000``
      - Base of the deterministic entity-id → TCP port map (``api_server``
        = base, ``conductor`` = base+1, ``worker_<rank>`` = base+100+rank).
+   * - ``MSTAR_DIST_TIMEOUT_S``
+     - config's ``dist_timeout_s``
+     - Timeout in seconds for the NCCL world group and its parallel
+       subgroups (:func:`mstar.distributed.communication.resolve_dist_timeout`).
+       Overrides the deployment config's ``dist_timeout_s``; with neither set,
+       PyTorch's default applies. Raise it only where weight load, JIT or
+       CUDA-graph capture can exceed that default (a 1T MoE at TP8 does) —
+       a hung collective takes correspondingly longer to abort. Must be set
+       before the conductor spawns workers, which inherit it.
    * - ``MSTAR_SHM_ARENA``
      - ``0``
      - SHM tensor-transport implementation. ``0``: per-uuid files.

docs/serving.rst

@@ -44,6 +44,10 @@ mstar serve
    * - ``--cache-dir``
      - HF default
      - HuggingFace weight cache directory.
+   * - ``--model-path``
+     - registry default
+     - Local checkpoint directory or HF repo id to load weights from. Use this
+       instead of hardcoding a path in the config YAML.
    * - ``--tensor-comm-protocol``
      - ``SHM``
      - Tensor transport: ``SHM`` (safe single-node default), ``TCP``, or ``RDMA``.
@@ -122,6 +126,16 @@ mstar-serve
    * - ``--cache-dir``
      - HF default
      - HuggingFace weight cache directory.
+   * - ``--model-path``
+     - registry default
+     - Local checkpoint directory or HF repo id to load weights from. Applies to
+       any model and overrides the registry default, so a config YAML never has
+       to hardcode one machine's filesystem layout:
+
+       .. code-block:: bash
+
+          mstar-serve --config configs/kimi_k2_7_code_tp8_shm.yaml \
+                      --model-path /dev/shm/kimi_k2_7_code
    * - ``--socket-path-prefix``
      - ``/tmp/mstar``
      - ZMQ IPC socket prefix (shared with conductor/workers).
@@ -178,6 +192,9 @@ A config maps the model's computation-graph nodes to physical GPU ranks. The key
        scoped to specific ``graph_walks`` and/or sharded with ``tp_size``.
    * - ``model_kwargs``
      - *(optional)* Server-init model parameters (see below).
+   * - ``dist_timeout_s``
+     - *(optional)* Timeout in seconds for the NCCL world group and its parallel
+       subgroups. Unset keeps PyTorch's default.
 
 Node names are model-specific — they are the keys of the model's
 ``get_node_engine_types`` (e.g. BAGEL's ``vit_encoder`` / ``vae_encoder`` / ``LLM``,
@@ -208,6 +225,11 @@ Because placement is config-only, the *same* model code runs single-GPU or fully
 disaggregated. ``configs/`` ships several layouts per model (``*_single_gpu``,
 ``*_colocated``, ``*_pd_disaggregated``, ``*_cfg_parallel``, …).
 
+``configs/synthetic/`` holds reduced, randomly-initialized deployments used to
+exercise the serving path without real weights (shape/plumbing checks, TP
+sanity). They are **not** production configs — they load a generated checkpoint
+and emit meaningless tokens.
+
 **Tensor parallelism.** Shard a node across GPUs with ``tp_size`` and that many ``ranks``:
 
 .. code-block:: yaml

mstar/api_server/entrypoint.py

@@ -52,6 +52,19 @@ def _detect_modality(filename: str) -> str:
 # Conductor process target (top-level for picklability with spawn)
 # ------------------------------------------------------------------
 
+def _resolve_model_path(model_name: str, model_path: str | None) -> str:
+    """Where to load weights from: ``--model-path`` if given, else the registry default.
+
+    ``--model-path`` takes a local directory or an HF repo id and applies to any
+    model, so a deployment config never has to hardcode one machine's filesystem
+    layout. Both the API-server-side model instance and the conductor-side one
+    resolve through here, so they cannot diverge.
+    """
+    if model_path:
+        return model_path
+    return HF_MODELS.get(model_name, {}).get("model_path_hf", "")
+
+
 def _conductor_process_target(
     model_name: str,
     config_path: str,
@@ -61,7 +74,8 @@ def _conductor_process_target(
     log_level: str = "INFO",
     cache_dir: str | None = None,
     tensor_comm_protocol=CommProtocol.RDMA,
-    tcp_transfer_device=""
+    tcp_transfer_device="",
+    model_path: str | None = None,
 ):
     """Runs DummyConductor.run() in a spawned process."""
     logging.basicConfig(
@@ -93,7 +107,7 @@ def _conductor_process_target(
         )
 
     model = get_model_class(model_name)(
-        model_path_hf=HF_MODELS.get(model_name, {}).get("model_path_hf", ""),
+        model_path_hf=_resolve_model_path(model_name, model_path),
         cache_dir=cache_dir,
         **yaml_model_kwargs,
     )
@@ -866,6 +880,12 @@ def main(argv: list[str] | None = None):
         "--cache-dir", type=str, default=None,
         help="Directory for caching downloaded HuggingFace model files",
     )
+    parser.add_argument(
+        "--model-path", type=str, default=None,
+        help="Where to load weights from — a local checkpoint directory or an HF "
+             "repo id. Overrides the model's registry default so deployment "
+             "configs need not hardcode a filesystem path.",
+    )
     parser.add_argument(
         "--log-level", type=str, default="INFO",
         choices=["DEBUG", "INFO", "WARNING", "ERROR"],
@@ -900,7 +920,7 @@ def main(argv: list[str] | None = None):
     # (tokenization only — no GPU weights needed)
     from mstar.model.registry import get_model_class
     model = get_model_class(model_name)(
-        model_path_hf=HF_MODELS.get(model_name, {}).get("model_path_hf", ""),
+        model_path_hf=_resolve_model_path(model_name, args.model_path),
         cache_dir=args.cache_dir,
         **yaml_model_kwargs,
     )
@@ -932,7 +952,8 @@ def main(argv: list[str] | None = None):
             args.log_level,
             args.cache_dir,
             CommProtocol(args.tensor_comm_protocol),
-            args.tcp_transfer_device
+            args.tcp_transfer_device,
+            args.model_path,
         ),
     )
     conductor_proc.start()

mstar/cli/main.py

@@ -38,6 +38,7 @@
     "whisper_large": "whisper_large.yaml",
     "higgs_audio": "higgs_audio.yaml",
     "wan22": "wan22.yaml",
+    "kimi_k2_7": "kimi_k2_7_code_tp8_shm.yaml",
 }
 
 
@@ -149,6 +150,8 @@ def _serve(args: argparse.Namespace) -> None:
     ]
     if args.cache_dir:
         argv += ["--cache-dir", args.cache_dir]
+    if args.model_path:
+        argv += ["--model-path", args.model_path]
     if args.log_stats:
         argv += ["--log-stats"]
     if args.log_stats_file:
@@ -172,6 +175,11 @@ def build_parser() -> argparse.ArgumentParser:
     serve.add_argument("--gpus", default=None, help="CUDA_VISIBLE_DEVICES, e.g. '0' or '0,1,2'")
     serve.add_argument("--config", default=None, help="override the default config (path to YAML)")
     serve.add_argument("--cache-dir", default=None, help="HuggingFace weight cache directory")
+    serve.add_argument(
+        "--model-path", default=None,
+        help="local checkpoint directory or HF repo id to load weights from "
+             "(overrides the model's registry default)",
+    )
     serve.add_argument("--socket-path-prefix", default=None, help="ZMQ IPC socket prefix")
     serve.add_argument("--upload-dir", default=None, help="temp dir for uploaded media")
     serve.add_argument(

mstar/conductor/conductor.py

@@ -414,6 +414,7 @@ def _derive_worker_info(self):
         self.parallel_config = GlobalParallelConfig(
             worker_graphs=self.worker_graphs,
             worker_ids=self.worker_ids,
+            dist_timeout_s=self.model_config.get("dist_timeout_s"),
         )
 
     def _launch_workers(self):

mstar/distributed/communication.py

@@ -1,9 +1,30 @@
+import os
 from dataclasses import dataclass, field
+from datetime import timedelta
 from typing import Any
 
 import torch
 import torch.distributed as dist
 
+DIST_TIMEOUT_ENV = "MSTAR_DIST_TIMEOUT_S"
+
+
+def resolve_dist_timeout(dist_timeout_s: float | None = None) -> dict[str, timedelta]:
+    """Build the ``timeout`` kwarg for ``init_process_group`` / ``new_group``."""
+    raw = os.environ.get(DIST_TIMEOUT_ENV, "").strip()
+    if raw:
+        try:
+            dist_timeout_s = float(raw)
+        except ValueError as exc:
+            raise ValueError(
+                f"{DIST_TIMEOUT_ENV} must be a number of seconds, got {raw!r}"
+            ) from exc
+    if dist_timeout_s is None:
+        return {}
+    if dist_timeout_s <= 0:
+        raise ValueError(f"Distributed timeout must be positive, got {dist_timeout_s}")
+    return {"timeout": timedelta(seconds=float(dist_timeout_s))}
+
 
 class CommGroup:
     """A communication group over one axis of the worker device mesh.
@@ -192,6 +213,8 @@ class WorkerParallelGroups:
     # projections).
     node_to_tp_group: dict[str, CommGroup] = field(default_factory=dict)
     node_to_sp_group: dict[str, CommGroup] = field(default_factory=dict)
+    # Process-group timeout in seconds, from the deployment config's
+    dist_timeout_s: float | None = None
 
     def add(self, node: str, comm_group: CommGroup):
         # disallow colocation of multiple comm groups on the same node
@@ -233,11 +256,13 @@ def init_dist(
         if not self.any_parallelism:
             return
 
+        timeout_kwargs = resolve_dist_timeout(self.dist_timeout_s)
         dist.init_process_group(
             backend="nccl",
             init_method=init_method,
             world_size=self.num_workers,
             rank=self.global_rank,
+            **timeout_kwargs,
         )
 
         # One subgroup per distinct rank tuple across BOTH mesh axes —
@@ -247,7 +272,9 @@ def init_dist(
         # an SP group (degenerate meshes) maps to one subgroup.
         rank_tuple_to_pg: dict[tuple[int, ...], "dist.ProcessGroup"] = {}
         for rank_tuple in self.world_parallel_groups:
-            rank_tuple_to_pg[rank_tuple] = dist.new_group(ranks=list(rank_tuple))
+            rank_tuple_to_pg[rank_tuple] = dist.new_group(
+                ranks=list(rank_tuple), **timeout_kwargs
+            )
 
         seen: set[int] = set()
         for comm_group in (
@@ -309,7 +336,8 @@ class GlobalParallelConfig:
     def __init__(
         # leaving type annotation as Any due to circular import
         self, worker_graphs: dict[str, Any],
-        worker_ids: list[str]
+        worker_ids: list[str],
+        dist_timeout_s: float | None = None,
     ):
         self.num_workers = len(worker_ids)
         any_parallelism = any(
@@ -337,6 +365,7 @@ def __init__(
                 global_rank=i, num_workers=self.num_workers,
                 any_parallelism=any_parallelism,
                 world_parallel_groups=world_parallel_groups,
+                dist_timeout_s=dist_timeout_s,
             ) for i, wid in enumerate(worker_ids)
         }
 
@@ -372,4 +401,3 @@ def __init__(
                         self.per_worker_config[worker_ids[rank]].add_sp(
                             node, self.sp_comm_groups[key]
                         )
-