Reconstruct the validated v0.5.17 rollout stack

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

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
 ```

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,16 +184,17 @@ 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
 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
 

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

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(

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(