From 5332a063e43d0172bef1a5b926c768a75e15f120 Mon Sep 17 00:00:00 2001
From: Nan <nan@modal.com>
Date: Mon, 10 Aug 2026 16:18:34 +0000
Subject: [PATCH 1/5] [cookbook] Pin the unified SGLang v0.5.17 runtime

---
 cookbook/common/serving_image.py              |  6 +++---
 .../miles_disagg/configs/kimi_k3_mxfp4.py     | 20 ++++++-------------
 2 files changed, 9 insertions(+), 17 deletions(-)

diff --git a/cookbook/common/serving_image.py b/cookbook/common/serving_image.py
index 7f1c021b..b50d4414 100644
--- a/cookbook/common/serving_image.py
+++ b/cookbook/common/serving_image.py
@@ -27,10 +27,10 @@ class SGLangRuntime:
 
 
 DEFAULT_SGLANG_RUNTIME = SGLangRuntime(
-    image="lmsysorg/sglang:v0.5.16",
+    image="lmsysorg/sglang:v0.5.17",
     repository="https://github.com/modal-projects/sglang.git",
-    branch="stitch-sglang-v0.5.16",
-    commit="a73ea9507fb981462768dbc5e869bdfeb5c48116",
+    branch="stitch-sglang-v0.5.17",
+    commit="0c79627e857eec795298a372adadf649209cdf2f",
 )
 
 _COOKBOOK_DIR = Path(__file__).resolve().parent.parent
diff --git a/cookbook/miles_disagg/configs/kimi_k3_mxfp4.py b/cookbook/miles_disagg/configs/kimi_k3_mxfp4.py
index 3a6d3e6b..d145931f 100644
--- a/cookbook/miles_disagg/configs/kimi_k3_mxfp4.py
+++ b/cookbook/miles_disagg/configs/kimi_k3_mxfp4.py
@@ -1,28 +1,20 @@
 """Kimi K3 rollout server using its native MXFP4 checkpoint."""
 
 from cookbook.common.config import ModalConfig
-from cookbook.common.serving_image import SGLangRuntime
 
 ROLLOUT_SOURCE_MODEL = "moonshotai/Kimi-K3"
 ROLLOUT_SOURCE_REVISION = "9f62e4e9fffbd0a83ddd60e1c209d828994b3569"
 ROLLOUT_NUM_GPUS_PER_ENGINE = 8
-SGLANG_DELTA_UPDATE_MODE = "disk"
-
-SGLANG_RUNTIME = SGLangRuntime(
-    image=(
-        "lmsysorg/sglang@"
-        "sha256:81a9c00654b3e4c7c681a4728a64fcb4853aa698dc9fea1959bbf4eb26bfb2e5"
-    ),
-    repository="https://github.com/modal-projects/sglang.git",
-    branch="stitch-sglang-kimi-k3",
-    commit="977e33dd4560463d3a9a981ba3ed764de289e858",
-)
+SGLANG_DELTA_UPDATE_MODE = "cpu"
 
 SGLANG_SERVER_ARGS = {
     "--trust-remote-code": "",
     "--load-format": "fastsafetensors",
     "--model-loader-extra-config": '{"enable_gds":false}',
     "--weight-loader-drop-cache-after-load": "",
+    "--enable-cpu-weight-cache": "",
+    "--cpu-weight-cache-max-compile-group-gb": "16",
+    "--cpu-weight-cache-canonical-checkpoint-dir": "/local-checkpoint/canonical",
     "--dist-timeout": "3600",
     "--context-length": "1048576",
     "--max-running-requests": "32",
@@ -44,7 +36,7 @@
     gpu="B300",
     rollout_target_inputs=32,
     rollout_ephemeral_disk_mib=2 * 1024 * 1024,
-    # Disk updates reconstruct the ~600 GiB target checkpoint on ephemeral NVMe;
-    # RAM stays generous as read cache for the canonical base.
+    # Keep the ~1.56 TB canonical checkpoint on NVMe. The eight rank-ready CPU
+    # images retain ~1.66 TB before engine and bounded staging overhead.
     rollout_memory_mib=(1024 * 1024, 3 * 1024 * 1024),
 )

From 5f70b9a7355a5a4186165c670e81c41334d58a67 Mon Sep 17 00:00:00 2001
From: Nan <nan@modal.com>
Date: Mon, 10 Aug 2026 16:18:42 +0000
Subject: [PATCH 2/5] [profiling] Cover the v0.5.17 weight-update paths

---
 tools/profiling/_delta_weight_update.py       | 145 ++++++++++++---
 tools/profiling/_hf_checkpoint.py             |  61 +++++++
 tools/profiling/_synthetic_delta.py           | 168 ------------------
 .../glm45_air_fp8_delta_weight_update.py      |  34 +++-
 .../glm5_2_nvfp4_delta_weight_update.py       |  72 +++++++-
 .../kimi_k2_6_nvfp4_delta_weight_update.py    | 108 +++++++----
 .../kimi_k3_mxfp4_delta_weight_update.py      | 142 +++++----------
 7 files changed, 394 insertions(+), 336 deletions(-)
 create mode 100644 tools/profiling/_hf_checkpoint.py

diff --git a/tools/profiling/_delta_weight_update.py b/tools/profiling/_delta_weight_update.py
index 11200cbc..43d459a1 100644
--- a/tools/profiling/_delta_weight_update.py
+++ b/tools/profiling/_delta_weight_update.py
@@ -22,6 +22,7 @@
 from typing import Any, Literal
 
 UpdateMode = Literal["disk", "cpu"]
+CanonicalStorage = Literal["memory", "disk"]
 
 
 @dataclass(frozen=True)
@@ -29,6 +30,7 @@ class WeightUpdateSpec:
     model_name: str
     base_checkpoint_dir: str
     local_target_checkpoint_dir: str
+    local_canonical_checkpoint_dir: str
     server_args: dict[str, str]
     tp_size: int = 4
     port: int = 8001
@@ -37,14 +39,27 @@ class WeightUpdateSpec:
 def server_args_for_mode(
     server_args: dict[str, str],
     update_mode: UpdateMode,
+    canonical_storage: CanonicalStorage | None,
+    canonical_checkpoint_dir: str,
 ) -> dict[str, str]:
     """Return direct SGLang arguments for one update mode."""
 
     result = dict(server_args)
     if update_mode == "cpu":
+        if canonical_storage not in {"memory", "disk"}:
+            raise ValueError(
+                "canonical_storage must be 'memory' or 'disk' for CPU updates"
+            )
         result["--enable-cpu-weight-cache"] = ""
         result.setdefault("--cpu-weight-cache-max-compile-group-gb", "8")
+        result.pop("--cpu-weight-cache-canonical-checkpoint-dir", None)
+        if canonical_storage == "disk":
+            result["--cpu-weight-cache-canonical-checkpoint-dir"] = (
+                canonical_checkpoint_dir
+            )
     elif update_mode == "disk":
+        if canonical_storage is not None:
+            raise ValueError("canonical_storage applies only to CPU updates")
         result.pop("--enable-cpu-weight-cache", None)
         result.pop("--cpu-weight-cache-max-compile-group-gb", None)
         result.pop("--cpu-weight-cache-canonical-checkpoint-dir", None)
@@ -59,6 +74,25 @@ def parse_update_mode(value: str) -> UpdateMode:
     return value
 
 
+def parse_canonical_storage(value: str | None) -> CanonicalStorage | None:
+    if value not in {None, "memory", "disk"}:
+        raise ValueError("canonical_storage must be 'memory' or 'disk'")
+    return value
+
+
+def parse_update_destination(
+    update_mode: str,
+    canonical_storage: str | None,
+) -> tuple[UpdateMode, CanonicalStorage | None]:
+    mode = parse_update_mode(update_mode)
+    storage = parse_canonical_storage(canonical_storage)
+    if mode == "cpu" and storage is None:
+        raise ValueError("CPU updates require --canonical-storage memory or disk")
+    if mode == "disk" and storage is not None:
+        raise ValueError("--canonical-storage applies only to CPU updates")
+    return mode, storage
+
+
 def modal_runtime_label() -> str:
     value = os.environ.get("MODAL_FUNCTION_RUNTIME")
     if not value:
@@ -113,14 +147,19 @@ def _generate(
         }
     ]
     if not fingerprint or not fingerprint_logprobs:
+        request = {
+            "model": model,
+            "messages": messages,
+            "temperature": 0,
+            "max_tokens": 96 if fingerprint else 80,
+        }
+        if fingerprint:
+            # Compare one deterministic DP worker rather than treating normal
+            # cross-rank numerical drift as a weight-update failure.
+            request["routed_dp_rank"] = 0
         response = httpx.post(
             f"{url}/v1/chat/completions",
-            json={
-                "model": model,
-                "messages": messages,
-                "temperature": 0,
-                "max_tokens": 96 if fingerprint else 80,
-            },
+            json=request,
             timeout=300,
             trust_env=False,
         )
@@ -144,6 +183,7 @@ def _generate(
         f"{url}/generate",
         json={
             "text": "Explain why the Moon has phases in exactly three short clauses.",
+            "routed_dp_rank": 0,
             "sampling_params": {
                 "temperature": 0,
                 "max_new_tokens": 48 if fingerprint else 80,
@@ -314,6 +354,51 @@ def summary(self) -> dict[str, Any]:
         }
 
 
+class _MemoryProbe:
+    def __init__(self, interval_s: float = 1.0) -> None:
+        self.interval_s = interval_s
+        self.stop = threading.Event()
+        self.samples: list[dict[str, int | str]] = []
+        self.thread = threading.Thread(
+            target=self._run,
+            name="memory-during-weight-stage",
+            daemon=True,
+        )
+
+    def _run(self) -> None:
+        while not self.stop.is_set():
+            self.samples.append(_memory_snapshot())
+            self.stop.wait(self.interval_s)
+
+    def __enter__(self) -> _MemoryProbe:
+        self.thread.start()
+        return self
+
+    def __exit__(self, exc_type, exc, traceback) -> None:
+        self.stop.set()
+        self.thread.join(timeout=max(5.0, 2 * self.interval_s))
+
+    def summary(self) -> dict[str, int | str | None]:
+        numeric_keys = {
+            key
+            for sample in self.samples
+            for key, value in sample.items()
+            if isinstance(value, int)
+        }
+        result: dict[str, int | str | None] = {"samples": len(self.samples)}
+        for key in sorted(numeric_keys):
+            values = [
+                value
+                for sample in self.samples
+                if isinstance(value := sample.get(key), int)
+            ]
+            if key == "MemAvailable_bytes":
+                result[f"{key}_min"] = min(values)
+            else:
+                result[f"{key}_max"] = max(values)
+        return result
+
+
 def _memory_snapshot() -> dict[str, int | str]:
     result: dict[str, int | str] = {}
     cgroup_files = {
@@ -408,12 +493,14 @@ def _print_profile_summary(results: dict[str, Any]) -> None:
             "model",
             "runtime",
             "update_mode",
+            "canonical_storage",
             "sample_id",
             "status",
             "initial_load_s",
             "destination_init_s",
             "destination_init_cpu",
             "generation_during_destination_init",
+            "memory_during_destination_init",
             "stage_s",
             "stage_cpu",
             "commit_rpc_s",
@@ -446,6 +533,7 @@ def run_delta_weight_update(
     source_dir: str,
     target_version: int,
     update_mode: UpdateMode,
+    canonical_storage: CanonicalStorage | None,
     runtime: str,
     sample_id: str,
 ) -> dict[str, Any]:
@@ -471,6 +559,7 @@ def run_delta_weight_update(
         "source_dir": source_dir,
         "target_version": target_version,
         "update_mode": update_mode,
+        "canonical_storage": canonical_storage,
         "runtime": runtime,
         "sample_id": sample_id,
         "tp_size": spec.tp_size,
@@ -483,7 +572,12 @@ def run_delta_weight_update(
         model_path=spec.base_checkpoint_dir,
         worker_port=spec.port,
         tp=spec.tp_size,
-        extra_server_args=server_args_for_mode(spec.server_args, update_mode),
+        extra_server_args=server_args_for_mode(
+            spec.server_args,
+            update_mode,
+            canonical_storage,
+            spec.local_canonical_checkpoint_dir,
+        ),
         # Allow for a cold, model-sized initial load. Destination initialization
         # is measured separately after the endpoint begins serving.
         health_timeout=2 * 60 * 60,
@@ -531,22 +625,28 @@ def run_delta_weight_update(
         with httpx.Client(timeout=None, trust_env=False) as client:
             destination_init_started = time.perf_counter()
             destination_init_cpu_started = _cgroup_cpu_usage_s()
-            with _GenerationProbe(url) as generation:
-                init_payload = {
-                    "base_checkpoint_dir": spec.base_checkpoint_dir,
-                    "target_version": 0,
-                    "destination": update_mode,
-                }
-                if update_mode == "disk":
-                    init_payload["local_checkpoint_dir"] = (
-                        spec.local_target_checkpoint_dir
+            generation = _GenerationProbe(url)
+            memory = _MemoryProbe()
+            try:
+                with generation, memory:
+                    init_payload = {
+                        "base_checkpoint_dir": spec.base_checkpoint_dir,
+                        "target_version": 0,
+                        "destination": update_mode,
+                    }
+                    if update_mode == "disk":
+                        init_payload["local_checkpoint_dir"] = (
+                            spec.local_target_checkpoint_dir
+                        )
+                    initialized = _post(
+                        client,
+                        url,
+                        "/stage_weight_update",
+                        init_payload,
                     )
-                initialized = _post(
-                    client,
-                    url,
-                    "/stage_weight_update",
-                    init_payload,
-                )
+            finally:
+                results["generation_during_destination_init"] = generation.summary()
+                results["memory_during_destination_init"] = memory.summary()
             results["destination_init_s"] = round(
                 time.perf_counter() - destination_init_started,
                 6,
@@ -557,7 +657,6 @@ def run_delta_weight_update(
                 results["destination_init_s"],
             )
             results["destination_init_rank_stats"] = initialized.get("rank_stats")
-            results["generation_during_destination_init"] = generation.summary()
             results["memory_after_destination_init"] = _memory_snapshot()
             if generation.errors or not generation.samples:
                 raise RuntimeError(
diff --git a/tools/profiling/_hf_checkpoint.py b/tools/profiling/_hf_checkpoint.py
new file mode 100644
index 00000000..2120ecbc
--- /dev/null
+++ b/tools/profiling/_hf_checkpoint.py
@@ -0,0 +1,61 @@
+"""Helpers for profiling public Hugging Face checkpoints on Modal Volumes."""
+
+from __future__ import annotations
+
+import shutil
+from collections.abc import Callable
+from pathlib import Path
+
+
+def download_snapshot(
+    repo_id: str,
+    revision: str,
+    cache_dir: str,
+    *,
+    commit: Callable[[], None],
+) -> str:
+    """Download and durably commit a model-sized snapshot in bounded batches."""
+
+    from huggingface_hub import HfApi, snapshot_download
+
+    files = HfApi().list_repo_files(repo_id=repo_id, revision=revision)
+    weights = sorted(path for path in files if path.endswith(".safetensors"))
+    batches = [sorted(set(files) - set(weights))]
+    batches.extend(weights[offset : offset + 4] for offset in range(0, len(weights), 4))
+    for index, batch in enumerate(batches, start=1):
+        snapshot_download(
+            repo_id=repo_id,
+            revision=revision,
+            cache_dir=cache_dir,
+            allow_patterns=batch,
+            max_workers=4,
+        )
+        commit()
+        print(f"Committed checkpoint download batch {index}/{len(batches)}")
+
+    path = snapshot_download(
+        repo_id=repo_id,
+        revision=revision,
+        cache_dir=cache_dir,
+        local_files_only=True,
+    )
+    print(f"Downloaded {repo_id}@{revision} to {path}")
+    return path
+
+
+def materialize_checkpoint_view(source_dir: str, target_dir: str) -> None:
+    """Expose cached weights and real sibling metadata as a local checkpoint."""
+
+    source = Path(source_dir)
+    target = Path(target_dir)
+    shutil.rmtree(target, ignore_errors=True)
+    target.mkdir(parents=True)
+    for path in source.rglob("*"):
+        destination = target / path.relative_to(source)
+        if path.is_dir():
+            destination.mkdir(parents=True, exist_ok=True)
+        elif path.suffix == ".safetensors":
+            destination.symlink_to(path.resolve())
+        else:
+            destination.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copy2(path, destination, follow_symlinks=True)
diff --git a/tools/profiling/_synthetic_delta.py b/tools/profiling/_synthetic_delta.py
index da6ea39d..22238e82 100644
--- a/tools/profiling/_synthetic_delta.py
+++ b/tools/profiling/_synthetic_delta.py
@@ -4,7 +4,6 @@
 
 import hashlib
 import json
-import math
 import os
 import shutil
 import struct
@@ -787,170 +786,3 @@ def prepare_standard_delta(
     print(f"Committed synthetic delta at {source_dir}")
     print(f"SYNTHETIC_DELTA={json.dumps(result, sort_keys=True)}")
     return result
-
-
-_FULL_COVERAGE_CHANGE_STRIDE_BYTES = 1024 * 1024
-_FULL_COVERAGE_ZERO_CHUNK = bytes(_STREAM_BYTES)
-_FULL_COVERAGE_XOR_CHUNK = bytearray(_FULL_COVERAGE_ZERO_CHUNK)
-_FULL_COVERAGE_XOR_CHUNK[::_FULL_COVERAGE_CHANGE_STRIDE_BYTES] = b"\x01" * (
-    _STREAM_BYTES // _FULL_COVERAGE_CHANGE_STRIDE_BYTES
-)
-_FULL_COVERAGE_XOR_CHUNK = bytes(_FULL_COVERAGE_XOR_CHUNK)
-
-
-def _compress_full_coverage_xor(byte_count: int) -> bytes:
-    import zstandard
-
-    compressor = zstandard.ZstdCompressor(level=1).compressobj()
-    output = []
-    remaining = byte_count
-    while remaining:
-        size = min(remaining, len(_FULL_COVERAGE_XOR_CHUNK))
-        output.append(compressor.compress(memoryview(_FULL_COVERAGE_XOR_CHUNK)[:size]))
-        remaining -= size
-    output.append(compressor.flush())
-    return b"".join(output)
-
-
-def _full_coverage_target_checksum(
-    handle: Any,
-    *,
-    offset: int,
-    byte_count: int,
-) -> str:
-    import xxhash
-
-    checksum = xxhash.xxh3_128()
-    handle.seek(offset)
-    remaining = byte_count
-    position = 0
-    while remaining:
-        chunk = handle.read(min(remaining, _STREAM_BYTES))
-        if not chunk:
-            raise RuntimeError("checkpoint tensor ended before its declared size")
-        first_change = (-position) % _FULL_COVERAGE_CHANGE_STRIDE_BYTES
-        if first_change < len(chunk):
-            changed = bytearray(chunk)
-            for index in range(
-                first_change,
-                len(changed),
-                _FULL_COVERAGE_CHANGE_STRIDE_BYTES,
-            ):
-                changed[index] ^= 1
-            chunk = changed
-        checksum.update(chunk)
-        remaining -= len(chunk)
-        position += len(chunk)
-    return checksum.hexdigest()
-
-
-def write_full_coverage_delta(
-    checkpoint_dir: str,
-    source_dir: str,
-    *,
-    workers: int = 8,
-) -> dict[str, Any]:
-    """Write one sparse DELTA with an element-level change in every tensor."""
-
-    import numpy as np
-    import safetensors.numpy
-
-    started = time.perf_counter()
-    checkpoint = Path(checkpoint_dir)
-    index = json.loads((checkpoint / "model.safetensors.index.json").read_text())
-    weight_map = index["weight_map"]
-    names_by_shard: dict[str, list[str]] = {}
-    for name, filename in weight_map.items():
-        names_by_shard.setdefault(filename, []).append(name)
-
-    target_dir = Path(source_dir) / "weight_v000001"
-    target_dir.mkdir(parents=True, exist_ok=False)
-
-    def write_shard(item: tuple[str, list[str]]) -> dict[str, int | str]:
-        filename, names = item
-        source_path = checkpoint / filename
-        data_start, header = _safetensors_header(source_path)
-        names.sort(key=lambda name: header[name]["data_offsets"][0])
-        compressed_tensors: dict[str, np.ndarray] = {}
-        checksums: dict[str, str] = {}
-        tensor_bytes = 0
-        changed_tensors = 0
-        changed_bytes = 0
-        with source_path.open("rb") as handle:
-            for name in names:
-                begin, end = header[name]["data_offsets"]
-                byte_count = end - begin
-                checksums[name] = _full_coverage_target_checksum(
-                    handle,
-                    offset=data_start + begin,
-                    byte_count=byte_count,
-                )
-                compressed_tensors[name] = np.frombuffer(
-                    _compress_full_coverage_xor(byte_count),
-                    dtype=np.uint8,
-                )
-                tensor_bytes += byte_count
-                changed_tensors += int(byte_count > 0)
-                changed_bytes += math.ceil(
-                    byte_count / _FULL_COVERAGE_CHANGE_STRIDE_BYTES
-                )
-            if hasattr(os, "posix_fadvise"):
-                try:
-                    os.posix_fadvise(
-                        handle.fileno(),
-                        0,
-                        0,
-                        os.POSIX_FADV_DONTNEED,
-                    )
-                except OSError:
-                    pass
-
-        target_path = target_dir / filename
-        safetensors.numpy.save_file(
-            compressed_tensors,
-            target_path,
-            metadata=checksums,
-        )
-        result: dict[str, int | str] = {
-            "filename": filename,
-            "tensors": len(names),
-            "tensor_bytes": tensor_bytes,
-            "changed_tensors": changed_tensors,
-            "changed_bytes": changed_bytes,
-            "wire_bytes": target_path.stat().st_size,
-        }
-        print(
-            f"SYNTHETIC_DELTA_SHARD={json.dumps(result, sort_keys=True)}",
-            flush=True,
-        )
-        return result
-
-    with ThreadPoolExecutor(
-        max_workers=min(workers, len(names_by_shard)),
-    ) as executor:
-        shards = list(executor.map(write_shard, names_by_shard.items()))
-
-    (target_dir / "model.safetensors.index.json").write_text(
-        json.dumps(
-            {
-                "metadata": {
-                    "version": "000001",
-                    "base_version": "000000",
-                    "delta_encoding": "xor",
-                    "compression_format": "zstd",
-                    "checksum_format": "xxh3-128",
-                },
-                "weight_map": weight_map,
-            }
-        )
-    )
-    return {
-        "scope": "all_tensors",
-        "tensors": len(weight_map),
-        "changed_tensors": sum(int(shard["changed_tensors"]) for shard in shards),
-        "tensor_bytes": sum(int(shard["tensor_bytes"]) for shard in shards),
-        "changed_bytes": sum(int(shard["changed_bytes"]) for shard in shards),
-        "change_stride_bytes": _FULL_COVERAGE_CHANGE_STRIDE_BYTES,
-        "wire_bytes": sum(int(shard["wire_bytes"]) for shard in shards),
-        "generation_s": round(time.perf_counter() - started, 6),
-    }
diff --git a/tools/profiling/glm45_air_fp8_delta_weight_update.py b/tools/profiling/glm45_air_fp8_delta_weight_update.py
index f31c05fc..579b0c3a 100644
--- a/tools/profiling/glm45_air_fp8_delta_weight_update.py
+++ b/tools/profiling/glm45_air_fp8_delta_weight_update.py
@@ -5,7 +5,7 @@
 
     MODAL_FUNCTION_RUNTIME=runc uv run --extra modal modal run -d \
       tools/profiling/glm45_air_fp8_delta_weight_update.py \
-      --update-mode cpu
+      --update-mode cpu --canonical-storage memory
 """
 
 from __future__ import annotations
@@ -23,6 +23,8 @@
 from tools.profiling._delta_weight_update import (
     WeightUpdateSpec,
     modal_runtime_label,
+    parse_canonical_storage,
+    parse_update_destination,
     parse_update_mode,
     run_delta_weight_update,
 )
@@ -49,6 +51,7 @@
 DELTA_SOURCE_DIR = f"{DELTA_MOUNT}/{DELTA_ID}"
 BASE_CHECKPOINT_DIR = str(model.ROLLOUT_CHECKPOINT_PATH)
 LOCAL_TARGET_CHECKPOINT_DIR = "/local-checkpoint/glm45-air-fp8/target"
+LOCAL_CANONICAL_CHECKPOINT_DIR = "/local-checkpoint/glm45-air-fp8/canonical"
 SGLANG_CACHE_PATH = "/root/.cache/sglang"
 SOURCE_MARKER = ".stitch-source.json"
 
@@ -180,24 +183,43 @@ def prepare_delta() -> dict:
     },
     timeout=2 * 60 * 60,
 )
-def benchmark(update_mode: str, runtime: str, sample_id: str) -> dict:
+def benchmark(
+    update_mode: str,
+    canonical_storage: str | None,
+    runtime: str,
+    sample_id: str,
+) -> dict:
     return run_delta_weight_update(
         WeightUpdateSpec(
             model_name="GLM-4.5-Air FP8",
             base_checkpoint_dir=BASE_CHECKPOINT_DIR,
             local_target_checkpoint_dir=LOCAL_TARGET_CHECKPOINT_DIR,
+            local_canonical_checkpoint_dir=LOCAL_CANONICAL_CHECKPOINT_DIR,
             server_args=SGLANG_SERVER_ARGS,
         ),
         source_dir=DELTA_SOURCE_DIR,
         target_version=1,
         update_mode=parse_update_mode(update_mode),
+        canonical_storage=parse_canonical_storage(canonical_storage),
         runtime=runtime,
         sample_id=sample_id,
     )
 
 
 @app.local_entrypoint()
-def main(update_mode: str = "cpu", sample_id: str = "1") -> None:
-    prepare_base.remote()
-    prepare_delta.remote()
-    benchmark.remote(parse_update_mode(update_mode), modal_runtime_label(), sample_id)
+def main(
+    update_mode: str = "disk",
+    canonical_storage: str | None = None,
+    sample_id: str = "1",
+    skip_preparation: bool = False,
+) -> None:
+    mode, storage = parse_update_destination(update_mode, canonical_storage)
+    if not skip_preparation:
+        prepare_base.remote()
+        prepare_delta.remote()
+    benchmark.remote(
+        mode,
+        storage,
+        modal_runtime_label(),
+        sample_id,
+    )
diff --git a/tools/profiling/glm5_2_nvfp4_delta_weight_update.py b/tools/profiling/glm5_2_nvfp4_delta_weight_update.py
index 7d278ada..94bd2092 100644
--- a/tools/profiling/glm5_2_nvfp4_delta_weight_update.py
+++ b/tools/profiling/glm5_2_nvfp4_delta_weight_update.py
@@ -3,7 +3,8 @@
 CPU destination:
 
     MODAL_FUNCTION_RUNTIME=runc uv run --extra modal modal run -d \
-      tools/profiling/glm5_2_nvfp4_delta_weight_update.py
+      tools/profiling/glm5_2_nvfp4_delta_weight_update.py \
+      --update-mode cpu --canonical-storage memory
 
 Disk destination:
 
@@ -25,6 +26,8 @@
 from tools.profiling._delta_weight_update import (
     WeightUpdateSpec,
     modal_runtime_label,
+    parse_canonical_storage,
+    parse_update_destination,
     parse_update_mode,
     run_delta_weight_update,
 )
@@ -37,19 +40,48 @@
 APP_NAME = "profile-glm5-2-nvfp4-delta-weight-update"
 EXPERIMENT = "glm5_2_nvfp4"
 DELTA_MOUNT = "/synthetic-delta"
-_PIPELINE_LAYER_ENDS = (14, 22, 30, 38, 46, 54, 62, 78)
+_TARGET_MODEL_LAYERS = 78
+
+
+def _pipeline_layer_ends() -> tuple[int, ...]:
+    stages = model.miles.pipeline_model_parallel_size
+    first = model.miles.decoder_first_pipeline_num_layers
+    last = model.miles.decoder_last_pipeline_num_layers
+    if stages < 2 or first is None or last is None:
+        raise ValueError("GLM-5.2 profiling requires explicit first/last PP stages")
+    middle_stages = stages - 2
+    remaining = _TARGET_MODEL_LAYERS - first - last
+    if middle_stages:
+        middle_layers, remainder = divmod(remaining, middle_stages)
+        if remainder:
+            raise ValueError("GLM-5.2 middle layers do not divide across PP stages")
+        layer_counts = (first, *(middle_layers for _ in range(middle_stages)), last)
+    else:
+        if remaining:
+            raise ValueError("GLM-5.2 first/last PP stages do not cover all layers")
+        layer_counts = (first, last)
+    ends = []
+    for count in layer_counts:
+        ends.append(count + (ends[-1] if ends else 0))
+    return tuple(ends)
+
+
+_PIPELINE_LAYER_ENDS = _pipeline_layer_ends()
 DELTA_SPEC = SyntheticDeltaSpec(
     checkpoint_format="nvfp4",
     quantized_value_density=0.00375,
     high_precision_value_density=0.01,
     output_shards=model.miles.pipeline_model_parallel_size,
-    output_shard_layout="miles-pp8-layer-placement-v1",
+    output_shard_layout=(
+        f"miles-pp{model.miles.pipeline_model_parallel_size}-layer-placement-v1"
+    ),
     # MTP is present in the immutable serving checkpoint but is not trained.
     immutable_prefixes=("model.layers.78.",),
 )
 DELTA_ID = f"glm5-2/{model.SOURCE_REVISION}/{synthetic_delta_profile_id(DELTA_SPEC)}"
 DELTA_SOURCE_DIR = f"{DELTA_MOUNT}/{DELTA_ID}"
 LOCAL_TARGET_CHECKPOINT_DIR = "/local-checkpoint/glm5-2-nvfp4/target"
+LOCAL_CANONICAL_CHECKPOINT_DIR = "/local-checkpoint/glm5-2-nvfp4/canonical"
 SGLANG_CACHE_PATH = "/root/.cache/sglang"
 
 app = modal.App(APP_NAME)
@@ -72,9 +104,12 @@
 prep_image = trainer_image.build_trainer_image(
     hf_cache_path=str(HF_CACHE_PATH),
     experiment=EXPERIMENT,
-    miles_repo_ref=model.MILES_REPO_REF,
     extra_pip_packages=model.TRAINER_EXTRA_PIP_PACKAGES,
     image_run_commands=model.TRAINER_IMAGE_RUN_COMMANDS,
+).add_local_dir(
+    str(Path(__file__).resolve().parents[1]),
+    remote_path="/root/tools",
+    ignore=["**/__pycache__", "**/*.pyc"],
 )
 serving_image = build_serving_image(
     hf_cache_path=str(HF_CACHE_PATH),
@@ -150,25 +185,44 @@ def prepare_delta() -> dict:
     },
     timeout=6 * 60 * 60,
 )
-def benchmark(update_mode: str, runtime: str, sample_id: str) -> dict:
+def benchmark(
+    update_mode: str,
+    canonical_storage: str | None,
+    runtime: str,
+    sample_id: str,
+) -> dict:
     return run_delta_weight_update(
         WeightUpdateSpec(
             model_name="GLM-5.2 mixed NVFP4/BF16",
             base_checkpoint_dir=str(model.ROLLOUT_CHECKPOINT_PATH),
             local_target_checkpoint_dir=LOCAL_TARGET_CHECKPOINT_DIR,
+            local_canonical_checkpoint_dir=LOCAL_CANONICAL_CHECKPOINT_DIR,
             server_args=model.SGLANG_SERVER_ARGS,
             tp_size=model.ROLLOUT_GPUS_PER_ENGINE,
         ),
         source_dir=DELTA_SOURCE_DIR,
         target_version=1,
         update_mode=parse_update_mode(update_mode),
+        canonical_storage=parse_canonical_storage(canonical_storage),
         runtime=runtime,
         sample_id=sample_id,
     )
 
 
 @app.local_entrypoint()
-def main(update_mode: str = "cpu", sample_id: str = "1") -> None:
-    prepare_base.remote()
-    prepare_delta.remote()
-    benchmark.remote(parse_update_mode(update_mode), modal_runtime_label(), sample_id)
+def main(
+    update_mode: str = "disk",
+    canonical_storage: str | None = None,
+    sample_id: str = "1",
+    skip_preparation: bool = False,
+) -> None:
+    mode, storage = parse_update_destination(update_mode, canonical_storage)
+    if not skip_preparation:
+        prepare_base.remote()
+        prepare_delta.remote()
+    benchmark.remote(
+        mode,
+        storage,
+        modal_runtime_label(),
+        sample_id,
+    )
diff --git a/tools/profiling/kimi_k2_6_nvfp4_delta_weight_update.py b/tools/profiling/kimi_k2_6_nvfp4_delta_weight_update.py
index 56ea6a9e..2dc1906d 100644
--- a/tools/profiling/kimi_k2_6_nvfp4_delta_weight_update.py
+++ b/tools/profiling/kimi_k2_6_nvfp4_delta_weight_update.py
@@ -1,12 +1,11 @@
 """Profile one Kimi K2.6 NVFP4 delta weight update on Modal.
 
-The entrypoint prepares the pinned serving checkpoint with the Miles NVFP4
-recipe, builds a standardized element-wise synthetic delta, and runs one
-verified update.
+The entrypoint downloads NVIDIA's pinned serving checkpoint, builds a
+standardized element-wise synthetic delta, and runs one verified update.
 
     MODAL_FUNCTION_RUNTIME=runc uv run --extra modal modal run -d \
       tools/profiling/kimi_k2_6_nvfp4_delta_weight_update.py \
-      --update-mode cpu
+      --update-mode cpu --canonical-storage disk
 """
 
 from __future__ import annotations
@@ -15,16 +14,20 @@
 
 import modal
 
-from cookbook.common.constants import CHECKPOINTS_PATH, HF_CACHE_PATH
+from cookbook.common.constants import HF_CACHE_PATH
 from cookbook.common.serving_image import build_serving_image
-from cookbook.miles_disagg import prep, trainer_image
-from cookbook.miles_disagg.configs import kimi_k2_6_nvfp4 as model
 from tools.profiling._delta_weight_update import (
     WeightUpdateSpec,
     modal_runtime_label,
+    parse_canonical_storage,
+    parse_update_destination,
     parse_update_mode,
     run_delta_weight_update,
 )
+from tools.profiling._hf_checkpoint import (
+    download_snapshot,
+    materialize_checkpoint_view,
+)
 from tools.profiling._synthetic_delta import (
     SyntheticDeltaSpec,
     prepare_standard_delta,
@@ -33,32 +36,37 @@
 
 APP_NAME = "profile-kimi-k2-6-nvfp4-delta-weight-update"
 EXPERIMENT = "kimi_k2_6_nvfp4"
+ROLLOUT_MODEL = "nvidia/Kimi-K2.6-NVFP4"
+ROLLOUT_REVISION = "2fd3a800dedd098b8327eb49e93ebc75f85da19f"
 DELTA_MOUNT = "/synthetic-delta"
 DELTA_SPEC = SyntheticDeltaSpec(
     checkpoint_format="nvfp4",
     quantized_value_density=0.003,
     high_precision_value_density=0.01,
     # Text-only RL leaves the vision encoder and projector fixed.
-    immutable_prefixes=("vision_tower.", "mm_projector."),
+    immutable_prefixes=("vision_tower.", "multi_modal_projector."),
 )
-DELTA_ID = f"kimi-k2-6/{model.SOURCE_REVISION}/{synthetic_delta_profile_id(DELTA_SPEC)}"
+DELTA_ID = f"kimi-k2-6/{ROLLOUT_REVISION}/{synthetic_delta_profile_id(DELTA_SPEC)}"
 DELTA_SOURCE_DIR = f"{DELTA_MOUNT}/{DELTA_ID}"
-BASE_CHECKPOINT_DIR = str(model.ROLLOUT_CHECKPOINT_PATH)
+HF_SNAPSHOT_DIR = (
+    f"{HF_CACHE_PATH}/models--nvidia--Kimi-K2.6-NVFP4/snapshots/{ROLLOUT_REVISION}"
+)
 LOCAL_CHECKPOINT_ROOT = "/local-checkpoint/kimi-k2-6-nvfp4"
+BASE_CHECKPOINT_DIR = f"{LOCAL_CHECKPOINT_ROOT}/base"
 LOCAL_TARGET_CHECKPOINT_DIR = f"{LOCAL_CHECKPOINT_ROOT}/target"
 LOCAL_CANONICAL_CHECKPOINT_DIR = f"{LOCAL_CHECKPOINT_ROOT}/canonical"
 SGLANG_CACHE_PATH = "/root/.cache/sglang"
 
 SGLANG_SERVER_ARGS = {
-    "--served-model-name": model.SOURCE_MODEL,
+    "--served-model-name": ROLLOUT_MODEL,
     "--load-format": "fastsafetensors",
     "--model-loader-extra-config": '{"enable_gds":false}',
     "--weight-loader-drop-cache-after-load": "",
-    "--cpu-weight-cache-canonical-checkpoint-dir": LOCAL_CANONICAL_CHECKPOINT_DIR,
     "--trust-remote-code": "",
     "--tool-call-parser": "kimi_k2",
     "--reasoning-parser": "kimi_k2",
     "--dist-timeout": "3600",
+    "--watchdog-timeout": "3600",
     "--kv-cache-dtype": "fp8_e4m3",
     "--attention-backend": "tokenspeed_mla",
     "--context-length": "32768",
@@ -77,18 +85,26 @@
 hf_cache_volume = modal.Volume.from_name(
     "huggingface-cache", create_if_missing=True, version=2
 )
-checkpoint_volume = modal.Volume.from_name(
-    "miles-checkpoints", create_if_missing=True, version=2
-)
 delta_volume = modal.Volume.from_name(
     "stitch-synthetic-deltas", create_if_missing=True, version=2
 )
 sglang_cache_volume = modal.Volume.from_name(
     "sglang-cache", create_if_missing=True, version=2
 )
-prep_image = trainer_image.build_trainer_image(
-    hf_cache_path=str(HF_CACHE_PATH),
-    experiment=EXPERIMENT,
+download_image = (
+    modal.Image.debian_slim(python_version="3.11")
+    .pip_install("huggingface_hub[hf_transfer]")
+    .env(
+        {
+            "HF_XET_HIGH_PERFORMANCE": "1",
+            "HF_HUB_ENABLE_HF_TRANSFER": "1",
+        }
+    )
+    .add_local_dir(
+        str(Path(__file__).resolve().parents[1]),
+        remote_path="/root/tools",
+        ignore=["**/__pycache__", "**/*.pyc"],
+    )
 )
 serving_image = build_serving_image(
     hf_cache_path=str(HF_CACHE_PATH),
@@ -101,18 +117,20 @@
 
 
 @app.function(
-    image=prep_image,
-    gpu=f"{model.modal.gpu}:1",
-    memory=model.modal.trainer_memory_mib,
-    volumes={
-        str(HF_CACHE_PATH): hf_cache_volume,
-        str(CHECKPOINTS_PATH): checkpoint_volume,
-    },
+    image=download_image,
+    cpu=32,
+    memory=(16 * 1024, 256 * 1024),
+    volumes={str(HF_CACHE_PATH): hf_cache_volume},
     secrets=[modal.Secret.from_name("huggingface-secret")],
     timeout=6 * 60 * 60,
 )
-def prepare_base() -> None:
-    prep.prepare_checkpoints(model, checkpoint_volume)
+def download_model() -> str:
+    return download_snapshot(
+        ROLLOUT_MODEL,
+        ROLLOUT_REVISION,
+        str(HF_CACHE_PATH),
+        commit=hf_cache_volume.commit,
+    )
 
 
 @app.function(
@@ -120,14 +138,14 @@ def prepare_base() -> None:
     cpu=64,
     memory=(64 * 1024, 512 * 1024),
     volumes={
-        str(CHECKPOINTS_PATH): checkpoint_volume.read_only(),
+        str(HF_CACHE_PATH): hf_cache_volume.read_only(),
         DELTA_MOUNT: delta_volume,
     },
     timeout=6 * 60 * 60,
 )
 def prepare_delta() -> dict:
     return prepare_standard_delta(
-        BASE_CHECKPOINT_DIR,
+        HF_SNAPSHOT_DIR,
         DELTA_SOURCE_DIR,
         spec=DELTA_SPEC,
         commit=delta_volume.commit,
@@ -142,30 +160,50 @@ def prepare_delta() -> dict:
     # Disk mode retains both the immutable base and a complete mutable target.
     ephemeral_disk=1_572_864,
     volumes={
-        str(CHECKPOINTS_PATH): checkpoint_volume.read_only(),
+        str(HF_CACHE_PATH): hf_cache_volume.read_only(),
         DELTA_MOUNT: delta_volume.read_only(),
         SGLANG_CACHE_PATH: sglang_cache_volume,
     },
     timeout=4 * 60 * 60,
 )
-def benchmark(update_mode: str, runtime: str, sample_id: str) -> dict:
+def benchmark(
+    update_mode: str,
+    canonical_storage: str | None,
+    runtime: str,
+    sample_id: str,
+) -> dict:
+    materialize_checkpoint_view(HF_SNAPSHOT_DIR, BASE_CHECKPOINT_DIR)
     return run_delta_weight_update(
         WeightUpdateSpec(
             model_name="Kimi K2.6 NVFP4",
             base_checkpoint_dir=BASE_CHECKPOINT_DIR,
             local_target_checkpoint_dir=LOCAL_TARGET_CHECKPOINT_DIR,
+            local_canonical_checkpoint_dir=LOCAL_CANONICAL_CHECKPOINT_DIR,
             server_args=SGLANG_SERVER_ARGS,
         ),
         source_dir=DELTA_SOURCE_DIR,
         target_version=1,
         update_mode=parse_update_mode(update_mode),
+        canonical_storage=parse_canonical_storage(canonical_storage),
         runtime=runtime,
         sample_id=sample_id,
     )
 
 
 @app.local_entrypoint()
-def main(update_mode: str = "cpu", sample_id: str = "1") -> None:
-    prepare_base.remote()
-    prepare_delta.remote()
-    benchmark.remote(parse_update_mode(update_mode), modal_runtime_label(), sample_id)
+def main(
+    update_mode: str = "disk",
+    canonical_storage: str | None = None,
+    sample_id: str = "1",
+    skip_preparation: bool = False,
+) -> None:
+    mode, storage = parse_update_destination(update_mode, canonical_storage)
+    if not skip_preparation:
+        download_model.remote()
+        prepare_delta.remote()
+    benchmark.remote(
+        mode,
+        storage,
+        modal_runtime_label(),
+        sample_id,
+    )
diff --git a/tools/profiling/kimi_k3_mxfp4_delta_weight_update.py b/tools/profiling/kimi_k3_mxfp4_delta_weight_update.py
index 3d64f803..8468c3bf 100644
--- a/tools/profiling/kimi_k3_mxfp4_delta_weight_update.py
+++ b/tools/profiling/kimi_k3_mxfp4_delta_weight_update.py
@@ -1,13 +1,14 @@
 """Download Kimi K3 and validate one complete MXFP4 delta update on Modal.
 
-Disk destination (the Kimi K3 config's declared update mode):
+Disk destination:
 
-    MODAL_FUNCTION_RUNTIME=runc uv run --extra modal modal run -d \
+    uv run --extra modal modal run -d \
       tools/profiling/kimi_k3_mxfp4_delta_weight_update.py
 
-CPU destination with the canonical checkpoint on local storage:
+CPU destination with the canonical checkpoint on local storage (the recipe's
+declared update mode):
 
-    MODAL_FUNCTION_RUNTIME=runc uv run --extra modal modal run -d \
+    uv run --extra modal modal run -d \
       tools/profiling/kimi_k3_mxfp4_delta_weight_update.py \
       --update-mode cpu --canonical-storage disk
 
@@ -19,8 +20,6 @@
 
 from __future__ import annotations
 
-import json
-import shutil
 from pathlib import Path
 
 import modal
@@ -30,10 +29,20 @@
 from tools.profiling._delta_weight_update import (
     WeightUpdateSpec,
     modal_runtime_label,
+    parse_canonical_storage,
+    parse_update_destination,
     parse_update_mode,
     run_delta_weight_update,
 )
-from tools.profiling._synthetic_delta import write_full_coverage_delta
+from tools.profiling._hf_checkpoint import (
+    download_snapshot,
+    materialize_checkpoint_view,
+)
+from tools.profiling._synthetic_delta import (
+    SyntheticDeltaSpec,
+    prepare_standard_delta,
+    synthetic_delta_profile_id,
+)
 
 APP_NAME = "profile-kimi-k3-mxfp4-delta-weight-update"
 EXPERIMENT = "kimi_k3_mxfp4"
@@ -43,12 +52,19 @@
     f"{model.ROLLOUT_SOURCE_REVISION}"
 )
 DELTA_MOUNT = "/synthetic-delta"
-DELTA_ID = f"kimi-k3/{model.ROLLOUT_SOURCE_REVISION}/full-coverage-v3"
+DELTA_SPEC = SyntheticDeltaSpec(
+    checkpoint_format="mxfp4",
+    quantized_value_density=0.003,
+    high_precision_value_density=0.01,
+    # Text-only RL leaves the vision encoder and projector fixed.
+    immutable_prefixes=("vision_tower.", "mm_projector."),
+)
+DELTA_ID = (
+    f"kimi-k3/{model.ROLLOUT_SOURCE_REVISION}/{synthetic_delta_profile_id(DELTA_SPEC)}"
+)
 DELTA_SOURCE_DIR = f"{DELTA_MOUNT}/{DELTA_ID}"
 BASE_CHECKPOINT_DIR = "/local-checkpoint/kimi-k3-mxfp4/base"
 LOCAL_TARGET_CHECKPOINT_DIR = "/local-checkpoint/kimi-k3-mxfp4/target"
-# CPU-mode-only overlay: the config is a clean disk config, so the profiler
-# injects the cpu-weight-cache args when profiling the cpu destination.
 CPU_CACHE_GROUP_GB = "16"
 CANONICAL_CHECKPOINT_DIR = "/local-checkpoint/kimi-k3-mxfp4/canonical"
 SGLANG_CACHE_PATH = "/root/.cache/sglang"
@@ -95,7 +111,6 @@
     hf_cache_path=HF_CACHE_PATH,
     experiment=EXPERIMENT,
     extra_env=getattr(model, "SGLANG_SERVER_ENV", None),
-    runtime=model.SGLANG_RUNTIME,
 ).add_local_dir(
     str(Path(__file__).resolve().parents[1]),
     remote_path="/root/tools",
@@ -112,38 +127,12 @@
     timeout=6 * 60 * 60,
 )
 def download_model() -> str:
-    from huggingface_hub import HfApi, snapshot_download
-
-    files = HfApi().list_repo_files(
-        repo_id=model.ROLLOUT_SOURCE_MODEL,
-        revision=model.ROLLOUT_SOURCE_REVISION,
-    )
-    weights = sorted(path for path in files if path.endswith(".safetensors"))
-    metadata = sorted(set(files) - set(weights))
-    batches = [metadata]
-    batches.extend(weights[offset : offset + 4] for offset in range(0, len(weights), 4))
-    for index, batch in enumerate(batches, start=1):
-        snapshot_download(
-            repo_id=model.ROLLOUT_SOURCE_MODEL,
-            revision=model.ROLLOUT_SOURCE_REVISION,
-            cache_dir=HF_CACHE_PATH,
-            allow_patterns=batch,
-            max_workers=4,
-        )
-        hf_cache_volume.commit()
-        print(f"Committed checkpoint download batch {index}/{len(batches)}")
-
-    path = snapshot_download(
-        repo_id=model.ROLLOUT_SOURCE_MODEL,
-        revision=model.ROLLOUT_SOURCE_REVISION,
-        cache_dir=HF_CACHE_PATH,
-        local_files_only=True,
-    )
-    print(
-        f"Downloaded {model.ROLLOUT_SOURCE_MODEL}"
-        f"@{model.ROLLOUT_SOURCE_REVISION} to {path}"
+    return download_snapshot(
+        model.ROLLOUT_SOURCE_MODEL,
+        model.ROLLOUT_SOURCE_REVISION,
+        HF_CACHE_PATH,
+        commit=hf_cache_volume.commit,
     )
-    return path
 
 
 @app.function(
@@ -157,43 +146,12 @@ def download_model() -> str:
     timeout=6 * 60 * 60,
 )
 def prepare_delta() -> dict:
-    metadata_path = Path(DELTA_SOURCE_DIR) / "controlled_delta.json"
-    index_path = (
-        Path(DELTA_SOURCE_DIR) / "weight_v000001" / "model.safetensors.index.json"
-    )
-    if metadata_path.is_file() and index_path.is_file():
-        result = json.loads(metadata_path.read_text())
-        print(f"Reusing synthetic delta at {DELTA_SOURCE_DIR}")
-        return result
-
-    shutil.rmtree(DELTA_SOURCE_DIR, ignore_errors=True)
-    Path(DELTA_SOURCE_DIR).mkdir(parents=True)
-    result = write_full_coverage_delta(
+    return prepare_standard_delta(
         HF_SNAPSHOT_DIR,
         DELTA_SOURCE_DIR,
+        spec=DELTA_SPEC,
+        commit=delta_volume.commit,
     )
-    metadata_path.write_text(json.dumps(result, sort_keys=True))
-    delta_volume.commit()
-    print(f"Committed synthetic delta at {DELTA_SOURCE_DIR}")
-    return result
-
-
-def _materialize_checkpoint_view() -> None:
-    """Give trusted remote code real sibling files without copying model weights."""
-
-    source = Path(HF_SNAPSHOT_DIR)
-    target = Path(BASE_CHECKPOINT_DIR)
-    shutil.rmtree(target, ignore_errors=True)
-    target.mkdir(parents=True)
-    for path in source.rglob("*"):
-        destination = target / path.relative_to(source)
-        if path.is_dir():
-            destination.mkdir(parents=True, exist_ok=True)
-        elif path.suffix == ".safetensors":
-            destination.symlink_to(path.resolve())
-        else:
-            destination.parent.mkdir(parents=True, exist_ok=True)
-            shutil.copy2(path, destination, follow_symlinks=True)
 
 
 @app.function(
@@ -215,31 +173,22 @@ def benchmark(
     runtime: str,
     sample_id: str,
 ) -> dict:
-    _materialize_checkpoint_view()
+    materialize_checkpoint_view(HF_SNAPSHOT_DIR, BASE_CHECKPOINT_DIR)
     server_args = dict(model.SGLANG_SERVER_ARGS)
-    if update_mode == "cpu":
-        server_args["--enable-cpu-weight-cache"] = ""
-        server_args["--cpu-weight-cache-max-compile-group-gb"] = CPU_CACHE_GROUP_GB
-        storage = canonical_storage if canonical_storage is not None else "disk"
-        if storage == "disk":
-            server_args["--cpu-weight-cache-canonical-checkpoint-dir"] = (
-                CANONICAL_CHECKPOINT_DIR
-            )
-        elif storage != "memory":
-            raise ValueError("canonical_storage must be 'memory' or 'disk'")
-    elif canonical_storage is not None:
-        raise ValueError("--canonical-storage applies only with --update-mode cpu")
+    server_args["--cpu-weight-cache-max-compile-group-gb"] = CPU_CACHE_GROUP_GB
     return run_delta_weight_update(
         WeightUpdateSpec(
             model_name="Kimi K3 MXFP4",
             base_checkpoint_dir=BASE_CHECKPOINT_DIR,
             local_target_checkpoint_dir=LOCAL_TARGET_CHECKPOINT_DIR,
+            local_canonical_checkpoint_dir=CANONICAL_CHECKPOINT_DIR,
             server_args=server_args,
             tp_size=model.ROLLOUT_NUM_GPUS_PER_ENGINE,
         ),
         source_dir=DELTA_SOURCE_DIR,
         target_version=1,
         update_mode=parse_update_mode(update_mode),
+        canonical_storage=parse_canonical_storage(canonical_storage),
         runtime=runtime,
         sample_id=sample_id,
     )
@@ -250,15 +199,18 @@ def main(
     update_mode: str = "disk",
     canonical_storage: str | None = None,
     sample_id: str = "1",
+    skip_preparation: bool = False,
 ) -> None:
-    parsed_mode = parse_update_mode(update_mode)
-    if parsed_mode == "disk" and canonical_storage is not None:
-        raise ValueError("--canonical-storage applies only with --update-mode cpu")
-    download_model.remote()
-    prepare_delta.remote()
+    parsed_mode, parsed_storage = parse_update_destination(
+        update_mode,
+        canonical_storage,
+    )
+    if not skip_preparation:
+        download_model.remote()
+        prepare_delta.remote()
     benchmark.remote(
         parsed_mode,
-        canonical_storage,
+        parsed_storage,
         modal_runtime_label(),
         sample_id,
     )

From 0b4655671e56c9d431282d838607e2e89aeec3a5 Mon Sep 17 00:00:00 2001
From: Nan <nan@modal.com>
Date: Mon, 10 Aug 2026 16:19:04 +0000
Subject: [PATCH 3/5] [cookbook] Update the pinned Miles runtime

---
 cookbook/miles_disagg/MILES_FORK.md           | 98 ++++++++-----------
 cookbook/miles_disagg/app.py                  | 17 ++--
 cookbook/miles_disagg/config.py               | 12 +--
 .../miles_disagg/configs/glm45_air_bf16.py    |  2 +-
 .../miles_disagg/configs/glm45_air_fp8.py     |  2 +-
 .../configs/glm47_flash_swebench_pro.py       | 27 +++--
 cookbook/miles_disagg/configs/glm5_2_nvfp4.py | 14 +--
 .../configs/kimi_k25_2layer_nvfp4.py          |  2 +-
 .../miles_disagg/configs/kimi_k2_6_nvfp4.py   |  2 +-
 .../miles_disagg/configs/moonlight_nvfp4.py   |  2 +-
 .../configs/qwen3_30b_a3b_nvfp4_46.py         |  2 +-
 cookbook/miles_disagg/prep.py                 | 42 +++++---
 cookbook/miles_disagg/trainer_image.py        |  2 +-
 13 files changed, 118 insertions(+), 106 deletions(-)

diff --git a/cookbook/miles_disagg/MILES_FORK.md b/cookbook/miles_disagg/MILES_FORK.md
index a46df75d..7fc83d40 100644
--- a/cookbook/miles_disagg/MILES_FORK.md
+++ b/cookbook/miles_disagg/MILES_FORK.md
@@ -1,60 +1,49 @@
-# Miles forks
+# Miles fork
 
-Stitch’s Miles recipes use a dated trainer image and an immutable Miles commit.
-[`trainer_image.py`](trainer_image.py) defines the default:
+Stitch’s Miles recipes install an immutable Miles revision over a dated trainer
+image. [`trainer_image.py`](trainer_image.py) defines the shared default:
 
 ```python
 MILES_IMAGE_TAG = "radixark/miles:dev-202607290235"
 MILES_REPO_URL = "https://github.com/modal-projects/miles.git"
-MILES_REPO_REF = "1eb7520018446cb94b7406715f66dff1a271b53b"
+MILES_REPO_REF = "f83a68b296ca6f42d04e77e2b42c070c1ab02c70"
 ```
 
-An experiment config may override `MILES_REPO_REF` without changing other
-recipes. GLM-5.2 does this because it runs on Miles’ fully-async SWE branch.
+The pin is `modal-projects/miles:stitch-miles-fully-async-swe`. It contains a
+Miles fully-async stack based on upstream main at `6bc45ad39`, followed by the
+Stitch external-fleet integration.
 
-## Standard weight-sync branch
+## Miles fully-async stack
 
-The branch `stitch-weight-sync-v0516` is upstream `radixark/miles` main at
-`52403d0e3` plus six commits:
+These commits are useful independently of Stitch:
 
 | Commit | Responsibility |
 | --- | --- |
-| `7995ec27c` | Route generation through one opaque fleet endpoint, publish versions without per-engine handles, gate individual requests through a generic hook, and use a finite rollout read timeout. |
-| `4572f97a6` | Emit canonical quantized checkpoint layouts for disk deltas, support GLM-Air channel FP8 and plain-language-model Kimi names, and preserve SGLang runtime layouts for P2P/broadcast. |
-| `cf56d32b0` | Encode zero-dimensional quantization scalars by flattening before their byte view. |
-| `748f1724a` | Use SGLang’s staged checkpoint API for Miles-managed disk-delta engines. |
-| `d814b87c3` | Replace the unused delta progress bar with an explicit baseline-snapshot log. |
-| `1eb752001` | Honor ModelOpt glob exclusions when exporting NVFP4 weights. |
-
-## GLM-5.2 fully-async SWE branch
-
-The GLM-5.2 and GLM-4.7 fully-async SWE recipes pin
-`stitch-miles-fully-async-swe` at `b1020b596`. The Stitch branch is stacked on
-`modal/feat/modal-swe-fully-async` at `d0c11a412`.
-
-The fully-async branch owns behavior that is useful without Stitch:
-
-| Commit | Responsibility |
-| --- | --- |
-| `de10d68d1` | Format the fully-async rollout files. |
-| `ffcad9557` | Match GLM router tensors to the canonical checkpoint dtypes. |
-| `c0661aa6a` | Validate canonical disk-delta layouts and encode scalar tensors safely. |
-| `8857d8114` | Shard routing replay by trainer topology with a lossless compact wire dtype. |
-| `8e140dbd9` | Apply the configured log-prob token budget when routing replay is enabled. |
-| `791ef9593` | Bound Modal agent submissions before they enter Ray actor mailboxes. |
-| `d0c11a412` | Format the added fully-async code. |
-
-The Stitch branch adds only the external-fleet integration:
+| `ad0411c85` | Assemble additional routing-replay rows during in-place updates. |
+| `35c78a78b` | Preserve in-flight requests while changing their weight version. |
+| `6d638ba8a` | Make distributed service startup and teardown deterministic. |
+| `9be219f58` | Report generation and queue staleness separately. |
+| `63da214e1` | Distinguish explicit checkpoint resume from a fresh run. |
+| `bd0216740` | Load optional P2P weight-sync dependencies lazily. |
+| `717802e21` | Validate canonical disk-delta tensor layouts. |
+| `13dd38252` | Report trainer-rollout policy drift accurately. |
+| `a658a505c` | Make session endpoint affinity and lifecycle explicit. |
+| `6dbc5ab45` | Add the Modal Sandbox v2 transport for SWE rollouts. |
+| `080e58555` | Honor rollout-only LoRA configuration. |
+| `4e98074fa` | Add and validate the mini-SWE rollout adapter. |
+| `113271f5e` | Add CISPO policy optimization. |
+| `d688df471` | Replay exact truncated-sampling support during training. |
+| `53c402487` | Compact routing-replay expert IDs for transport. |
+| `719e2844d` | Suppress per-request session access logs. |
+| `45a2c5b25` | Reject aborted generations at the session protocol boundary. |
+
+## Stitch integration
 
 | Commit | Responsibility |
 | --- | --- |
-| `2fa28cde4` | Route session rollouts through an external fleet with request gating, version constraints, and finite timeouts. |
-| `7bfb4a69a` | Publish disk deltas without Miles-managed rollout-engine handles. |
-| `23b262bca` | Overlap the external fleet's initial delta snapshot with the first rollout. |
-| `b1020b596` | Sort imports in the external-fleet adapter. |
-
-This branch is additive and only backs fully-async SWE experiments. It does not
-replace the standard recipe pin.
+| `69244e487` | Route fully-async generation through an external fleet with version constraints and finite timeouts. |
+| `b54a23588` | Publish disk deltas without Miles-managed rollout-engine handles. |
+| `f83a68b29` | Overlap external weight updates with active rollout. |
 
 The dated image supplies Megatron-LM, TransformerEngine, CUDA, and other
 compiled dependencies. Miles is installed over it with `--no-deps`, so the
@@ -80,16 +69,15 @@ encoding is dtype- and quantization-agnostic; model-specific converters are
 responsible only for producing the same tensor names, dtypes, shapes, and byte
 layouts as the base checkpoint.
 
-## Re-porting
-
-When updating a Miles pin:
-
-1. start from the intended upstream branch;
-2. reapply only the responsibilities still missing upstream;
-3. keep canonical checkpoint layout changes scoped to `disk-delta`;
-4. use a dated Miles image whose Megatron and TransformerEngine match that
-   upstream revision;
-5. run Miles pre-commit plus the focused GPU tests for export, endpoint routing,
-   and staged engine requests;
-6. run a multi-step rollout test with a replica joining mid-run; and
-7. update the immutable repository SHA and this file.
+## Updating the pin
+
+1. Start from the intended upstream Miles revision.
+2. Reapply only behavior still missing upstream, keeping general Miles changes
+   below the Stitch-specific integration commits.
+3. Keep canonical checkpoint layout changes scoped to `disk-delta`.
+4. Use a dated image whose Megatron-LM and TransformerEngine match the target
+   Miles revision.
+5. Run Miles pre-commit and the focused tests for export, endpoint routing, and
+   staged engine requests.
+6. Run a multi-step rollout test with a replica joining mid-run.
+7. Update the immutable SHA and this file.
diff --git a/cookbook/miles_disagg/app.py b/cookbook/miles_disagg/app.py
index 41bce656..0d5e762b 100644
--- a/cookbook/miles_disagg/app.py
+++ b/cookbook/miles_disagg/app.py
@@ -393,14 +393,15 @@ def train(self, payload: dict) -> None:
 
 def _build_train_cmd(cfg: MilesConfig) -> str:
     train_script = f"{MILES_ROOT}/{'train_async.py' if cfg.async_mode else 'train.py'}"
-    model_script = cfg.miles_model_script
-    if model_script:
-        inner = (
-            f"source {MILES_ROOT}/{model_script} && "
-            f"python3 {train_script} ${{MODEL_ARGS[@]}} {shlex.join(cfg.cli_args())}"
-        )
-        return f"bash -c {shlex.quote(inner)}"
-    return f"python3 {train_script} {shlex.join(cfg.cli_args())}"
+    args = cfg.cli_args()
+    if cfg.megatron_model_type:
+        from miles.utils.external_utils.model_args_utils import load_model_args
+
+        args = [
+            *shlex.split(load_model_args(cfg.megatron_model_type)),
+            *args,
+        ]
+    return shlex.join(["python3", train_script, *args])
 
 
 # ── Entrypoints (preparation lives in a separate app: cookbook.miles_disagg.prep_app) ──
diff --git a/cookbook/miles_disagg/config.py b/cookbook/miles_disagg/config.py
index b0457275..2b2a3288 100644
--- a/cookbook/miles_disagg/config.py
+++ b/cookbook/miles_disagg/config.py
@@ -2,7 +2,7 @@
 
 Every public, non-callable attribute becomes a miles CLI arg via ``cli_args`` (miles
 wraps Megatron's parser, so Megatron args pass straight through); ``environment`` /
-``async_mode`` / ``miles_model_script`` are launcher instructions, not CLI args. The
+``async_mode`` / ``megatron_model_type`` are launcher instructions, not CLI args. The
 Modal-infra half of an experiment is ``common.config.ModalConfig``.
 """
 
@@ -10,7 +10,7 @@
 
 from typing import Any
 
-_MILES_SKIP = {"environment", "async_mode", "miles_model_script"}
+_MILES_SKIP = {"environment", "async_mode", "megatron_model_type"}
 # Fields miles reads as YAML files; inline dicts are materialized before launch.
 # (te_precision_config_file is handled separately in app.py — it needs an identical
 # node-local path on every Ray actor, not a per-launch tmpdir.)
@@ -23,9 +23,7 @@ class MilesConfig:
 
     environment: dict = {}
     async_mode: bool = False  # True -> train_async.py
-    miles_model_script: str = (
-        ""  # shell script (relative to the miles root) defining MODEL_ARGS
-    )
+    megatron_model_type: str = ""
 
     def __init__(self, **kwargs: Any) -> None:
         self.environment = dict(
@@ -89,7 +87,7 @@ def to_payload(self) -> dict[str, Any]:
             "fields": self._fields(),
             "environment": dict(self.environment),
             "async_mode": self.async_mode,
-            "miles_model_script": self.miles_model_script,
+            "megatron_model_type": self.megatron_model_type,
         }
 
     @classmethod
@@ -97,5 +95,5 @@ def from_payload(cls, payload: dict[str, Any]) -> MilesConfig:
         cfg = cls(**payload["fields"])
         cfg.environment = dict(payload["environment"])
         cfg.async_mode = payload["async_mode"]
-        cfg.miles_model_script = payload["miles_model_script"]
+        cfg.megatron_model_type = payload["megatron_model_type"]
         return cfg
diff --git a/cookbook/miles_disagg/configs/glm45_air_bf16.py b/cookbook/miles_disagg/configs/glm45_air_bf16.py
index 036f3d7d..6536f066 100644
--- a/cookbook/miles_disagg/configs/glm45_air_bf16.py
+++ b/cookbook/miles_disagg/configs/glm45_air_bf16.py
@@ -67,7 +67,7 @@
 
 
 class _Miles(MilesConfig):
-    miles_model_script = "scripts/models/glm4.5-106B-A12B.sh"
+    megatron_model_type = "glm4.5-106B-A12B"
 
     hf_checkpoint = str(ROLLOUT_CHECKPOINT_PATH)
     ref_load = str(TORCH_DIST_CHECKPOINT_PATH)
diff --git a/cookbook/miles_disagg/configs/glm45_air_fp8.py b/cookbook/miles_disagg/configs/glm45_air_fp8.py
index 72a7dfa3..948080ce 100644
--- a/cookbook/miles_disagg/configs/glm45_air_fp8.py
+++ b/cookbook/miles_disagg/configs/glm45_air_fp8.py
@@ -63,7 +63,7 @@
 
 
 class _Miles(MilesConfig):
-    miles_model_script = "scripts/models/glm4.5-106B-A12B.sh"
+    megatron_model_type = "glm4.5-106B-A12B"
 
     hf_checkpoint = str(ROLLOUT_CHECKPOINT_PATH)
     ref_load = str(TORCH_DIST_CHECKPOINT_PATH)
diff --git a/cookbook/miles_disagg/configs/glm47_flash_swebench_pro.py b/cookbook/miles_disagg/configs/glm47_flash_swebench_pro.py
index fc2fa765..4d144a77 100644
--- a/cookbook/miles_disagg/configs/glm47_flash_swebench_pro.py
+++ b/cookbook/miles_disagg/configs/glm47_flash_swebench_pro.py
@@ -9,7 +9,6 @@
 
 APP_NAME = "stitch-glm47-flash-swebench-pro"
 EXPERIMENT_VOLUME_NAME = "stitch-miles-glm47-flash-swebench-pro"
-MILES_REPO_REF = "b1020b5961657ef1bb8c9f56bda49bc12899fa57"
 LOCAL_CHECKPOINT_PATH = None
 
 SOURCE_MODEL = "zai-org/GLM-4.7-Flash"
@@ -34,12 +33,12 @@
     "/root/cookbook/miles_disagg/patches/megatron-hdo-dp-reshardable-step.patch",
 ]
 
-TRAINER_NODES = 2
+TRAINER_NODES = 4
 GPUS_PER_TRAINER_NODE = 8
 ROLLOUT_GPUS_PER_ENGINE = 1
-ROLLOUT_MIN_CONTAINERS = 20
 ROLLOUT_INPUTS_PER_ENGINE = 20
-ROLLOUT_CONCURRENT_SAMPLES = 544
+ROLLOUT_CONCURRENT_SAMPLES = 640  # saturates the 40-worker rollout harness
+ROLLOUT_MIN_CONTAINERS = 48  # 48 rollout GPUs at one GPU per engine
 ROLLOUT_MAX_RUNNING_REQUESTS = 28
 ROLLOUT_MAX_QUEUED_REQUESTS = 4  # backpressure
 MAX_SEQ_LEN = 65_536
@@ -70,16 +69,17 @@
     "--decode-log-interval": "1000",
     "--log-level-http": "warning",
     "--enable-return-routed-experts": "",
+    "--sampling-mask-max-tokens": "8192",
 }
 
 modal = ModalConfig(
     gpu="H200",
     rollout_gpu="H200",
     cloud="aws",
-    trainer_memory_mib=(1024 * 1024, 2 * 1024 * 1024),
+    trainer_memory_mib=(1024 * 1024, 3 * 1024 * 1024),
     rollout_memory_mib=(512 * 1024, 1024 * 1024),
     rollout_min_containers=ROLLOUT_MIN_CONTAINERS,
-    rollout_max_containers=None,
+    rollout_max_containers=ROLLOUT_MIN_CONTAINERS,
     rollout_target_inputs=ROLLOUT_INPUTS_PER_ENGINE,
     rollout_ephemeral_disk_mib=524_288,
     trainer_ephemeral_disk_mib=1_048_576,
@@ -99,7 +99,7 @@
 
 
 class _Miles(MilesConfig):
-    miles_model_script = "scripts/models/glm4.7-flash.sh"
+    megatron_model_type = "glm4.7-flash"
     async_mode = True
 
     hf_checkpoint = str(ROLLOUT_CHECKPOINT_PATH)
@@ -142,7 +142,8 @@ class _Miles(MilesConfig):
     balance_data = True
 
     fully_async = True
-    rollout_sample_completion_backfill = True
+    pause_generation_mode = "in_place"
+    rollout_submission_granularity = "sample"
     custom_rollout_log_function_path = "modal_swe_metrics.log_rollout_data"
     custom_generate_function_path = (
         "miles.rollout.generate_hub.agentic_tool_call.generate"
@@ -150,10 +151,11 @@ class _Miles(MilesConfig):
     custom_agent_function_path = "modal_swe_agent_function.run"
     custom_rm_path = "modal_swe_agent_function.reward_func"
     tito_model = "glm47"
-    use_session_server = True
+    use_session_server = "v2"
+    session_sample_picker_path = "modal_swe_agent_function.pick_latest_leaf"
+    session_sample_postprocessor_path = "modal_swe_agent_function.postprocess_samples"
     session_server_port = [30000, 30064]
     session_server_startup_timeout_seconds = 180
-    tito_session_mismatch_sample_rate = 0.0625
 
     num_rollout = 500
     save_interval = 20
@@ -163,10 +165,15 @@ class _Miles(MilesConfig):
     global_batch_size = 256
     rollout_temperature = 1.0
     rollout_top_p = 0.95
+    rollout_top_k = 8192
     rollout_max_response_len = 8192
     max_seq_len = MAX_SEQ_LEN
     max_weight_staleness = 6
     async_max_concurrent_samples = ROLLOUT_CONCURRENT_SAMPLES
+    async_data_buffer_capacity_factor = (
+        0.5  # at most half a batch of completed R3 state
+    )
+    async_unused_samples_handler = "drop"
     eval_interval = None
 
     use_rollout_routing_replay = True
diff --git a/cookbook/miles_disagg/configs/glm5_2_nvfp4.py b/cookbook/miles_disagg/configs/glm5_2_nvfp4.py
index 43bdc2b1..0a2a4673 100644
--- a/cookbook/miles_disagg/configs/glm5_2_nvfp4.py
+++ b/cookbook/miles_disagg/configs/glm5_2_nvfp4.py
@@ -9,9 +9,6 @@
 
 APP_NAME = "stitch-glm5-2-nvfp4"
 EXPERIMENT_VOLUME_NAME = "stitch-miles-glm5-2-nvfp4"
-# This experiment layers Stitch integration onto Miles' fully-async SWE runtime.
-# Other cookbook experiments keep the standard Miles pin in trainer_image.py.
-MILES_REPO_REF = "b1020b5961657ef1bb8c9f56bda49bc12899fa57"
 LOCAL_CHECKPOINT_PATH = None
 TRAINER_EXTRA_PIP_PACKAGES = (
     "harbor[modal,huggingface]==0.20.0",
@@ -193,7 +190,7 @@
 
 
 class _Miles(MilesConfig):
-    miles_model_script = "scripts/models/glm5.2-744B-A40B.sh"
+    megatron_model_type = "glm5.2-744B-A40B"
     async_mode = True
 
     hf_checkpoint = str(ROLLOUT_CHECKPOINT_PATH)
@@ -284,7 +281,7 @@ class _Miles(MilesConfig):
     balance_data = True
 
     fully_async = True
-    rollout_sample_completion_backfill = True
+    rollout_submission_granularity = "sample"
     custom_rollout_log_function_path = "modal_swe_metrics.log_rollout_data"
     custom_generate_function_path = (
         "miles.rollout.generate_hub.agentic_tool_call.generate"
@@ -292,10 +289,11 @@ class _Miles(MilesConfig):
     custom_agent_function_path = "modal_swe_agent_function.run"
     custom_rm_path = "modal_swe_agent_function.reward_func"
     tito_model = "glm47"
-    use_session_server = True
+    use_session_server = "v2"
+    session_sample_picker_path = "modal_swe_agent_function.pick_latest_leaf"
+    session_sample_postprocessor_path = "modal_swe_agent_function.postprocess_samples"
     session_server_port = [30000, 30064]
     session_server_startup_timeout_seconds = 180
-    tito_session_mismatch_sample_rate = 0.0625
 
     num_rollout = 500
     save_interval = 10
@@ -309,6 +307,8 @@ class _Miles(MilesConfig):
     use_dynamic_global_batch_size = True
     max_weight_staleness = 6
     async_max_concurrent_samples = ROLLOUT_CONCURRENT_SAMPLES
+    async_data_buffer_capacity_factor = 3.0
+    async_unused_samples_handler = "drop"
     eval_interval = None
 
     use_rollout_routing_replay = True
diff --git a/cookbook/miles_disagg/configs/kimi_k25_2layer_nvfp4.py b/cookbook/miles_disagg/configs/kimi_k25_2layer_nvfp4.py
index f5e51c87..1b5a46ae 100644
--- a/cookbook/miles_disagg/configs/kimi_k25_2layer_nvfp4.py
+++ b/cookbook/miles_disagg/configs/kimi_k25_2layer_nvfp4.py
@@ -59,7 +59,7 @@
 
 
 class _Miles(MilesConfig):
-    miles_model_script = "scripts/models/kimi-k25_2layer.sh"
+    megatron_model_type = "kimi-k25_2layer"
 
     hf_checkpoint = str(ROLLOUT_CHECKPOINT_PATH)
     ref_load = str(TORCH_DIST_CHECKPOINT_PATH)
diff --git a/cookbook/miles_disagg/configs/kimi_k2_6_nvfp4.py b/cookbook/miles_disagg/configs/kimi_k2_6_nvfp4.py
index 06e11ce8..4f24bcaa 100644
--- a/cookbook/miles_disagg/configs/kimi_k2_6_nvfp4.py
+++ b/cookbook/miles_disagg/configs/kimi_k2_6_nvfp4.py
@@ -74,7 +74,7 @@
 
 class _Miles(MilesConfig):
     # Arch comes from the model script (shared with Kimi-K2-Thinking; K2.6 matches).
-    miles_model_script = "scripts/models/kimi-k2-thinking.sh"
+    megatron_model_type = "kimi-k2-thinking"
 
     hf_checkpoint = str(ROLLOUT_CHECKPOINT_PATH)
     ref_load = str(TORCH_DIST_CHECKPOINT_PATH)
diff --git a/cookbook/miles_disagg/configs/moonlight_nvfp4.py b/cookbook/miles_disagg/configs/moonlight_nvfp4.py
index 4676a7ab..3617562b 100644
--- a/cookbook/miles_disagg/configs/moonlight_nvfp4.py
+++ b/cookbook/miles_disagg/configs/moonlight_nvfp4.py
@@ -58,7 +58,7 @@
 
 class _Miles(MilesConfig):
     # Arch comes from the model script; do NOT inline arch attrs here.
-    miles_model_script = "scripts/models/moonlight.sh"
+    megatron_model_type = "moonlight"
 
     hf_checkpoint = str(ROLLOUT_CHECKPOINT_PATH)
     ref_load = str(BF16_CHECKPOINT_PATH)
diff --git a/cookbook/miles_disagg/configs/qwen3_30b_a3b_nvfp4_46.py b/cookbook/miles_disagg/configs/qwen3_30b_a3b_nvfp4_46.py
index 92f00fc5..fd5e27b1 100644
--- a/cookbook/miles_disagg/configs/qwen3_30b_a3b_nvfp4_46.py
+++ b/cookbook/miles_disagg/configs/qwen3_30b_a3b_nvfp4_46.py
@@ -130,7 +130,7 @@
 
 
 class _Miles(MilesConfig):
-    miles_model_script = "scripts/models/qwen3-30B-A3B.sh"
+    megatron_model_type = "qwen3-30B-A3B"
 
     # Bridge mode: ref_load is the bf16 HF masters directly (no torch_dist prep).
     hf_checkpoint = str(ROLLOUT_CHECKPOINT_PATH)
diff --git a/cookbook/miles_disagg/prep.py b/cookbook/miles_disagg/prep.py
index 40d04634..ac3ea5da 100644
--- a/cookbook/miles_disagg/prep.py
+++ b/cookbook/miles_disagg/prep.py
@@ -7,6 +7,7 @@
 
 import json
 import os
+import shlex
 import shutil
 import subprocess
 import threading
@@ -171,8 +172,11 @@ def prepare_torch_dist(exp, checkpoint_volume, *, rank: int, master_addr: str) -
     ):
         print(f"reusing existing torch_dist {torch_dist_dir}")
         return
-    if not exp.miles.miles_model_script:
-        raise SystemExit("prepare_torch_dist requires miles_model_script (MODEL_ARGS)")
+    if not exp.miles.megatron_model_type:
+        raise SystemExit("prepare_torch_dist requires megatron_model_type")
+    from miles.utils.external_utils.model_args_utils import load_model_args
+
+    model_args = shlex.split(load_model_args(exp.miles.megatron_model_type))
     nodes = exp.modal.torch_dist_prep_nodes
     use_wrapper = nodes > 1 and getattr(exp, "USE_MODAL_TORCH_DIST_WRAPPER", False)
     convert = (
@@ -180,15 +184,28 @@ def prepare_torch_dist(exp, checkpoint_volume, *, rank: int, master_addr: str) -
         if use_wrapper
         else f"{MILES_ROOT}/tools/convert_hf_to_torch_dist.py"
     )
-    inner = (
-        f"source {MILES_ROOT}/{exp.miles.miles_model_script} && "
-        f"PYTHONPATH={MEGATRON_PATH} torchrun"
-        f" --nnodes {nodes} --node-rank {rank} --master-addr {master_addr} --master-port 29500"
-        f" --nproc-per-node {exp.modal.torch_dist_prep_gpus_per_node}"
-        f" {convert} ${{MODEL_ARGS[@]}}"
-        f" --hf-checkpoint {bf16_dir} --save {torch_dist_dir} --megatron-to-hf-mode raw"
-        f" {exp.modal.torch_dist_convert_extra_args}"
-    )
+    command = [
+        "torchrun",
+        "--nnodes",
+        str(nodes),
+        "--node-rank",
+        str(rank),
+        "--master-addr",
+        master_addr,
+        "--master-port",
+        "29500",
+        "--nproc-per-node",
+        str(exp.modal.torch_dist_prep_gpus_per_node),
+        convert,
+        *model_args,
+        "--hf-checkpoint",
+        bf16_dir,
+        "--save",
+        torch_dist_dir,
+        "--megatron-to-hf-mode",
+        "raw",
+        *shlex.split(exp.modal.torch_dist_convert_extra_args),
+    ]
     env = {**os.environ}
     if use_wrapper:
         env["SKIP_RELEASE_RENAME"] = "1"
@@ -196,7 +213,8 @@ def prepare_torch_dist(exp, checkpoint_volume, *, rank: int, master_addr: str) -
         f"converting bf16 masters -> torch_dist ref_load ({nodes}-node torchrun, rank {rank})...",
         flush=True,
     )
-    subprocess.run(["bash", "-c", inner], check=True, env=env)
+    env["PYTHONPATH"] = f"{MEGATRON_PATH}:{env.get('PYTHONPATH', '')}"
+    subprocess.run(command, check=True, env=env)
     # Every node commits its own distcp shards (disjoint files merge on the Volume);
     # a rank-0-only commit would drop the other nodes' shards.
     checkpoint_volume.commit()
diff --git a/cookbook/miles_disagg/trainer_image.py b/cookbook/miles_disagg/trainer_image.py
index 92b0e041..2abd6aea 100644
--- a/cookbook/miles_disagg/trainer_image.py
+++ b/cookbook/miles_disagg/trainer_image.py
@@ -20,7 +20,7 @@
 # a moved mutable tag, so `latest` silently serves whatever was first pulled.
 MILES_IMAGE_TAG = "radixark/miles:dev-202607290235"
 MILES_REPO_URL = "https://github.com/modal-projects/miles.git"
-MILES_REPO_REF = "1eb7520018446cb94b7406715f66dff1a271b53b"  # stitch-weight-sync-v0516
+MILES_REPO_REF = "f83a68b296ca6f42d04e77e2b42c070c1ab02c70"
 
 MILES_ROOT = "/root/miles"
 # Source-only megatron.training must be on PYTHONPATH.

From 799214b984e7a2596eb5f91f653772cf83b4e8d2 Mon Sep 17 00:00:00 2001
From: Nan <nan@modal.com>
Date: Mon, 10 Aug 2026 16:19:26 +0000
Subject: [PATCH 4/5] [cookbook] Route sessions within rollout capacity

---
 cookbook/common/constants.py                  |  4 ++
 cookbook/common/hooks.py                      |  6 +-
 cookbook/common/hooks_test.py                 |  1 -
 cookbook/common/router.py                     | 66 +++++++++++--------
 cookbook/common/router_test.py                | 37 +++++++----
 cookbook/miles_disagg/app.py                  |  1 +
 .../miles_disagg/configs/glm45_air_bf16.py    |  1 -
 .../miles_disagg/configs/glm45_air_fp8.py     |  1 -
 .../configs/glm47_flash_swebench_pro.py       |  1 -
 cookbook/miles_disagg/configs/glm5_2_nvfp4.py |  1 -
 .../configs/kimi_k25_2layer_nvfp4.py          |  1 -
 .../miles_disagg/configs/kimi_k2_6_nvfp4.py   |  1 -
 .../miles_disagg/configs/moonlight_nvfp4.py   |  1 -
 .../configs/qwen3_30b_a3b_nvfp4_46.py         |  1 -
 cookbook/slime_disagg/app.py                  |  1 +
 .../slime_disagg/configs/kimi_k2_6_int4.py    |  1 -
 cookbook/slime_disagg/configs/moonlight.py    |  1 -
 .../slime_disagg/configs/moonlight_int4.py    |  1 -
 .../configs/qwen3_4b_delta_flash.py           |  2 -
 19 files changed, 71 insertions(+), 58 deletions(-)

diff --git a/cookbook/common/constants.py b/cookbook/common/constants.py
index 8660c41d..208feb29 100644
--- a/cookbook/common/constants.py
+++ b/cookbook/common/constants.py
@@ -19,6 +19,10 @@
 SGLANG_PORT = 8001  # the private sglang server behind the sidecar
 RAY_PORT = 6379
 
+# Modal's native sticky-routing header. The same session ID is also used by the
+# cookbook router when selecting a rollout replica.
+MODAL_SESSION_ID_HEADER = "Modal-Session-ID"
+
 # Timeouts.
 MINUTES = 60
 SERVER_STARTUP_TIMEOUT = 60 * MINUTES
diff --git a/cookbook/common/hooks.py b/cookbook/common/hooks.py
index 179726d7..7fc23790 100644
--- a/cookbook/common/hooks.py
+++ b/cookbook/common/hooks.py
@@ -19,6 +19,7 @@
 from stitch.types import PointerRewind, VersionRef
 
 from . import process
+from .constants import MODAL_SESSION_ID_HEADER
 
 logger = logging.getLogger(__name__)
 
@@ -126,9 +127,6 @@ async def gated_rollout_request_hook(
     weights beyond its lag bound."""
     payload, headers = request["payload"], dict(request.get("headers") or {})
     mode = str(getattr(args, "rollout_request_weight_version_mode", "min"))
-    affinity = str(
-        getattr(args, "rollout_session_affinity_header", "x-session-affinity")
-    )
 
     latest = exact = None
     lag = 0
@@ -146,7 +144,7 @@ async def gated_rollout_request_hook(
         lag=lag,
         exact=exact,
         session_id=sample_affinity_key(sample),
-        affinity_header=affinity,
+        affinity_header=MODAL_SESSION_ID_HEADER,
     )
     request["headers"] = headers
     request["max_retries"] = int(
diff --git a/cookbook/common/hooks_test.py b/cookbook/common/hooks_test.py
index e6205f4d..a95b8231 100644
--- a/cookbook/common/hooks_test.py
+++ b/cookbook/common/hooks_test.py
@@ -211,7 +211,6 @@ def test_request_hook_min_lag() -> None:
             str(root),
             rollout_request_weight_version_lag=2,
             rollout_request_retry_attempts=900,
-            rollout_session_affinity_header="Modal-Session-ID",
         )
         request = {"payload": {}}
         asyncio.run(
diff --git a/cookbook/common/router.py b/cookbook/common/router.py
index 95725339..6858a987 100644
--- a/cookbook/common/router.py
+++ b/cookbook/common/router.py
@@ -10,7 +10,7 @@
 own sticky routing turns the first saturated replicas into 503 attractors — sessions stuck
 on a full replica keep retrying it while the rest of the pool starves. Here, the registry
 polls every replica's live queue depth (``/v1/loads`` + a ``/health`` zombie check); the
-router pins each ``modal-session-id`` to the least-loaded healthy replica via the
+router pins each ``modal-session-id`` to a healthy replica with spare capacity via the
 ``modal-flash-upstream`` header, and a 503 from a pinned replica evicts it from rotation and
 retries on a healthier one, so load spreads instead of sticking.
 
@@ -26,7 +26,6 @@
 import asyncio
 import contextlib
 import logging
-import math
 import random
 import threading
 import time
@@ -42,11 +41,11 @@
 
 from stitch.pools.modal_flash import list_flash_containers_async
 
+from .constants import MODAL_SESSION_ID_HEADER
+
 logger = logging.getLogger(__name__)
 
 SESSION_ROUTE_TTL_SECONDS = 4 * 60 * 60
-SESSION_ROUTE_IMBALANCE_THRESHOLD = 0.5
-SESSION_ROUTE_OVERLOAD_FLOOR = 3
 SESSION_ROUTE_MAX_UPSTREAMS = 10
 
 CONTAINER_POLL_INTERVAL_SECONDS = 1.0
@@ -120,7 +119,7 @@ def filter_headers(headers: dict[str, str]) -> dict[str, str]:
         "modal-flash-upstream",
         "modal-key",
         "modal-secret",
-        "modal-session-id",
+        MODAL_SESSION_ID_HEADER.lower(),
         "x-forwarded-for",
         "x-forwarded-host",
         "x-forwarded-port",
@@ -131,20 +130,33 @@ def filter_headers(headers: dict[str, str]) -> dict[str, str]:
     return {k: v for k, v in headers.items() if k.lower() not in removed}
 
 
-def select_least_loaded_container(
-    containers: dict[str, ContainerInfo],
+def select_underloaded_container(
+    containers: dict[str, ContainerInfo], overload_threshold: int
 ) -> ContainerInfo:
-    minimum = min(c.load for c in containers.values())
-    return random.choice([c for c in containers.values() if c.load == minimum])
+    """Spread new sessions across replicas with headroom.
+
+    Registry loads are snapshots shared by multiple router replicas. Selecting the exact
+    minimum makes every router converge on a newly ready zero-load replica before the next
+    snapshot, creating a thundering herd. Random choice from the healthy underloaded set
+    preserves load shedding without requiring distributed per-request reservations.
+    """
+    candidates = [
+        container
+        for container in containers.values()
+        if container.load < overload_threshold
+    ]
+    return random.choice(candidates or list(containers.values()))
 
 
 async def route_session(
-    session_routes: modal.Dict, session_id: str, containers: dict[str, ContainerInfo]
+    session_routes: modal.Dict,
+    session_id: str,
+    containers: dict[str, ContainerInfo],
+    overload_threshold: int,
 ) -> ContainerInfo:
     """Pick the replica for one session: the most-recently-used of its known replicas
-    that isn't overloaded, else the pool's least-loaded. A replica is overloaded when its
-    load is ≥50% above the pool's mean (floored), so a hot replica sheds new session
-    traffic without a hard capacity number. Mutates and persists the session's routes."""
+    that has room below the pool's configured soft capacity, else a random replica with
+    headroom. Mutates and persists the session's routes."""
     current_time = time.time()
 
     routes: list[RouteEntry] = RouteEntryList.validate_python(
@@ -160,15 +172,6 @@ async def route_session(
     ]
     routes.sort(key=lambda entry: entry.last_sent, reverse=True)
 
-    overload_threshold = max(
-        math.ceil(
-            sum(c.load for c in containers.values())
-            / len(containers)
-            * (1 + SESSION_ROUTE_IMBALANCE_THRESHOLD)
-        ),
-        SESSION_ROUTE_OVERLOAD_FLOOR,
-    )
-
     async def save_routes() -> None:
         await session_routes.put.aio(
             session_id,
@@ -184,7 +187,7 @@ async def save_routes() -> None:
             await save_routes()
             return container
 
-    selected = select_least_loaded_container(containers)
+    selected = select_underloaded_container(containers, overload_threshold)
     reason = (
         f"previous upstreams [{', '.join(entry.task_id for entry in routes)}] overloaded"
         f" (load >= {overload_threshold})"
@@ -271,12 +274,14 @@ def serve_router(
     registry_url: str,
     upstream_url: str,
     session_routes: modal.Dict,
+    overload_threshold: int,
 ) -> None:
     """Start the session-routing proxy on a ``Router`` container (@modal.enter)."""
     router = _ProxyApp(
         registry_url=registry_url,
         upstream_url=upstream_url,
         session_routes=session_routes,
+        overload_threshold=overload_threshold,
     )
     router.start()
     replica._router_server = router
@@ -390,11 +395,17 @@ class _ProxyApp(_UvicornApp):
     """Proxies requests to rollout replicas with session-affinity routing."""
 
     def __init__(
-        self, *, registry_url: str, upstream_url: str, session_routes: modal.Dict
+        self,
+        *,
+        registry_url: str,
+        upstream_url: str,
+        session_routes: modal.Dict,
+        overload_threshold: int,
     ) -> None:
         self.registry_url = registry_url.rstrip("/")
         self.upstream_url = upstream_url.rstrip("/")
         self.session_routes = session_routes
+        self.overload_threshold = overload_threshold
         self.containers: dict[str, ContainerInfo] = {}
 
     @contextlib.contextmanager
@@ -510,7 +521,7 @@ async def forward(request: Request, path: str) -> Any:
                 return Response(content=b"", status_code=200, media_type="text/plain")
 
             body = await request.body()
-            session_id = request.headers.get("modal-session-id")
+            session_id = request.headers.get(MODAL_SESSION_ID_HEADER)
             log_prefix = f"[request {uuid.uuid4()}, session {session_id}]"
 
             try:
@@ -520,7 +531,10 @@ async def forward(request: Request, path: str) -> Any:
                     container = None
                     if session_id and self.containers:
                         container = await route_session(
-                            self.session_routes, session_id, self.containers
+                            self.session_routes,
+                            session_id,
+                            self.containers,
+                            self.overload_threshold,
                         )
                         headers["modal-flash-upstream"] = container.upstream
 
diff --git a/cookbook/common/router_test.py b/cookbook/common/router_test.py
index cc7c5e1c..5b414462 100644
--- a/cookbook/common/router_test.py
+++ b/cookbook/common/router_test.py
@@ -1,5 +1,5 @@
 """Router harness: the pure routing helpers — ``route_session`` stickiness / overload
-fallback / TTL eviction, ``select_least_loaded_container``, ``filter_headers``, and
+fallback / TTL eviction, ``select_underloaded_container``, ``filter_headers``, and
 ``_container_addr`` — against a fake session-routes dict (no Modal involved)."""
 
 from __future__ import annotations
@@ -17,7 +17,7 @@
     _ProxyApp,
     filter_headers,
     route_session,
-    select_least_loaded_container,
+    select_underloaded_container,
 )
 
 
@@ -56,31 +56,31 @@ def _seeded(routes: FakeRoutes, session: str, entries: list[dict]) -> None:
     )
 
 
-def test_route_session_pins_least_loaded_on_first_request() -> None:
+def test_route_session_pins_only_available_underloaded_replica() -> None:
     routes, containers = FakeRoutes(), _containers(0, 5, 5)
-    picked = asyncio.run(route_session(routes, "s1", containers))
+    picked = asyncio.run(route_session(routes, "s1", containers, 4))
     assert picked.task_id == "ta-0"
     assert routes.store["s1"][0]["task_id"] == "ta-0"
 
 
 def test_route_session_is_sticky_for_known_healthy_replica() -> None:
     routes, containers = FakeRoutes(), _containers(0, 5, 5)
-    first = asyncio.run(route_session(routes, "s1", containers))
+    first = asyncio.run(route_session(routes, "s1", containers, 4))
     # ta-1 becomes strictly less loaded, but stickiness holds while the pinned
-    # replica stays below the overload threshold (3 < ceil(avg × 1.5) = 4).
+    # replica stays below the configured overload threshold.
     containers["ta-0"] = containers["ta-0"].model_copy(update={"load": 3})
     containers["ta-1"] = containers["ta-1"].model_copy(update={"load": 0})
-    second = asyncio.run(route_session(routes, "s1", containers))
+    second = asyncio.run(route_session(routes, "s1", containers, 4))
     assert second.task_id == first.task_id == "ta-0"
 
 
-def test_route_session_sheds_overloaded_replica_to_least_loaded() -> None:
+def test_route_session_sheds_overloaded_replica_to_replica_with_headroom() -> None:
     routes, containers = FakeRoutes(), _containers(0, 0, 20)
-    first = asyncio.run(route_session(routes, "s1", containers))
+    first = asyncio.run(route_session(routes, "s1", containers, 10))
     containers[first.task_id] = containers[first.task_id].model_copy(
         update={"load": 20}
     )
-    second = asyncio.run(route_session(routes, "s1", containers))
+    second = asyncio.run(route_session(routes, "s1", containers, 10))
     assert second.task_id != first.task_id
     assert second.load == 0
 
@@ -98,15 +98,23 @@ def test_route_session_drops_expired_and_undiscovered_routes() -> None:
             {"task_id": "ta-gone", "last_sent": time.time()},
         ],
     )
-    picked = asyncio.run(route_session(routes, "s1", containers))
+    picked = asyncio.run(route_session(routes, "s1", containers, 4))
     assert picked.task_id == "ta-0"
     assert [entry["task_id"] for entry in routes.store["s1"]] == ["ta-0"]
 
 
-def test_select_least_loaded_container_picks_minimum() -> None:
+def test_select_underloaded_container_spreads_across_headroom(monkeypatch) -> None:
     containers = _containers(7, 2, 3, 2)
-    picked = select_least_loaded_container(containers)
-    assert picked.load == 2
+    candidates = []
+
+    def choose(options):
+        candidates.extend(options)
+        return options[0]
+
+    monkeypatch.setattr("cookbook.common.router.random.choice", choose)
+    picked = select_underloaded_container(containers, overload_threshold=4)
+    assert picked.load < 4
+    assert {container.task_id for container in candidates} == {"ta-1", "ta-2", "ta-3"}
 
 
 def test_filter_headers_drops_routing_and_hop_by_hop() -> None:
@@ -164,6 +172,7 @@ async def run() -> None:
             registry_url="https://registry",
             upstream_url="https://upstream",
             session_routes=FakeRoutes(),
+            overload_threshold=4,
         )
         client = FakeClient()
         app.client = client
diff --git a/cookbook/miles_disagg/app.py b/cookbook/miles_disagg/app.py
index 0d5e762b..65e57ba3 100644
--- a/cookbook/miles_disagg/app.py
+++ b/cookbook/miles_disagg/app.py
@@ -254,6 +254,7 @@ def enter(self) -> None:
             registry_url=RouterRegistry.get_url(),
             upstream_url=Server.get_url(),
             session_routes=session_routes,
+            overload_threshold=ROLLOUT_CONCURRENCY,
         )
 
     @modal.exit()
diff --git a/cookbook/miles_disagg/configs/glm45_air_bf16.py b/cookbook/miles_disagg/configs/glm45_air_bf16.py
index 6536f066..5a67b3bb 100644
--- a/cookbook/miles_disagg/configs/glm45_air_bf16.py
+++ b/cookbook/miles_disagg/configs/glm45_air_bf16.py
@@ -91,7 +91,6 @@ class _Miles(MilesConfig):
         "rollout_request_weight_version_lag": 1,
         "rollout_request_retry_attempts": 900,
         "rollout_request_retry_sleep": 1.0,
-        "rollout_session_affinity_header": "Modal-Session-ID",
         "rollout_request_timeout_secs": 300,
     }
 
diff --git a/cookbook/miles_disagg/configs/glm45_air_fp8.py b/cookbook/miles_disagg/configs/glm45_air_fp8.py
index 948080ce..6c72fc37 100644
--- a/cookbook/miles_disagg/configs/glm45_air_fp8.py
+++ b/cookbook/miles_disagg/configs/glm45_air_fp8.py
@@ -88,7 +88,6 @@ class _Miles(MilesConfig):
         "rollout_request_weight_version_lag": 1,
         "rollout_request_retry_attempts": 900,
         "rollout_request_retry_sleep": 1.0,
-        "rollout_session_affinity_header": "Modal-Session-ID",
         "rollout_request_timeout_secs": 300,
     }
 
diff --git a/cookbook/miles_disagg/configs/glm47_flash_swebench_pro.py b/cookbook/miles_disagg/configs/glm47_flash_swebench_pro.py
index 4d144a77..440f8fad 100644
--- a/cookbook/miles_disagg/configs/glm47_flash_swebench_pro.py
+++ b/cookbook/miles_disagg/configs/glm47_flash_swebench_pro.py
@@ -123,7 +123,6 @@ class _Miles(MilesConfig):
         "rollout_request_weight_version_lag": 1,
         "rollout_request_retry_attempts": 1200,
         "rollout_request_retry_sleep": 1.0,
-        "rollout_session_affinity_header": "Modal-Session-ID",
         "rollout_request_timeout_secs": 300,
     }
 
diff --git a/cookbook/miles_disagg/configs/glm5_2_nvfp4.py b/cookbook/miles_disagg/configs/glm5_2_nvfp4.py
index 0a2a4673..37172d6c 100644
--- a/cookbook/miles_disagg/configs/glm5_2_nvfp4.py
+++ b/cookbook/miles_disagg/configs/glm5_2_nvfp4.py
@@ -222,7 +222,6 @@ class _Miles(MilesConfig):
         "rollout_request_weight_version_lag": 1,
         "rollout_request_retry_attempts": 1200,
         "rollout_request_retry_sleep": 1.0,
-        "rollout_session_affinity_header": "Modal-Session-ID",
         "rollout_request_timeout_secs": 300,
     }
 
diff --git a/cookbook/miles_disagg/configs/kimi_k25_2layer_nvfp4.py b/cookbook/miles_disagg/configs/kimi_k25_2layer_nvfp4.py
index 1b5a46ae..bfa9d4e0 100644
--- a/cookbook/miles_disagg/configs/kimi_k25_2layer_nvfp4.py
+++ b/cookbook/miles_disagg/configs/kimi_k25_2layer_nvfp4.py
@@ -83,7 +83,6 @@ class _Miles(MilesConfig):
         "rollout_request_weight_version_lag": 1,
         "rollout_request_retry_attempts": 240,
         "rollout_request_retry_sleep": 1.0,
-        "rollout_session_affinity_header": "Modal-Session-ID",
     }
 
     async_mode = True
diff --git a/cookbook/miles_disagg/configs/kimi_k2_6_nvfp4.py b/cookbook/miles_disagg/configs/kimi_k2_6_nvfp4.py
index 4f24bcaa..02fd1bb1 100644
--- a/cookbook/miles_disagg/configs/kimi_k2_6_nvfp4.py
+++ b/cookbook/miles_disagg/configs/kimi_k2_6_nvfp4.py
@@ -103,7 +103,6 @@ class _Miles(MilesConfig):
         # 1200x1s = 20 min, outlasts a ~16 min cold-load.
         "rollout_request_retry_attempts": 1200,
         "rollout_request_retry_sleep": 1.0,
-        "rollout_session_affinity_header": "Modal-Session-ID",
         # finite read timeout, else a request to a scaled-down container hangs forever.
         "rollout_request_timeout_secs": 300,
     }
diff --git a/cookbook/miles_disagg/configs/moonlight_nvfp4.py b/cookbook/miles_disagg/configs/moonlight_nvfp4.py
index 3617562b..8ff8e1d8 100644
--- a/cookbook/miles_disagg/configs/moonlight_nvfp4.py
+++ b/cookbook/miles_disagg/configs/moonlight_nvfp4.py
@@ -87,7 +87,6 @@ class _Miles(MilesConfig):
         "rollout_request_weight_version_lag": 1,
         "rollout_request_retry_attempts": 240,
         "rollout_request_retry_sleep": 1.0,
-        "rollout_session_affinity_header": "Modal-Session-ID",
     }
 
     async_mode = True
diff --git a/cookbook/miles_disagg/configs/qwen3_30b_a3b_nvfp4_46.py b/cookbook/miles_disagg/configs/qwen3_30b_a3b_nvfp4_46.py
index fd5e27b1..856abb5d 100644
--- a/cookbook/miles_disagg/configs/qwen3_30b_a3b_nvfp4_46.py
+++ b/cookbook/miles_disagg/configs/qwen3_30b_a3b_nvfp4_46.py
@@ -155,7 +155,6 @@ class _Miles(MilesConfig):
         "rollout_request_weight_version_lag": 1,
         "rollout_request_retry_attempts": 240,
         "rollout_request_retry_sleep": 1.0,
-        "rollout_session_affinity_header": "Modal-Session-ID",
     }
 
     async_mode = True
diff --git a/cookbook/slime_disagg/app.py b/cookbook/slime_disagg/app.py
index 56d09125..b903bf6d 100644
--- a/cookbook/slime_disagg/app.py
+++ b/cookbook/slime_disagg/app.py
@@ -242,6 +242,7 @@ def enter(self) -> None:
             registry_url=RouterRegistry.get_url(),
             upstream_url=Server.get_url(),
             session_routes=session_routes,
+            overload_threshold=ROLLOUT_CONCURRENCY,
         )
 
     @modal.exit()
diff --git a/cookbook/slime_disagg/configs/kimi_k2_6_int4.py b/cookbook/slime_disagg/configs/kimi_k2_6_int4.py
index da7774a9..e36258a2 100644
--- a/cookbook/slime_disagg/configs/kimi_k2_6_int4.py
+++ b/cookbook/slime_disagg/configs/kimi_k2_6_int4.py
@@ -84,7 +84,6 @@ class _Slime(SlimeConfig):
     rollout_request_weight_version_lag = 1
     rollout_request_retry_attempts = 240
     rollout_request_retry_sleep = 1.0
-    rollout_session_affinity_header = "Modal-Session-ID"
 
     async_mode = True
     update_weights_interval = 1
diff --git a/cookbook/slime_disagg/configs/moonlight.py b/cookbook/slime_disagg/configs/moonlight.py
index dfada624..af3a1ccd 100644
--- a/cookbook/slime_disagg/configs/moonlight.py
+++ b/cookbook/slime_disagg/configs/moonlight.py
@@ -63,7 +63,6 @@ class _Slime(SlimeConfig):
     rollout_request_weight_version_lag = 1
     rollout_request_retry_attempts = 240
     rollout_request_retry_sleep = 1.0
-    rollout_session_affinity_header = "Modal-Session-ID"
 
     async_mode = True
     update_weights_interval = 1
diff --git a/cookbook/slime_disagg/configs/moonlight_int4.py b/cookbook/slime_disagg/configs/moonlight_int4.py
index 18067c19..e9aa4c18 100644
--- a/cookbook/slime_disagg/configs/moonlight_int4.py
+++ b/cookbook/slime_disagg/configs/moonlight_int4.py
@@ -65,7 +65,6 @@ class _Slime(SlimeConfig):
     rollout_request_weight_version_lag = 1
     rollout_request_retry_attempts = 240
     rollout_request_retry_sleep = 1.0
-    rollout_session_affinity_header = "Modal-Session-ID"
 
     async_mode = True
     update_weights_interval = 1
diff --git a/cookbook/slime_disagg/configs/qwen3_4b_delta_flash.py b/cookbook/slime_disagg/configs/qwen3_4b_delta_flash.py
index 9df34815..73204492 100644
--- a/cookbook/slime_disagg/configs/qwen3_4b_delta_flash.py
+++ b/cookbook/slime_disagg/configs/qwen3_4b_delta_flash.py
@@ -57,8 +57,6 @@ class _Slime(SlimeConfig):
     rollout_request_weight_version_lag = 0
     rollout_request_retry_attempts = 240
     rollout_request_retry_sleep = 1.0
-    # session affinity so GRPO siblings co-locate on one Flash replica.
-    rollout_session_affinity_header = "Modal-Session-ID"
 
     # disk-delta publish-only: slime writes weight_v{N}/ + `latest`; the hook commits and wakes the pool.
     update_weight_mode = "delta"

From b40df83fda611eae19748da35f3ea69fc8ed502c Mon Sep 17 00:00:00 2001
From: Nan <nan@modal.com>
Date: Mon, 10 Aug 2026 16:19:40 +0000
Subject: [PATCH 5/5] Document the validated v0.5.17 rollout stack

---
 README.md                      | 50 ++++++++++++++-------
 cookbook/README.md             | 18 +++++---
 cookbook/common/SGLANG_FORK.md | 80 +++++++++++++++-------------------
 3 files changed, 79 insertions(+), 69 deletions(-)

diff --git a/README.md b/README.md
index ab49e5ac..4f6060b3 100644
--- a/README.md
+++ b/README.md
@@ -43,23 +43,41 @@ recover without becoming part of the trainer's process lifecycle.
 
 ## Measured delta updates
 
-These single-update measurements use the pinned cookbook stack, 64 vCPUs, and
-element-wise XOR deltas. Remote transfer, delta generation, and one-time CPU
-destination initialization are excluded. Preparation runs while inference
-remains available; only activation pauses the engine.
-
-| Model | TP | Canonical checkpoint | Preparation | Engine pause | Total update |
+These are verified single-update measurements from the pinned v0.5.17 cookbook
+stack. Every row completed checksum verification, generated successfully before,
+during, and after the update, changed from the base, and reproduced the exact
+post-update text, tokens, and logprobs. The synthetic XOR deltas are element-wise
+over rollout-visible values and change nearly every trained tensor.
+Quantized-value densities are 0.6% for FP8, 0.3% for Kimi K2.6 NVFP4 and Kimi
+K3 MXFP4, and 0.375% for mixed GLM-5.2; high-precision values change at 1%.
+Remote transfer, delta generation, and one-time CPU destination initialization
+are excluded. Preparation runs while inference remains available; only
+activation pauses the engine.
+
+| Model | TP | Update path | Preparation | Engine pause | Total update |
 | --- | ---: | --- | ---: | ---: | ---: |
-| GLM-4.5-Air FP8 | 4 | RAM | 41.2 s | 1.0 s | 42.2 s |
-| Kimi K2.6 NVFP4 | 4 | RAM | 55.6 s | 2.78 s | 58.4 s |
-| Kimi K2.6 NVFP4 | 4 | NVMe | 102.3 s | 2.76 s | 105.1 s |
-| GLM-5.2 all-NVFP4 checkpoint | 4 | RAM | 59.3 s | 2.14 s | 61.5 s |
-| GLM-5.2 mixed NVFP4/BF16 | 4 | RAM | 86.2 s | 3.01 s | 89.2 s |
-| Kimi K3 MXFP4 | 8 | RAM | 122.3 s | 3.82 s | 126.1 s |
-| Kimi K3 MXFP4 | 8 | NVMe | 283.8 s | 3.79 s | 287.5 s |
-
-The mixed GLM-5.2 profile changes 0.375% of rollout-visible NVFP4 values and
-1% of BF16 values; its compressed delta is 10.06 GB.
+| GLM-4.5-Air FP8 | 4 | CPU cache; canonical in RAM | 26.6 s | 0.99 s | 27.5 s |
+| GLM-4.5-Air FP8 | 4 | CPU cache; canonical on NVMe | 85.5 s | 0.98 s | 86.5 s |
+| GLM-4.5-Air FP8 | 4 | Disk checkpoint | 18.5 s | 27.15 s | 45.7 s |
+| Kimi K2.6 NVFP4 | 4 | CPU cache; canonical in RAM | 72.5 s | 2.82 s | 75.3 s |
+| Kimi K2.6 NVFP4 | 4 | CPU cache; canonical on NVMe | 279.0 s | 3.23 s | 282.2 s |
+| Kimi K2.6 NVFP4 | 4 | Disk checkpoint | 149.0 s | 165.09 s | 314.1 s |
+| GLM-5.2 mixed NVFP4/BF16 | 4 | CPU cache; canonical in RAM | 116.4 s | 3.30 s | 119.7 s |
+| GLM-5.2 mixed NVFP4/BF16 | 4 | CPU cache; canonical on NVMe | 164.5 s | 3.26 s | 167.7 s |
+| GLM-5.2 mixed NVFP4/BF16 | 4 | Disk checkpoint | 128.0 s | 299.38 s | 427.4 s |
+| Kimi K3 MXFP4 | 8 | CPU cache; canonical in RAM | 120.5 s | 3.84 s | 124.3 s |
+| Kimi K3 MXFP4 | 8 | CPU cache; canonical on NVMe | 1,089.7 s | 3.83 s | 1,093.6 s |
+| Kimi K3 MXFP4 | 8 | Disk checkpoint | 704.0 s | 300.37 s | 1,004.4 s |
+
+These are wall-clock samples, not a hardware distribution. NVMe preparation
+reads and writes a complete canonical checkpoint and therefore tracks the
+assigned host's local-storage bandwidth; the Kimi K2.6 and Kimi K3 NVMe samples
+are therefore host-specific rather than model-only transformation costs. K3's
+canonical checkpoint and eight rank images occupy 3.22 TB before engine and
+staging overhead; the all-RAM sample reached 3.29 TB after staging. Its supplied
+recipe therefore keeps the canonical checkpoint on NVMe to preserve operating
+headroom. The mixed GLM-5.2 and K3 deltas are 10.04 GB and 24.13 GB compressed,
+respectively.
 
 Each profiler reconstructs and checksums the complete target and validates
 generation before, during, and after activation. See
diff --git a/cookbook/README.md b/cookbook/README.md
index 673b6be0..dcf3e5a0 100644
--- a/cookbook/README.md
+++ b/cookbook/README.md
@@ -124,8 +124,8 @@ fully asynchronous Miles training on SWE-bench Pro.
 
 | Component | Configuration |
 | --- | --- |
-| Trainer | 2 nodes × 8 H200 GPUs |
-| Rollout | 20 warm replicas × 1 H200 GPU, with autoscaling above the warm fleet |
+| Trainer | 4 nodes × 8 H200 GPUs |
+| Rollout | 48 replicas × 1 H200 GPU |
 | Model | `zai-org/GLM-4.7-Flash`, pinned BF16 revision |
 | Dataset | SWE-bench Pro, including task environments and verifiers |
 | Weight sync | Checksummed XOR deltas with CPU preparation and in-place activation |
@@ -147,7 +147,7 @@ uv run --extra modal python -m cookbook.miles_disagg.launch
 ```
 
 Checkpoint preparation materializes the pinned BF16 model. TorchDist
-preparation converts it for the two-node trainer. Dataset preparation writes
+preparation converts it for the four-node trainer. Dataset preparation writes
 the pinned prompts, task environments, verifiers, and source manifest.
 
 ### Weight-update performance
@@ -244,12 +244,14 @@ the replica. Bundled MTP heads do not need a separate volume.
 ## Profile a weight update
 
 The model profilers prepare their pinned base checkpoint and synthetic delta,
-then run with `--update-mode disk|cpu`:
+then run with `--update-mode disk|cpu`. CPU runs also select
+`--canonical-storage memory|disk`; `disk` uses host-local NVMe.
 
 ```bash
 uv run --extra modal modal run -d \
   tools/profiling/glm45_air_fp8_delta_weight_update.py \
-  --update-mode cpu
+  --update-mode cpu \
+  --canonical-storage memory
 ```
 
 Prepared artifacts are reused. The profilers generate during staging, pause
@@ -259,10 +261,12 @@ for DFlash. DSpark still rejects logprob-returning requests, so its profiles
 compare repeated deterministic text instead of token IDs and logprobs.
 
 The K3 profiler downloads the pinned public checkpoint and constructs a
-checksummed XOR publication covering every checkpoint tensor:
+checksummed XOR publication over mutable, rollout-visible values. The fixed
+vision tower and projector are excluded.
 
 ```bash
 uv run --extra modal modal run -d \
   tools/profiling/kimi_k3_mxfp4_delta_weight_update.py \
-  --update-mode cpu
+  --update-mode cpu \
+  --canonical-storage disk
 ```
diff --git a/cookbook/common/SGLANG_FORK.md b/cookbook/common/SGLANG_FORK.md
index 1dd3c6db..7187a1fc 100644
--- a/cookbook/common/SGLANG_FORK.md
+++ b/cookbook/common/SGLANG_FORK.md
@@ -11,51 +11,30 @@ default runtime:
 
 ```python
 DEFAULT_SGLANG_RUNTIME = SGLangRuntime(
-    image="lmsysorg/sglang:v0.5.16",
+    image="lmsysorg/sglang:v0.5.17",
     repository="https://github.com/modal-projects/sglang.git",
-    branch="stitch-sglang-v0.5.16",
-    commit="a73ea9507fb981462768dbc5e869bdfeb5c48116",
+    branch="stitch-sglang-v0.5.17",
+    commit="0c79627e857eec795298a372adadf649209cdf2f",
 )
 ```
 
-The branch is upstream v0.5.16 plus:
+The branch is upstream v0.5.17 plus four independently reviewable layers:
 
-| Commit | Responsibility |
+| Layer | Responsibility |
 | --- | --- |
-| `49031cb24c` | Configure fastsafetensors with or without GDS and honor post-load cache release. |
-| `5f6e1e613f` | Materialize and verify complete targets on host-local disk. |
-| `a7e20596ba` | Restore checkpoint-facing quantized layouts for complete weight loading. |
-| `c867782f3e` | Build verified, rank-ready CPU weight images from canonical targets. |
-| `111a804d2e` | Expose asynchronous disk/CPU staging and CPU-to-GPU target-model commit APIs. |
-| `e0859b7390` | Stream CPU delta lineages through bounded memory. |
-| `77eca472e6` | Fold disk XOR lineages with bounded positional I/O. |
-| `526e0ddca2` | Fail cache-flushing CPU commits before GPU mutation when the engine is busy. |
-| `a562908a10` | Normalize native ModelOpt FP4 expert tensors through their existing loader path. |
-| `1a4a4fd6b5` | Return aligned verifier logprobs for DFlash rollout tokens. |
-| `607e107b44` | Store the canonical CPU-cache checkpoint on NVMe and overlap verified persistence with bounded rank-image compilation. |
-| `1051a95a6a` | Balance CPU delta transforms across persistent worker tasks. |
-| `0094b725b9` | Support top-p-only sampling masks through the native generation and chat-completions APIs. |
-| `7b09ce9f77` | Return aligned top-p sampling masks for tokens accepted by DFlash SpecV2. |
-| `a50de4fe3e` | Monitor data-parallel scheduler subprocesses and fail when one exits. |
-| `e02a07c905` | Preserve routed-expert and indexer top-k state-capture outputs through DFlash. |
-| `af563ae597` | Remove attention-TP alignment rows before materializing DFlash prefill KV state. |
-| `270c78efaa` | Resolve exact and prefix aborts that arrive while requests are still held by the tokenizer. |
-| `325abb5afa` | Keep a tokenizer-held abort result alive until its request waiter consumes it. |
-| `a73ea9507f` | Accept top-p sampling-mask requests preserved in compatibility metadata by typed routers. |
-
-The image and branch must use the same SGLang release because Stitch overlays
-Python code onto the image’s existing CUDA and C++ extensions.
-
-Models that require another upstream SGLang line set `SGLANG_RUNTIME` in their
-configuration. The image, fork branch, and immutable commit stay together so
-the Python overlay remains ABI-compatible with the image.
-
-Kimi K3 MXFP4 recipes pin the public K3 image and `stitch-sglang-kimi-k3` fork.
-That fork ports the same weight-sync responsibilities onto SGLang’s public
-`kimi-k3` branch; other recipes continue to use the v0.5.16 default. Its
-K3-native loader narrows expert lookup, batches safe copies, scopes post-load
-work to loaded modules, and transforms Blackwell MXFP4 runtime layouts on GPU
-before caching rank-ready host images.
+| Reload lifecycle | Restore checkpoint-facing layouts, run each quantization method's native loader and post-load hooks, and fail closed if a partially mutated model cannot be rolled back. |
+| Verified materialization | Apply and fold complete XOR delta lineages in canonical checkpoint space, verify the published checksum, and durably materialize disk targets. |
+| CPU staging | Build bounded rank-ready host images while serving, optionally keep the canonical checkpoint on local NVMe, then commit every runtime storage in place. |
+| Serving correctness | Preserve routed-expert state and sampling masks across data-parallel and speculative paths, classify client cancellations, and surface scheduler-process failures. |
+
+The branch history keeps these physical responsibilities in separate commits;
+the immutable pin above is the executable definition of the stack.
+
+The image and immutable source pin stay together so the Python overlay remains
+ABI-compatible with the image's CUDA and C++ extensions. SGLang v0.5.17 includes
+Kimi K3, so all cookbook recipes now use this one runtime line. The fork's MXFP4
+staging path transforms runtime layouts on GPU before caching rank-ready host
+images.
 
 ## API
 
@@ -129,8 +108,10 @@ CPU mode keeps rank-ready images in RAM for the shortest commit:
 1. After v0 begins serving, SGLang allocates one complete rank-ready image per
    local TP rank and either caches one canonical checkpoint per host in RAM or
    materializes it on host-local storage.
-2. It builds v0 through the model’s ordinary weight loader and quantization
-   hooks and verifies that the prepared runtime storages match the active model.
+2. When the base is the boot checkpoint, it captures the already-realized active
+   runtime storages into the rank images instead of repeating a model-sized
+   load. A different base goes through the model's ordinary loader and
+   quantization hooks and must match the requested checkpoint before use.
 3. For every delta lineage, it reconstructs and checksums the canonical target,
    then builds every next rank image while inference continues. The in-memory
    path streams deltas through a bounded work budget; the storage-backed path
@@ -177,8 +158,9 @@ retain those reclaimable pages for later reads.
 
 The group bound limits transient loader work; it does not tune correctness or
 assume a model architecture. An indivisible module larger than the requested
-bound remains intact and is reported. The K3 recipe uses a 16 GiB bound to
-balance its module granularity against transient RAM.
+bound remains intact and is reported. Each staging clone is reclaimed at its
+group boundary, so CUDA-required post-load transforms cannot accumulate a
+second model-sized device copy.
 
 With the default in-memory canonical checkpoint, persistent host RAM is:
 
@@ -187,6 +169,11 @@ one canonical checkpoint per host
 + one rank-local runtime image per local TP rank
 ```
 
+The canonical checkpoint is interleaved across the host's allowed NUMA nodes so
+it cannot exhaust one GPU-local node while capacity remains elsewhere. Rank
+images remain GPU-local because they are the source of the latency-sensitive
+CPU-to-GPU commit.
+
 With a storage-backed canonical checkpoint, persistent host RAM is the rank
 images; local storage holds one canonical checkpoint. File-cache pages used
 during preparation are reclaimable.
@@ -197,7 +184,7 @@ Measured component sizes are:
 | --- | ---: | ---: | ---: |
 | GLM-4.5-Air FP8 | 4 | 112.6 GB | 27.2 GB × 4 |
 | Kimi K2.6 NVFP4 | 4 | about 595 GB | about 151 GB × 4 |
-| GLM-5.2 mixed NVFP4/BF16 | 4 | 617.6 GB | 156.3 GB × 4 |
+| GLM-5.2 mixed NVFP4/BF16 | 4 | 617.6 GB | 179.3 GB × 4 |
 | Kimi K3 MXFP4 | 8 | 1.561 TB | 207.5 GB × 8 |
 
 Allow additional memory for the engine process, delta decoding, and bounded
@@ -205,8 +192,9 @@ loader staging. The supplied GLM-4.5 recipe requests `(512 GiB, 2 TiB)`;
 GLM-5.2, Kimi K2.6, and Kimi K3 request `(1 TiB, 3 TiB)`, expressed as
 `(request, limit)`.
 
-All runtime storages are prepared and committed. Element-wise sparsity only
-reduces the compressed delta transport and XOR work.
+All runtime storages are prepared and committed. Element-wise sparsity reduces
+the compressed delta transport and storage, but not the full-target checksum,
+sharding, runtime-layout conversion, or CPU-to-GPU commit.
 
 ## Correctness