From 830de9b49c19ff25b95b39c9b747f4b091872808 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 14 Oct 2026 12:57:10 +0000
Subject: [PATCH 01/88] feat(weights): sparse XOR deltas for the filesystem
 weight broadcast

---
 docs/scaling.md                               |   4 +
 .../src/prime_rl/configs/shared.py            |   9 +
 src/prime_rl/inference/vllm/server.py         |  12 ++
 .../inference/vllm/worker/filesystem.py       | 122 ++++++++++--
 .../inference/vllm/worker/weight_transfer.py  | 132 +++++++++++++
 src/prime_rl/orchestrator/clients.py          |  14 ++
 src/prime_rl/transports/weights/filesystem.py | 177 +++++++++++++++++-
 src/prime_rl/utils/pathing.py                 |  16 +-
 src/prime_rl/utils/sparse_weights.py          |  89 +++++++++
 tests/unit/utils/test_sparse_weights.py       |  54 ++++++
 10 files changed, 604 insertions(+), 25 deletions(-)
 create mode 100644 src/prime_rl/utils/sparse_weights.py
 create mode 100644 tests/unit/utils/test_sparse_weights.py

diff --git a/docs/scaling.md b/docs/scaling.md
index 73bbf68027c..f7968b09665 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -274,6 +274,10 @@ ModelExpress exchanges peer metadata during startup. Weight updates reuse prepar
 
 By default, the trainer and inference worker each allocate one transfer arena. Set `weight_broadcast.overlap_transfer_and_replay = true` to allocate two arenas on both sides and replay one weight group while receiving the next. The additional arena is the size of the largest transfer group per GPU; allocation errors are reported instead of silently disabling overlap.
 
+### Sparse filesystem weight broadcast
+
+With `[weight_broadcast] type = "filesystem"` and `sparse = true`, the trainer writes a full checkpoint only every `anchor_interval` steps (default 25) and at its first two broadcasts. In between, it writes the XOR of the bf16 weights against the previous broadcast. Each trainer rank diffs its own shards, so no weights are gathered; with AdamW at RL learning rates about 1% of bf16 elements change per step, and a delta is about 1% of the checkpoint size. Engines read and decode a delta while they keep serving, then patch their weights in place during the pause. An engine that is not at a delta's base step reloads the last anchor and replays the deltas since. On anchor steps, engines apply the delta first and check it against the full checkpoint bit for bit, so a model whose weights cannot be patched in place fails at the second broadcast. The delta goes through vLLM's own weight loader; weights that are transformed after loading (quantized or re-laid-out) are not supported. LoRA is not supported.
+
 ### Custom Templates
 
 For unusual partitions, module loads, or environment setup, supply your own Jinja2 template:
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
index 1ce168817a9..bd23b47b955 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
@@ -48,6 +48,15 @@ class BaseWeightBroadcastConfig(BaseConfig):
 class FileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig):
     type: Literal["filesystem"] = "filesystem"
 
+    sparse: bool = False
+    """Between anchors, publish only the XOR delta of the wire-dtype weights against the previous
+    broadcast. Each trainer rank diffs its own shards; engines decode the delta while serving and
+    patch their weights in place during the pause. Not supported with LoRA."""
+
+    anchor_interval: int = Field(25, ge=1)
+    """With ``sparse``, publish a full checkpoint every this many steps (and at the first broadcast).
+    Engines that fall off the delta chain reload the last anchor and replay the deltas since."""
+
 
 class InMemoryWeightBroadcastConfig(BaseWeightBroadcastConfig):
     host: str = "localhost"
diff --git a/src/prime_rl/inference/vllm/server.py b/src/prime_rl/inference/vllm/server.py
index 7fea3bb4e6c..394c1a66993 100644
--- a/src/prime_rl/inference/vllm/server.py
+++ b/src/prime_rl/inference/vllm/server.py
@@ -60,6 +60,18 @@ async def update_weights(request: Request):
     return {"status": "ok"}
 
 
+@router.post("/prefetch_weights")
+async def prefetch_weights(request: Request):
+    """Start decoding a sparse weight delta in the workers and wait for it without
+    blocking generation: the workers decode on a background thread."""
+    data = await request.json()
+    client = engine_client(request)
+    await client.collective_rpc("prefetch_weights_from_path", args=(data.get("weight_dir"),))
+    while not all(await client.collective_rpc("prefetch_weights_done")):
+        await asyncio.sleep(0.05)
+    return {"status": "ok"}
+
+
 @router.post("/load_lora_adapter")
 async def load_lora_adapter(lora_request: LoadLoRAAdapterRequest, raw_request: Request):
     """Wrapper around vLLM's /v1/load_lora_adapter.
diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 53471969df4..9bf0ba5e931 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -1,9 +1,21 @@
+import itertools
+import time
+from concurrent.futures import Future, ThreadPoolExecutor
+from pathlib import Path
 from typing import TYPE_CHECKING
 
+import torch
 from torch.nn import Module
+from vllm.logger import init_logger
 from vllm.model_executor.model_loader import DefaultModelLoader, get_model_loader
 
-from prime_rl.inference.vllm.worker.weight_transfer import load_weights_checkpoint_layerwise
+from prime_rl.inference.vllm.worker.weight_transfer import (
+    CheckpointIndexMap,
+    load_weights_checkpoint_layerwise,
+    update_mla_absorbed_weights,
+    weight_checksums,
+)
+from prime_rl.utils.sparse_weights import DELTA_DIR, MANIFEST, bits_dtype, decode_gaps, read_delta, read_manifest
 
 # This is to get type hints for the Worker class but not actually extend it at runtime as this is required by vLLM worker extension
 if TYPE_CHECKING:
@@ -13,30 +25,116 @@
 else:
     Worker = object
 
+logger = init_logger("vllm.inference.vllm.worker_filesystem")
+
+_prefetch_pool = ThreadPoolExecutor(1)
+
 
 class FileSystemWeightUpdateWorker(Worker):
-    """vLLM worker extension for updating weights in-place using shared filesystem."""
+    """vLLM worker extension for updating weights in-place using shared filesystem.
+
+    A broadcast step dir holds a full HF checkpoint (anchor), a sparse delta
+    against an earlier step, or both. The worker tracks which step it holds and
+    applies a delta only on top of its base step; otherwise it reloads the
+    chain's anchor and replays the deltas since."""
+
+    weight_step: int | None = None
+    prefetched: tuple[str, Future] | None = None
+    index_map: CheckpointIndexMap | None = None
 
     def liveness_probe(self) -> None:
         """No-op RPC used by the API server liveness endpoint."""
         return None
 
-    def update_weights_from_path(self, weight_path: str) -> None:
-        """Update weights from a specified path in shared filesystem containing a HF-compatible checkpoint."""
-        # Get vLLM model runner and model
+    def _model(self) -> Module:
         # When enforce_eager=True, model isn't wrapped by torch.compile so no .runnable attr
-        model_runner = self.model_runner
-        if hasattr(model_runner.model, "runnable"):
-            model = model_runner.model.runnable
-        else:
-            model = model_runner.model
+        model = self.model_runner.model
+        if hasattr(model, "runnable"):
+            model = model.runnable
         assert isinstance(model, Module)
+        return model
+
+    def prefetch_weights_from_path(self, weight_path: str) -> None:
+        """Start decoding the sparse delta in ``weight_path`` into pinned host memory."""
+        self.prefetched = (weight_path, _prefetch_pool.submit(read_delta, Path(weight_path)))
+
+    def prefetch_weights_done(self) -> bool:
+        return self.prefetched is None or self.prefetched[1].done()
+
+    def update_weights_from_path(self, weight_path: str) -> None:
+        """Move to the broadcast step in ``weight_path`` (a full HF checkpoint and/or a sparse delta)."""
+        step_dir = Path(weight_path)
+        step = int(step_dir.name.removeprefix("step_"))
+        if not (step_dir / DELTA_DIR / MANIFEST).exists():
+            self._load_checkpoint(step_dir)
+        elif read_manifest(step_dir)["base_step"] == self.weight_step:
+            self._apply_delta(step_dir)
+            if any(step_dir.glob("*.safetensors")):
+                # Anchor step: check that the delta chain reproduced the full checkpoint bit for bit.
+                chained = weight_checksums(self._model())
+                self._load_checkpoint(step_dir)
+                mismatched = [name for name, value in weight_checksums(self._model()).items() if chained[name] != value]
+                if mismatched:
+                    raise RuntimeError(f"Sparse weight deltas diverged from the v{step} anchor: {mismatched[:8]}")
+                logger.info(f"Sparse weight delta chain matches the v{step} anchor")
+        else:
+            self._replay_from_anchor(step_dir)
+        self.weight_step = step
+
+    def _replay_from_anchor(self, step_dir: Path) -> None:
+        chain = [step_dir]
+        manifest = read_manifest(step_dir)
+        while manifest["base_step"] != manifest["anchor_step"]:
+            chain.append(step_dir.parent / f"step_{manifest['base_step']}")
+            manifest = read_manifest(chain[-1])
+        logger.info(f"Replaying {len(chain)} sparse weight deltas on top of the v{manifest['anchor_step']} anchor")
+        self._load_checkpoint(step_dir.parent / f"step_{manifest['anchor_step']}")
+        for delta_dir in reversed(chain):
+            self._apply_delta(delta_dir)
+
+    def _apply_delta(self, step_dir: Path) -> None:
+        start = time.perf_counter()
+        prefetched, self.prefetched = self.prefetched, None
+        if prefetched is not None and prefetched[0] == step_dir.as_posix():
+            delta = prefetched[1].result()
+        else:
+            delta = read_delta(step_dir)
+        read_time = time.perf_counter() - start
+        tensors = read_manifest(step_dir)["tensors"]
+        model = self._model()
+        device = next(model.parameters()).device
+        if self.index_map is None:
+            self.index_map = CheckpointIndexMap(model, tensors)
+            logger.info(f"Built the checkpoint index map in {time.perf_counter() - start - read_time:.2f}s")
+        for entries, gaps, values in delta:
+            counts = torch.tensor([count for _, count in entries], device=device)
+            indices = decode_gaps(gaps.to(device, non_blocking=True), counts)
+            values = values.to(device, non_blocking=True)
+            offset = num_bytes = 0
+            for bits, group in itertools.groupby(entries, key=lambda entry: bits_dtype(tensors[entry[0]][1])):
+                run = list(group)
+                num_changes = sum(count for _, count in run)
+                itemsize = bits.itemsize
+                xor = values[num_bytes : num_bytes + num_changes * itemsize]
+                # A run after an odd number of narrower values is misaligned for its dtype view.
+                xor = xor.view(bits) if num_bytes % itemsize == 0 else xor.clone().view(bits)
+                self.index_map.apply_xor(run, indices[offset : offset + num_changes], xor)
+                offset += num_changes
+                num_bytes += num_changes * itemsize
+        update_mla_absorbed_weights(model)
+        torch.cuda.synchronize()
+        logger.info(
+            f"Applied sparse weight delta {step_dir.name}: {sum(len(gaps) for _, gaps, _ in delta)} changes, "
+            f"read {read_time:.2f}s (0 if prefetched), apply {time.perf_counter() - start - read_time:.2f}s"
+        )
 
-        # Get vLLM model loader
+    def _load_checkpoint(self, step_dir: Path) -> None:
+        """Reload every weight from the HF-compatible checkpoint in ``step_dir``."""
+        model = self._model()
         model_loader = get_model_loader(self.load_config)
         assert isinstance(model_loader, DefaultModelLoader)
         local_source = DefaultModelLoader.Source(
-            weight_path,
+            step_dir.as_posix(),
             revision=None,  # TODO: Check that this is correct or if we should use the default (model_config.revision)
             prefix="",
             fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True),
diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py
index dadbad95469..e861d7233ca 100644
--- a/src/prime_rl/inference/vllm/worker/weight_transfer.py
+++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py
@@ -1,3 +1,4 @@
+from collections import defaultdict
 from typing import Iterable
 
 import torch
@@ -6,6 +7,8 @@
 from vllm.logger import init_logger
 from vllm.model_executor.model_loader.reload import finalize_layerwise_reload, initialize_layerwise_reload
 
+from prime_rl.utils.sparse_weights import as_bits
+
 logger = init_logger("vllm.inference.vllm.worker_weight_transfer")
 
 
@@ -23,6 +26,135 @@ def load_weights_checkpoint_layerwise(
         finalize_layerwise_reload(model, model_config)
 
 
+class CheckpointIndexMap:
+    """Maps flat checkpoint (HF) indices to flat indices of the live vLLM params.
+
+    Built once by running the model's own loader on placeholder checkpoint
+    tensors with ``Tensor.copy_`` replaced by a recorder, so it covers the
+    loader's name mapping, packing (qkv, gate_up, experts) and TP/EP slicing.
+    Supported: every write is one ``copy_`` of a row/column block of a 1-D or
+    2-D checkpoint tensor into a param of the same dtype, with no post-load
+    transform of the weights.
+    """
+
+    def __init__(self, model: Module, tensors: dict[str, tuple[list[int], str]]):
+        device = next(model.parameters()).device
+        params: dict[int, torch.Tensor] = {}
+        for param in model.parameters():
+            if param.storage_offset() == 0 and param.is_contiguous():
+                params.setdefault(param.untyped_storage().data_ptr(), param)
+        self.params: list[torch.Tensor] = list(params.values())
+        param_index = {ptr: index for index, ptr in enumerate(params)}
+        source_names: dict[int, str] = {}
+        # name -> [(row0, rows, col0, cols, param index, offset, row stride, col stride)]
+        records: dict[str, list[tuple[int, ...]]] = defaultdict(list)
+
+        def record_copy_(destination: torch.Tensor, source: torch.Tensor, non_blocking: bool = False) -> torch.Tensor:
+            name = source_names.get(source.untyped_storage().data_ptr())
+            index = param_index.get(destination.untyped_storage().data_ptr())
+            shape = tensors[name][0] if name is not None else None
+            num_cols = shape[-1] if shape else 1
+            source_2d = source.view(1, -1) if source.dim() == 1 else source
+            destination_2d = destination.view(1, -1) if destination.dim() == 1 else destination
+            if (
+                name is None
+                or index is None
+                or len(shape) > 2
+                or source.dtype != destination.dtype
+                or source.shape != destination.shape
+                or source.dim() > 2
+                or source_2d.stride(-1) != 1
+                or (source_2d.shape[0] > 1 and source_2d.stride(0) != num_cols)
+            ):
+                raise NotImplementedError("Sparse weight deltas need plain block copies from checkpoint tensors")
+            row0, col0 = divmod(source.storage_offset(), num_cols)
+            records[name].append(
+                (row0, source_2d.shape[0], col0, source_2d.shape[1], index, destination.storage_offset())
+                + tuple(destination_2d.stride())
+            )
+            return destination
+
+        def placeholders():
+            for name, (shape, dtype) in tensors.items():
+                tensor = torch.empty(shape, dtype=getattr(torch, dtype), device=device)
+                source_names[tensor.untyped_storage().data_ptr()] = name
+                yield name, tensor
+
+        original_copy = torch.Tensor.copy_
+        torch.Tensor.copy_ = record_copy_
+        try:
+            model.load_weights(placeholders())  # type: ignore
+        finally:
+            torch.Tensor.copy_ = original_copy
+
+        self.num_cols = {name: (shape[-1] if shape else 1) for name, (shape, _) in tensors.items()}
+        self.records = dict(records)
+
+    @torch.no_grad()
+    def apply_xor(
+        self, entries: list[tuple[str, int]], indices: torch.Tensor, xor: torch.Tensor, chunk: int = 1 << 25
+    ) -> None:
+        """XOR ``xor`` (bits) into the params at the checkpoint ``indices`` of ``entries``."""
+        device = indices.device
+        max_records = max((len(self.records.get(name, ())) for name, _ in entries), default=0)
+        # Per entry and write k: [num_cols, row0, rows, col0, cols, param, offset, row stride, col stride].
+        tables = [
+            torch.tensor(
+                [
+                    [self.num_cols[name], *self.records[name][k]]
+                    if len(self.records.get(name, ())) > k
+                    else [1, 0, 0, 0, 0, 0, 0, 0, 0]
+                    for name, _ in entries
+                ],
+                dtype=torch.int64,
+                device=device,
+            )
+            for k in range(max_records)
+        ]
+        entry_ends = torch.tensor([count for _, count in entries], device=device).cumsum(0)
+        for start in range(0, indices.numel(), chunk):
+            index = indices[start : start + chunk]
+            source = xor[start : start + chunk]
+            entry = torch.searchsorted(
+                entry_ends, torch.arange(start, start + index.numel(), device=device), right=True
+            )
+            for table in tables:
+                num_cols = table[entry, 0]
+                row = index // num_cols - table[entry, 1]
+                col = index % num_cols - table[entry, 3]
+                hit = (row >= 0) & (row < table[entry, 2]) & (col >= 0) & (col < table[entry, 4])
+                target = (table[entry, 6] + row * table[entry, 7] + col * table[entry, 8])[hit]
+                owner = table[entry, 5][hit]
+                self._scatter_xor(owner, target, source[hit])
+
+    def _scatter_xor(self, owner: torch.Tensor, target: torch.Tensor, source: torch.Tensor) -> None:
+        order = torch.argsort(owner, stable=True)
+        owner, target, source = owner[order], target[order], source[order]
+        counts = torch.bincount(owner, minlength=len(self.params)).tolist()
+        end = 0
+        for index, count in enumerate(counts):
+            begin, end = end, end + count
+            if not count:
+                continue
+            flat = as_bits(self.params[index].data).view(-1)
+            positions = target[begin:end]
+            flat.index_put_((positions,), flat[positions] ^ source[begin:end])
+
+
+@torch.no_grad()
+def weight_checksums(model: Module, chunk: int = 1 << 24) -> dict[str, int]:
+    """Position-weighted sum of every parameter's raw bits (exact, any dtype)."""
+    sums = {}
+    for name, param in model.named_parameters():
+        bits = as_bits(param.data).reshape(-1)
+        total = torch.zeros((), dtype=torch.int64, device=bits.device)
+        for start in range(0, bits.numel(), chunk):
+            part = bits[start : start + chunk].long()
+            total += (part * torch.arange(start + 1, start + 1 + part.numel(), device=bits.device)).sum()
+        sums[name] = total
+    return {name: int(total) for name, total in sums.items()}
+
+
 @torch.no_grad()
 def update_mla_absorbed_weights(model: Module) -> None:
     """Recompute MLA absorbed KV weights after in-place kv_b_proj updates."""
diff --git a/src/prime_rl/orchestrator/clients.py b/src/prime_rl/orchestrator/clients.py
index 46de44dd08e..956f8f892b4 100644
--- a/src/prime_rl/orchestrator/clients.py
+++ b/src/prime_rl/orchestrator/clients.py
@@ -175,6 +175,20 @@ async def initialize_client(admin_client: AsyncClient, rank_offset: int) -> None
             )
         )
 
+    async def prefetch_weights(self, weight_dir: Path) -> None:
+        """Have every engine read and decode a sparse weight delta while serving."""
+        await asyncio.gather(
+            *[
+                _admin_post(
+                    admin_client,
+                    "/prefetch_weights",
+                    json={"weight_dir": weight_dir.as_posix()},
+                    timeout_s=UPDATE_WEIGHTS_TIMEOUT_S,
+                )
+                for admin_client in self.clients
+            ]
+        )
+
     async def update_weights(
         self,
         weight_dir: Path | None,
diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index 9f139b9c5c7..31fe0607f5a 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -1,7 +1,14 @@
+import json
+import time
+from collections import defaultdict
+from math import prod
 from pathlib import Path
 
+import torch
+import torch.distributed as dist
 import torch.nn as nn
 from torch.distributed.tensor import DTensor
+from torch.distributed.tensor._utils import compute_local_shape_and_global_offset
 
 from prime_rl.configs.shared import FileSystemWeightBroadcastConfig
 from prime_rl.configs.trainer import LoRAConfig
@@ -9,17 +16,116 @@
 from prime_rl.trainer.lora import get_lora_state, save_lora_config
 from prime_rl.transports.weights.base import FINISHED_MARKER, WeightReceiver, WeightSender
 from prime_rl.utils.pathing import wait_for_path
+from prime_rl.utils.sparse_weights import DELTA_DIR, MANIFEST, as_bits, encode_gaps, write_rank_file
 from prime_rl.utils.weights import (
     convert_state_dict_to_hf,
     gather_weights_parallel,
+    resolve_fqn,
+    resolve_wire_dtype,
     save_state_dict,
     save_state_dict_parallel,
 )
 
 
+class SparseDeltaEncoder:
+    """Rank-local XOR deltas of the trainer's weight shards, in HF coordinates.
+
+    Every rank keeps the wire-dtype bits of its own shards from the previous
+    broadcast and diffs against them on device, so no weights are gathered.
+    The prime->HF conversion is replayed once on meta tensors: every HF tensor
+    must come out as a contiguous view of one trainer tensor (renames, expert
+    unstacking), which maps a changed trainer element to its HF name and index.
+    """
+
+    def __init__(self, model: nn.Module):
+        keep_in_fp32 = getattr(model, "keep_in_fp32_for_weight_transfer", None)
+        self.wire_dtypes: dict[str, torch.dtype] = {}
+        meta: dict[str, torch.Tensor] = {}
+        for key, value in model.state_dict().items():
+            dtype = resolve_wire_dtype(keep_in_fp32, key, torch.bfloat16) if isinstance(value, DTensor) else value.dtype
+            self.wire_dtypes[key] = dtype
+            meta[resolve_fqn(model, key)] = torch.empty(value.shape, dtype=dtype, device="meta")
+        owner = {id(tensor): fqn for fqn, tensor in meta.items()}
+
+        segments = defaultdict(list)
+        self.hf_tensors: dict[str, tuple[list[int], str]] = {}
+        for name, tensor in convert_state_dict_to_hf(model, dict(meta)).items():
+            base = tensor if tensor._base is None else tensor._base
+            if id(base) not in owner or not tensor.is_contiguous():
+                raise NotImplementedError(f"Sparse weight broadcast cannot map {name} to a contiguous trainer slice")
+            if tensor.numel() >= 2**32:
+                raise NotImplementedError(f"Sparse weight broadcast indexes tensors below 2^32 elements ({name})")
+            segments[owner[id(base)]].append((tensor.storage_offset(), name))
+            self.hf_tensors[name] = (list(tensor.shape), str(tensor.dtype).removeprefix("torch."))
+        self.segments = {fqn: sorted(segs) for fqn, segs in segments.items()}
+        self.segment_starts: dict[str, torch.Tensor] = {}
+        self.snapshot: dict[str, torch.Tensor] = {}
+        self.initialized = False
+
+    def _local_shard(self, key: str, value: torch.Tensor, is_master: bool) -> tuple[torch.Tensor, int] | None:
+        """This rank's contiguous flat slice of ``value`` and its global flat offset,
+        or None when another rank owns the elements."""
+        if not isinstance(value, DTensor):
+            return (value.detach().reshape(-1), 0) if is_master else None
+        if all(placement.is_replicate() for placement in value.placements):
+            return (value.to_local().detach().reshape(-1), 0) if is_master else None
+        local_shape, global_offset = compute_local_shape_and_global_offset(
+            value.shape, value.device_mesh, value.placements
+        )
+        if tuple(local_shape[1:]) != tuple(value.shape[1:]) or any(global_offset[1:]):
+            raise NotImplementedError(f"Sparse weight broadcast needs dim-0 shards ({key}: {value.placements})")
+        local = value.to_local().detach()[: local_shape[0]]
+        return local.reshape(-1), global_offset[0] * prod(value.shape[1:])
+
+    @torch.no_grad()
+    def encode(
+        self, model: nn.Module, is_master: bool
+    ) -> tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor] | None:
+        """Diff the current shards against the snapshot and advance it. Returns
+        ``(entries, gaps, xor value bytes)`` on the host, or None on the first call."""
+        first = not self.initialized
+        self.initialized = True
+        names, counts, gaps, values = [], [], [], []
+        for key, value in model.state_dict().items():
+            fqn = resolve_fqn(model, key)
+            shard = self._local_shard(key, value, is_master) if fqn in self.segments else None
+            if shard is None or shard[0].numel() == 0:
+                continue
+            local, offset = shard
+            bits = as_bits(local.to(self.wire_dtypes[key]))
+            if first:
+                self.snapshot[key] = bits.clone()
+                continue
+            previous = self.snapshot[key]
+            xor = bits ^ previous
+            previous.copy_(bits)
+            changed = xor.nonzero().squeeze(1)
+            if fqn not in self.segment_starts:
+                self.segment_starts[fqn] = torch.tensor([start for start, _ in self.segments[fqn]], device=xor.device)
+            starts = self.segment_starts[fqn]
+            position = changed + offset
+            segment = torch.searchsorted(starts, position, right=True) - 1
+            entry_starts = torch.ones_like(segment, dtype=torch.bool)
+            entry_starts[1:] = segment[1:] != segment[:-1]
+            gaps.append(encode_gaps(position - starts[segment], entry_starts))
+            values.append(xor[changed].view(torch.uint8))
+            counts.append(torch.bincount(segment, minlength=len(starts)))
+            names.extend(name for _, name in self.segments[fqn])
+        if first:
+            return None
+        counts = torch.cat(counts).tolist() if counts else []
+        entries = [(name, count) for name, count in zip(names, counts) if count]
+        empty = torch.empty(0, dtype=torch.uint8)
+        gaps_host = torch.cat(gaps).cpu() if gaps else empty
+        values_host = torch.cat(values).cpu() if values else empty
+        return entries, gaps_host, values_host
+
+
 class FileSystemWeightSender(WeightSender):
     """Broadcast weights by saving a HF-compatible checkpoint (or, for LoRA
-    runs, the PEFT-shaped adapter) to a shared filesystem."""
+    runs, the PEFT-shaped adapter) to a shared filesystem. With ``sparse``,
+    steps between anchors only carry the XOR delta against the previous
+    broadcast (see ``prime_rl.utils.sparse_weights``)."""
 
     def __init__(
         self,
@@ -28,7 +134,14 @@ def __init__(
         lora_config: LoRAConfig | None = None,
     ):
         super().__init__(output_dir, config.timeout)
+        if config.sparse and lora_config is not None:
+            raise ValueError("Sparse weight broadcast does not support LoRA")
+        self.config = config
         self.lora_config = lora_config
+        self.encoder: SparseDeltaEncoder | None = None
+        self.last_step: int | None = None
+        self.anchor_step: int | None = None
+        self.num_broadcasts = 0
         self.logger.debug("Initialized filesystem weight broadcast")
 
     def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
@@ -50,11 +163,62 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
                     alpha=self.lora_config.alpha,
                     dropout=self.lora_config.dropout,
                 )
-        else:
+            return
+
+        # The second broadcast is an anchor too, so engines check the first delta right away.
+        self.num_broadcasts += 1
+        anchor = not self.config.sparse or self.num_broadcasts <= 2 or step % self.config.anchor_interval == 0
+        if self.config.sparse:
+            if self.encoder is None:
+                self.encoder = SparseDeltaEncoder(model)
+            # Anchors also carry the delta, so engines can check the chain against them.
+            self._save_delta(model, step, step_dir)
+            self.last_step = step
+        if anchor:
             state_dict = gather_weights_parallel(model)
             state_dict = convert_state_dict_to_hf(model, state_dict)
             self.logger.debug(f"Saving weights to {step_dir}")
             save_state_dict_parallel(state_dict, step_dir)
+            self.anchor_step = step
+
+    def _save_delta(self, model: nn.Module, step: int, step_dir: Path) -> None:
+        assert self.encoder is not None
+        start = time.perf_counter()
+        delta = self.encoder.encode(model, self.world.is_master)
+        if delta is None:
+            return
+        entries, gaps, values = delta
+        encode_time = time.perf_counter() - start
+        file_name = f"rank_{self.world.rank}.safetensors"
+        delta_dir = step_dir / DELTA_DIR
+        if self.world.is_master:
+            delta_dir.mkdir()
+        dist.barrier()
+        num_bytes = write_rank_file(delta_dir / file_name, entries, gaps.numpy(), values.numpy())
+        stats = (file_name, num_bytes, len(gaps), encode_time, time.perf_counter() - start)
+        all_stats: list = [None] * self.world.world_size
+        dist.all_gather_object(all_stats, stats)
+        if not self.world.is_master:
+            return
+        manifest = {
+            "step": step,
+            "base_step": self.last_step,
+            "anchor_step": self.anchor_step,
+            "files": [file_name for file_name, *_ in all_stats],
+            "tensors": self.encoder.hf_tensors,
+        }
+        (delta_dir / MANIFEST).write_text(json.dumps(manifest))
+        num_changes = sum(stat[2] for stat in all_stats)
+        num_elements = sum(prod(shape) for shape, _ in self.encoder.hf_tensors.values())
+        self.logger.info(
+            f"Sparse weight delta v{step} (base v{self.last_step}): {num_changes / num_elements:.3%} changed, "
+            f"{sum(stat[1] for stat in all_stats) / 2**20:.1f} MiB, "
+            f"encode {max(stat[3] for stat in all_stats):.2f}s, total {max(stat[4] for stat in all_stats):.2f}s"
+        )
+
+    def _clean(self, step: int) -> None:
+        # Sparse chains replay from their anchor, so keep everything since it.
+        super()._clean(step if self.anchor_step is None else min(step, self.anchor_step + 1))
 
 
 class FileSystemWeightReceiver(WeightReceiver):
@@ -70,5 +234,10 @@ async def receive(self, step: int) -> None:
         await wait_for_path(weights_dir / FINISHED_MARKER)
         if (weights_dir / "adapter_config.json").exists():
             await load_lora_adapter(self.admin_plane, self.model_name, weights_dir)
-        else:
-            await self.admin_plane.update_weights(weights_dir, transport="filesystem", step=step)
+            return
+        if (weights_dir / DELTA_DIR / MANIFEST).exists():
+            # Read and decode the delta while the engines keep serving; the pause only applies it.
+            start = time.perf_counter()
+            await self.admin_plane.prefetch_weights(weights_dir)
+            self.logger.debug(f"Prefetched sparse weight delta v{step} in {time.perf_counter() - start:.2f}s")
+        await self.admin_plane.update_weights(weights_dir, transport="filesystem", step=step)
diff --git a/src/prime_rl/utils/pathing.py b/src/prime_rl/utils/pathing.py
index a02c19f6006..fd169d7d6b9 100644
--- a/src/prime_rl/utils/pathing.py
+++ b/src/prime_rl/utils/pathing.py
@@ -435,15 +435,13 @@ def sync_wait_for_path(path: Path, interval: int = 1, log_interval: int = 10) ->
         wait_time += interval
 
 
-async def wait_for_path(path: Path, interval: int = 1, log_interval: int = 10) -> None:
+async def wait_for_path(path: Path, interval: float = 0.1, log_interval: float = 10) -> None:
     logger = get_logger()
-    wait_time = 0
     logger.debug(f"Waiting for path `{path}`")
-    while True:
-        if path.exists():
-            logger.debug(f"Found path `{path}`")
-            break
-        if wait_time % log_interval == 0 and wait_time > 0:  # Every log_interval seconds
-            logger.debug(f"Waiting for path `{path}` for {wait_time} seconds")
+    start = last_log = time.monotonic()
+    while not path.exists():
+        if time.monotonic() - last_log >= log_interval:
+            last_log = time.monotonic()
+            logger.debug(f"Waiting for path `{path}` for {last_log - start:.0f} seconds")
         await asyncio.sleep(interval)
-        wait_time += interval
+    logger.debug(f"Found path `{path}`")
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
new file mode 100644
index 00000000000..c5c804b0e26
--- /dev/null
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -0,0 +1,89 @@
+"""Sparse weight deltas for the filesystem weight broadcast.
+
+A delta lives in the ``delta/`` subdir of a broadcast step dir (next to the
+full checkpoint on anchor steps): one ``rank_{r}.safetensors`` per trainer rank
+and a ``manifest.json``, written last. A rank file stores two zstd streams over its
+entries ``(hf_name, count)``: ``gaps`` (uint32 gaps between sorted flat indices,
+restarted per entry) and ``values`` (the XOR of the new and the previous raw
+bits). XOR is exact for any dtype, but applying it twice reverts it, so a delta
+must be applied on top of exactly its ``base_step``.
+"""
+
+import json
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import numpy as np
+import pyzstd
+import torch
+from safetensors import safe_open
+from safetensors.numpy import load_file, save_file
+
+DELTA_DIR = "delta"
+MANIFEST = "manifest.json"
+
+_BITS_DTYPE = {1: torch.uint8, 2: torch.int16, 4: torch.int32, 8: torch.int64}
+
+
+def as_bits(tensor: torch.Tensor) -> torch.Tensor:
+    return tensor.view(_BITS_DTYPE[tensor.element_size()])
+
+
+def bits_dtype(dtype: str) -> torch.dtype:
+    """The integer dtype holding the raw bits of the dtype named ``dtype``."""
+    return _BITS_DTYPE[getattr(torch, dtype).itemsize]
+
+
+def encode_gaps(indices: torch.Tensor, entry_starts: torch.Tensor) -> torch.Tensor:
+    """Sorted flat indices -> ``index - previous - 1``, restarting at every
+    ``entry_starts`` position (where the gap is the index itself)."""
+    previous = torch.empty_like(indices)
+    previous[1:] = indices[:-1]
+    previous[entry_starts] = -1
+    return (indices - previous - 1).to(torch.int32)
+
+
+def write_rank_file(path: Path, entries: list[tuple[str, int]], gaps: np.ndarray, values: np.ndarray) -> int:
+    options = {
+        pyzstd.CParameter.compressionLevel: 1,
+        pyzstd.CParameter.nbWorkers: 8,
+        pyzstd.CParameter.checksumFlag: 1,
+    }
+    arrays = {
+        "gaps": np.frombuffer(pyzstd.compress(gaps.tobytes(), options), np.uint8),
+        "values": np.frombuffer(pyzstd.compress(values.tobytes(), options), np.uint8),
+    }
+    save_file(arrays, str(path), metadata={"entries": json.dumps(entries)})
+    return sum(array.nbytes for array in arrays.values())
+
+
+def decode_gaps(gaps: torch.Tensor, counts: torch.Tensor) -> torch.Tensor:
+    """Inverse of ``encode_gaps`` for consecutive entries of ``counts`` indices each."""
+    indices = torch.cumsum(gaps.long() + 1, 0) - 1
+    entry_starts = torch.cumsum(counts, 0) - counts
+    restart = torch.where(entry_starts > 0, indices[(entry_starts - 1).clamp(min=0)] + 1, 0)
+    return indices - torch.repeat_interleave(restart, counts)
+
+
+def _read_rank_file(path: Path) -> tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]:
+    with safe_open(str(path), framework="np") as f:
+        entries = json.loads(f.metadata()["entries"])
+    arrays = load_file(str(path))
+    gaps = torch.frombuffer(bytearray(pyzstd.decompress(arrays["gaps"])), dtype=torch.int32)
+    values = torch.frombuffer(bytearray(pyzstd.decompress(arrays["values"])), dtype=torch.uint8)
+    if torch.cuda.is_available():
+        gaps, values = gaps.pin_memory(), values.pin_memory()
+    return entries, gaps, values
+
+
+def read_manifest(step_dir: Path) -> dict:
+    return json.loads((step_dir / DELTA_DIR / MANIFEST).read_text())
+
+
+def read_delta(step_dir: Path) -> list[tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]]:
+    """Read and decompress every rank file of a delta step dir into ``(entries,
+    gaps, xor value bytes)`` in (pinned) host memory; ``decode_gaps`` turns the
+    gaps into flat indices on the device that applies them."""
+    files = read_manifest(step_dir)["files"]
+    with ThreadPoolExecutor(len(files)) as pool:
+        return list(pool.map(_read_rank_file, [step_dir / DELTA_DIR / name for name in files]))
diff --git a/tests/unit/utils/test_sparse_weights.py b/tests/unit/utils/test_sparse_weights.py
new file mode 100644
index 00000000000..f2a3faad76c
--- /dev/null
+++ b/tests/unit/utils/test_sparse_weights.py
@@ -0,0 +1,54 @@
+import json
+
+import torch
+
+from prime_rl.utils.sparse_weights import (
+    DELTA_DIR,
+    MANIFEST,
+    as_bits,
+    decode_gaps,
+    encode_gaps,
+    read_delta,
+    write_rank_file,
+)
+
+
+def test_sparse_delta_roundtrip_is_exact(tmp_path):
+    torch.manual_seed(0)
+    previous = {"a": torch.randn(64, 32, dtype=torch.bfloat16), "b": torch.randn(1000)}
+    current = {name: tensor.clone() for name, tensor in previous.items()}
+    current["a"].view(-1)[torch.randperm(2048)[:100]] += 1
+    current["b"].view(-1)[torch.randperm(1000)[:10]] *= 2
+
+    # Two ranks, each holding a row slice of "a"; rank 1 also holds "b".
+    delta_dir = tmp_path / DELTA_DIR
+    delta_dir.mkdir()
+    for rank, names in enumerate([[("a", slice(0, 32))], [("a", slice(32, 64)), ("b", slice(None))]]):
+        entries, gaps, values = [], [], []
+        for name, rows in names:
+            offset = 0 if rows.start is None else rows.start * 32
+            xor = as_bits(current[name][rows]).reshape(-1) ^ as_bits(previous[name][rows]).reshape(-1)
+            changed = xor.nonzero().squeeze(1)
+            starts = torch.zeros_like(changed, dtype=torch.bool)
+            starts[:1] = True
+            entries.append((name, len(changed)))
+            gaps.append(encode_gaps(changed + offset, starts))
+            values.append(xor[changed].view(torch.uint8))
+        write_rank_file(
+            delta_dir / f"rank_{rank}.safetensors", entries, torch.cat(gaps).numpy(), torch.cat(values).numpy()
+        )
+    tensors = {name: [list(tensor.shape), str(tensor.dtype).removeprefix("torch.")] for name, tensor in current.items()}
+    files = ["rank_0.safetensors", "rank_1.safetensors"]
+    (delta_dir / MANIFEST).write_text(json.dumps({"files": files, "tensors": tensors}))
+
+    patched = {name: as_bits(tensor.clone()).view(-1) for name, tensor in previous.items()}
+    for entries, gaps, values in read_delta(tmp_path):
+        indices = decode_gaps(gaps, torch.tensor([count for _, count in entries]))
+        i = b = 0
+        for name, count in entries:
+            itemsize = patched[name].element_size()
+            patched[name][indices[i : i + count]] ^= values[b : b + count * itemsize].clone().view(patched[name].dtype)
+            i += count
+            b += count * itemsize
+    for name, tensor in current.items():
+        assert torch.equal(patched[name], as_bits(tensor).view(-1))

From 41836bb13535a5a9248ef5d008cc568a4fa19838 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 14 Oct 2026 13:24:58 +0000
Subject: [PATCH 02/88] feat(weights): engine workers prefetch only the sparse
 delta entries and index ranges they hold

---
 .../inference/vllm/worker/filesystem.py       |  8 +-
 .../inference/vllm/worker/weight_transfer.py  |  8 ++
 src/prime_rl/transports/weights/filesystem.py |  4 +-
 src/prime_rl/utils/sparse_weights.py          | 75 ++++++++++++++-----
 tests/unit/utils/test_sparse_weights.py       | 35 ++++++---
 5 files changed, 96 insertions(+), 34 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 9bf0ba5e931..1d7fc12ca19 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -56,7 +56,11 @@ def _model(self) -> Module:
 
     def prefetch_weights_from_path(self, weight_path: str) -> None:
         """Start decoding the sparse delta in ``weight_path`` into pinned host memory."""
-        self.prefetched = (weight_path, _prefetch_pool.submit(read_delta, Path(weight_path)))
+        self.prefetched = (weight_path, _prefetch_pool.submit(read_delta, Path(weight_path), self._spans()))
+
+    def _spans(self) -> dict[str, tuple[int, int]] | None:
+        """The checkpoint index ranges this worker's params hold, once the index map exists."""
+        return None if self.index_map is None else self.index_map.spans
 
     def prefetch_weights_done(self) -> bool:
         return self.prefetched is None or self.prefetched[1].done()
@@ -98,7 +102,7 @@ def _apply_delta(self, step_dir: Path) -> None:
         if prefetched is not None and prefetched[0] == step_dir.as_posix():
             delta = prefetched[1].result()
         else:
-            delta = read_delta(step_dir)
+            delta = read_delta(step_dir, self._spans())
         read_time = time.perf_counter() - start
         tensors = read_manifest(step_dir)["tensors"]
         model = self._model()
diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py
index e861d7233ca..c88ceda7d9f 100644
--- a/src/prime_rl/inference/vllm/worker/weight_transfer.py
+++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py
@@ -89,6 +89,14 @@ def placeholders():
 
         self.num_cols = {name: (shape[-1] if shape else 1) for name, (shape, _) in tensors.items()}
         self.records = dict(records)
+        # [first, last + 1) flat checkpoint index each tensor's writes cover
+        self.spans = {
+            name: (
+                min(row0 * self.num_cols[name] + col0 for row0, _, col0, *_ in writes),
+                max((row0 + rows - 1) * self.num_cols[name] + col0 + cols for row0, rows, col0, cols, *_ in writes),
+            )
+            for name, writes in self.records.items()
+        }
 
     @torch.no_grad()
     def apply_xor(
diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index 31fe0607f5a..363bb9530ec 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -194,7 +194,9 @@ def _save_delta(self, model: nn.Module, step: int, step_dir: Path) -> None:
         if self.world.is_master:
             delta_dir.mkdir()
         dist.barrier()
-        num_bytes = write_rank_file(delta_dir / file_name, entries, gaps.numpy(), values.numpy())
+        num_bytes = write_rank_file(
+            delta_dir / file_name, entries, gaps.numpy(), values.numpy(), self.encoder.hf_tensors
+        )
         stats = (file_name, num_bytes, len(gaps), encode_time, time.perf_counter() - start)
         all_stats: list = [None] * self.world.world_size
         dist.all_gather_object(all_stats, stats)
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index c5c804b0e26..e5a3579c6a3 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -17,7 +17,7 @@
 import pyzstd
 import torch
 from safetensors import safe_open
-from safetensors.numpy import load_file, save_file
+from safetensors.numpy import save_file
 
 DELTA_DIR = "delta"
 MANIFEST = "manifest.json"
@@ -43,17 +43,24 @@ def encode_gaps(indices: torch.Tensor, entry_starts: torch.Tensor) -> torch.Tens
     return (indices - previous - 1).to(torch.int32)
 
 
-def write_rank_file(path: Path, entries: list[tuple[str, int]], gaps: np.ndarray, values: np.ndarray) -> int:
-    options = {
-        pyzstd.CParameter.compressionLevel: 1,
-        pyzstd.CParameter.nbWorkers: 8,
-        pyzstd.CParameter.checksumFlag: 1,
-    }
+def write_rank_file(
+    path: Path, entries: list[tuple[str, int]], gaps: np.ndarray, values: np.ndarray, tensors: dict
+) -> int:
+    """Write every entry's gaps and values as separate zstd frames, so readers can skip entries."""
+    options = {pyzstd.CParameter.compressionLevel: 1, pyzstd.CParameter.checksumFlag: 1}
+    counts = [count for _, count in entries]
+    value_counts = [count * bits_dtype(tensors[name][1]).itemsize for name, count in entries]
+    gap_parts = np.split(gaps, np.cumsum(counts)[:-1])
+    value_parts = np.split(values, np.cumsum(value_counts)[:-1])
+    with ThreadPoolExecutor(8) as pool:
+        gap_frames = list(pool.map(lambda part: pyzstd.compress(part.tobytes(), options), gap_parts))
+        value_frames = list(pool.map(lambda part: pyzstd.compress(part.tobytes(), options), value_parts))
     arrays = {
-        "gaps": np.frombuffer(pyzstd.compress(gaps.tobytes(), options), np.uint8),
-        "values": np.frombuffer(pyzstd.compress(values.tobytes(), options), np.uint8),
+        "gaps": np.frombuffer(b"".join(gap_frames), np.uint8),
+        "values": np.frombuffer(b"".join(value_frames), np.uint8),
     }
-    save_file(arrays, str(path), metadata={"entries": json.dumps(entries)})
+    frames = [[name, count, len(g), len(v)] for (name, count), g, v in zip(entries, gap_frames, value_frames)]
+    save_file(arrays, str(path), metadata={"entries": json.dumps(frames)})
     return sum(array.nbytes for array in arrays.values())
 
 
@@ -65,12 +72,36 @@ def decode_gaps(gaps: torch.Tensor, counts: torch.Tensor) -> torch.Tensor:
     return indices - torch.repeat_interleave(restart, counts)
 
 
-def _read_rank_file(path: Path) -> tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]:
+def _read_rank_file(
+    path: Path, tensors: dict, spans: dict[str, tuple[int, int]] | None
+) -> tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]:
     with safe_open(str(path), framework="np") as f:
-        entries = json.loads(f.metadata()["entries"])
-    arrays = load_file(str(path))
-    gaps = torch.frombuffer(bytearray(pyzstd.decompress(arrays["gaps"])), dtype=torch.int32)
-    values = torch.frombuffer(bytearray(pyzstd.decompress(arrays["values"])), dtype=torch.uint8)
+        frames = json.loads(f.metadata()["entries"])
+        gaps_bytes, values_bytes = f.get_slice("gaps"), f.get_slice("values")
+        entries, gaps, values = [], [], []
+        gaps_offset = values_offset = 0
+        for name, count, num_gap_bytes, num_value_bytes in frames:
+            gaps_end, values_end = gaps_offset + num_gap_bytes, values_offset + num_value_bytes
+            if spans is None or name in spans:
+                entry_gaps = np.frombuffer(pyzstd.decompress(gaps_bytes[gaps_offset:gaps_end].tobytes()), np.int32)
+                entry_values = np.frombuffer(
+                    pyzstd.decompress(values_bytes[values_offset:values_end].tobytes()), np.uint8
+                )
+                if spans is not None:
+                    indices = np.cumsum(entry_gaps, dtype=np.int64) + np.arange(1, count + 1) - 1
+                    begin, end = np.searchsorted(indices, spans[name])
+                    itemsize = bits_dtype(tensors[name][1]).itemsize
+                    entry_values = entry_values[begin * itemsize : end * itemsize]
+                    entry_gaps = entry_gaps[begin:end].copy()
+                    if end > begin:
+                        entry_gaps[0] = indices[begin]
+                if len(entry_gaps):
+                    entries.append((name, len(entry_gaps)))
+                    gaps.append(entry_gaps)
+                    values.append(entry_values)
+            gaps_offset, values_offset = gaps_end, values_end
+    gaps = torch.from_numpy(np.concatenate(gaps) if gaps else np.empty(0, np.int32))
+    values = torch.from_numpy(np.concatenate(values) if values else np.empty(0, np.uint8))
     if torch.cuda.is_available():
         gaps, values = gaps.pin_memory(), values.pin_memory()
     return entries, gaps, values
@@ -80,10 +111,14 @@ def read_manifest(step_dir: Path) -> dict:
     return json.loads((step_dir / DELTA_DIR / MANIFEST).read_text())
 
 
-def read_delta(step_dir: Path) -> list[tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]]:
+def read_delta(
+    step_dir: Path, spans: dict[str, tuple[int, int]] | None = None
+) -> list[tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]]:
     """Read and decompress every rank file of a delta step dir into ``(entries,
     gaps, xor value bytes)`` in (pinned) host memory; ``decode_gaps`` turns the
-    gaps into flat indices on the device that applies them."""
-    files = read_manifest(step_dir)["files"]
-    with ThreadPoolExecutor(len(files)) as pool:
-        return list(pool.map(_read_rank_file, [step_dir / DELTA_DIR / name for name in files]))
+    gaps into flat indices on the device that applies them. With ``spans``, keep
+    only the listed tensors and, within each, the flat indices in ``[start, end)``."""
+    manifest = read_manifest(step_dir)
+    paths = [step_dir / DELTA_DIR / name for name in manifest["files"]]
+    with ThreadPoolExecutor(len(paths)) as pool:
+        return list(pool.map(lambda path: _read_rank_file(path, manifest["tensors"], spans), paths))
diff --git a/tests/unit/utils/test_sparse_weights.py b/tests/unit/utils/test_sparse_weights.py
index f2a3faad76c..9841f607575 100644
--- a/tests/unit/utils/test_sparse_weights.py
+++ b/tests/unit/utils/test_sparse_weights.py
@@ -20,6 +20,7 @@ def test_sparse_delta_roundtrip_is_exact(tmp_path):
     current["a"].view(-1)[torch.randperm(2048)[:100]] += 1
     current["b"].view(-1)[torch.randperm(1000)[:10]] *= 2
 
+    tensors = {name: [list(tensor.shape), str(tensor.dtype).removeprefix("torch.")] for name, tensor in current.items()}
     # Two ranks, each holding a row slice of "a"; rank 1 also holds "b".
     delta_dir = tmp_path / DELTA_DIR
     delta_dir.mkdir()
@@ -35,20 +36,32 @@ def test_sparse_delta_roundtrip_is_exact(tmp_path):
             gaps.append(encode_gaps(changed + offset, starts))
             values.append(xor[changed].view(torch.uint8))
         write_rank_file(
-            delta_dir / f"rank_{rank}.safetensors", entries, torch.cat(gaps).numpy(), torch.cat(values).numpy()
+            delta_dir / f"rank_{rank}.safetensors", entries, torch.cat(gaps).numpy(), torch.cat(values).numpy(), tensors
         )
-    tensors = {name: [list(tensor.shape), str(tensor.dtype).removeprefix("torch.")] for name, tensor in current.items()}
     files = ["rank_0.safetensors", "rank_1.safetensors"]
     (delta_dir / MANIFEST).write_text(json.dumps({"files": files, "tensors": tensors}))
 
-    patched = {name: as_bits(tensor.clone()).view(-1) for name, tensor in previous.items()}
-    for entries, gaps, values in read_delta(tmp_path):
-        indices = decode_gaps(gaps, torch.tensor([count for _, count in entries]))
-        i = b = 0
-        for name, count in entries:
-            itemsize = patched[name].element_size()
-            patched[name][indices[i : i + count]] ^= values[b : b + count * itemsize].clone().view(patched[name].dtype)
-            i += count
-            b += count * itemsize
+    def apply(spans):
+        patched = {name: as_bits(tensor.clone()).view(-1) for name, tensor in previous.items()}
+        for entries, gaps, values in read_delta(tmp_path, spans):
+            indices = decode_gaps(gaps, torch.tensor([count for _, count in entries]))
+            i = b = 0
+            for name, count in entries:
+                itemsize = patched[name].element_size()
+                patched[name][indices[i : i + count]] ^= (
+                    values[b : b + count * itemsize].clone().view(patched[name].dtype)
+                )
+                i += count
+                b += count * itemsize
+        return patched
+
+    patched = apply(None)
     for name, tensor in current.items():
         assert torch.equal(patched[name], as_bits(tensor).view(-1))
+
+    # A reader that holds rows 16-47 of "a" (spanning both rank files) and none of "b".
+    patched = apply({"a": (16 * 32, 48 * 32)})
+    expected = as_bits(previous["a"].clone())
+    expected[16:48] = as_bits(current["a"][16:48])
+    assert torch.equal(patched["a"], expected.view(-1))
+    assert torch.equal(patched["b"], as_bits(previous["b"]).view(-1))

From 87669512b53db7a1b3f1350166597905196803ef Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 14 Oct 2026 14:38:17 +0000
Subject: [PATCH 03/88] fix(weights): allow checkpoint tensors the engine casts
 on load (e.g. fp32 routers) as long as they never change

---
 .../inference/vllm/worker/weight_transfer.py   | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py
index c88ceda7d9f..d49df747dc9 100644
--- a/src/prime_rl/inference/vllm/worker/weight_transfer.py
+++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py
@@ -34,7 +34,8 @@ class CheckpointIndexMap:
     loader's name mapping, packing (qkv, gate_up, experts) and TP/EP slicing.
     Supported: every write is one ``copy_`` of a row/column block of a 1-D or
     2-D checkpoint tensor into a param of the same dtype, with no post-load
-    transform of the weights.
+    transform of the weights. Tensors the loader casts (e.g. an fp32 router
+    from a bf16 checkpoint) are only allowed to stay unchanged.
     """
 
     def __init__(self, model: Module, tensors: dict[str, tuple[list[int], str]]):
@@ -48,6 +49,7 @@ def __init__(self, model: Module, tensors: dict[str, tuple[list[int], str]]):
         source_names: dict[int, str] = {}
         # name -> [(row0, rows, col0, cols, param index, offset, row stride, col stride)]
         records: dict[str, list[tuple[int, ...]]] = defaultdict(list)
+        unpatchable: set[str] = set()
 
         def record_copy_(destination: torch.Tensor, source: torch.Tensor, non_blocking: bool = False) -> torch.Tensor:
             name = source_names.get(source.untyped_storage().data_ptr())
@@ -56,11 +58,14 @@ def record_copy_(destination: torch.Tensor, source: torch.Tensor, non_blocking:
             num_cols = shape[-1] if shape else 1
             source_2d = source.view(1, -1) if source.dim() == 1 else source
             destination_2d = destination.view(1, -1) if destination.dim() == 1 else destination
+            if name is not None and source.dtype != destination.dtype:
+                # e.g. an fp32 router loaded from a bf16 checkpoint: fine as long as it never changes
+                unpatchable.add(name)
+                return destination
             if (
                 name is None
                 or index is None
                 or len(shape) > 2
-                or source.dtype != destination.dtype
                 or source.shape != destination.shape
                 or source.dim() > 2
                 or source_2d.stride(-1) != 1
@@ -89,20 +94,25 @@ def placeholders():
 
         self.num_cols = {name: (shape[-1] if shape else 1) for name, (shape, _) in tensors.items()}
         self.records = dict(records)
-        # [first, last + 1) flat checkpoint index each tensor's writes cover
+        self.unpatchable = unpatchable
+        # [first, last + 1) flat checkpoint index each tensor's writes cover; unpatchable tensors are
+        # read whole so that a change to them is caught in apply_xor
         self.spans = {
             name: (
                 min(row0 * self.num_cols[name] + col0 for row0, _, col0, *_ in writes),
                 max((row0 + rows - 1) * self.num_cols[name] + col0 + cols for row0, rows, col0, cols, *_ in writes),
             )
             for name, writes in self.records.items()
-        }
+        } | {name: (0, 2**62) for name in unpatchable}
 
     @torch.no_grad()
     def apply_xor(
         self, entries: list[tuple[str, int]], indices: torch.Tensor, xor: torch.Tensor, chunk: int = 1 << 25
     ) -> None:
         """XOR ``xor`` (bits) into the params at the checkpoint ``indices`` of ``entries``."""
+        changed = [name for name, _ in entries if name in self.unpatchable]
+        if changed:
+            raise NotImplementedError(f"Sparse weight deltas cannot patch {changed[:8]}: the engine casts them on load")
         device = indices.device
         max_records = max((len(self.records.get(name, ())) for name, _ in entries), default=0)
         # Per entry and write k: [num_cols, row0, rows, col0, cols, param, offset, row stride, col stride].

From 52c4d7f22d8d3c1ff4d6de45863f1e259881b454 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 14 Oct 2026 14:05:58 +0000
Subject: [PATCH 04/88] feat(weights): fp8-domain sparse deltas for engines
 that quantize online to fp8_per_block

---
 .../inference/vllm/worker/filesystem.py       |  51 +++++-
 .../inference/vllm/worker/weight_transfer.py  | 129 +++++++++++++---
 src/prime_rl/transports/weights/filesystem.py | 145 +++++++++++++++---
 src/prime_rl/utils/sparse_weights.py          |  30 ++++
 4 files changed, 313 insertions(+), 42 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 1d7fc12ca19..68a11ce3830 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -1,11 +1,15 @@
 import itertools
+import json
+import os
 import time
 from concurrent.futures import Future, ThreadPoolExecutor
 from pathlib import Path
 from typing import TYPE_CHECKING
 
 import torch
+from safetensors import safe_open
 from torch.nn import Module
+from vllm.distributed.parallel_state import get_ep_group, get_world_group
 from vllm.logger import init_logger
 from vllm.model_executor.model_loader import DefaultModelLoader, get_model_loader
 
@@ -15,7 +19,15 @@
     update_mla_absorbed_weights,
     weight_checksums,
 )
-from prime_rl.utils.sparse_weights import DELTA_DIR, MANIFEST, bits_dtype, decode_gaps, read_delta, read_manifest
+from prime_rl.utils.sparse_weights import (
+    DELTA_DIR,
+    ENGINE_LAYOUT,
+    MANIFEST,
+    bits_dtype,
+    decode_gaps,
+    read_delta,
+    read_manifest,
+)
 
 # This is to get type hints for the Worker class but not actually extend it at runtime as this is required by vLLM worker extension
 if TYPE_CHECKING:
@@ -29,6 +41,8 @@
 
 _prefetch_pool = ThreadPoolExecutor(1)
 
+_SAFETENSORS_DTYPES = {"BF16": "bfloat16", "F16": "float16", "F32": "float32"}
+
 
 class FileSystemWeightUpdateWorker(Worker):
     """vLLM worker extension for updating weights in-place using shared filesystem.
@@ -41,6 +55,7 @@ class FileSystemWeightUpdateWorker(Worker):
     weight_step: int | None = None
     prefetched: tuple[str, Future] | None = None
     index_map: CheckpointIndexMap | None = None
+    wrote_layout = False
 
     def liveness_probe(self) -> None:
         """No-op RPC used by the API server liveness endpoint."""
@@ -84,6 +99,40 @@ def update_weights_from_path(self, weight_path: str) -> None:
         else:
             self._replay_from_anchor(step_dir)
         self.weight_step = step
+        if (step_dir / DELTA_DIR).is_dir() and not self.wrote_layout and any(step_dir.glob("*.safetensors")):
+            self._write_engine_layout(step_dir)
+            self.wrote_layout = True
+
+    def _write_engine_layout(self, step_dir: Path) -> None:
+        """Tell the trainer which checkpoint tensors this engine holds in fp8 (and with which
+        scales), so it can send their deltas in the fp8 domain. Needs a full checkpoint."""
+        files = {}
+        tensors = {}
+        for path in step_dir.glob("*.safetensors"):
+            with safe_open(path, framework="pt") as f:
+                for name in f.keys():
+                    tensor_slice = f.get_slice(name)
+                    files[name] = path
+                    tensors[name] = (tensor_slice.get_shape(), _SAFETENSORS_DTYPES[tensor_slice.get_dtype()])
+        if self.index_map is None:
+            self.index_map = CheckpointIndexMap(self._model(), tensors)
+
+        def read_tensor(name: str) -> torch.Tensor:
+            with safe_open(files[name], framework="pt") as f:
+                return f.get_tensor(name)
+
+        # Each worker reports the tensors it holds; the trainer waits for all of them.
+        if self.vllm_config.parallel_config.enable_expert_parallel:
+            group = get_ep_group()
+        else:
+            group = get_world_group()
+        layout = {"world_size": group.world_size, "fp8": self.index_map.fp8_layout(read_tensor)}
+        path = step_dir.parent / ENGINE_LAYOUT / f"rank_{group.rank_in_group}.json"
+        path.parent.mkdir(exist_ok=True)
+        tmp = path.with_suffix(f".tmp{os.getpid()}")
+        tmp.write_text(json.dumps(layout))
+        os.replace(tmp, path)
+        logger.info(f"Wrote the engine layout: {len(layout['fp8'])} checkpoint tensors held in fp8")
 
     def _replay_from_anchor(self, step_dir: Path) -> None:
         chain = [step_dir]
diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py
index d49df747dc9..566da796473 100644
--- a/src/prime_rl/inference/vllm/worker/weight_transfer.py
+++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py
@@ -6,8 +6,15 @@
 from vllm.config import set_current_vllm_config
 from vllm.logger import init_logger
 from vllm.model_executor.model_loader.reload import finalize_layerwise_reload, initialize_layerwise_reload
+from vllm.model_executor.model_loader.reload.layerwise import (
+    LAYERWISE_INFO,
+    _get_original_loader,
+    _get_weight_loader,
+)
+from vllm.model_executor.model_loader.reload.meta import restore_layer_on_meta
+from vllm.model_executor.model_loader.reload.utils import get_layer_params_buffers, get_layer_tensors
 
-from prime_rl.utils.sparse_weights import as_bits
+from prime_rl.utils.sparse_weights import as_bits, quantize_fp8_blocks
 
 logger = init_logger("vllm.inference.vllm.worker_weight_transfer")
 
@@ -32,28 +39,24 @@ class CheckpointIndexMap:
     Built once by running the model's own loader on placeholder checkpoint
     tensors with ``Tensor.copy_`` replaced by a recorder, so it covers the
     loader's name mapping, packing (qkv, gate_up, experts) and TP/EP slicing.
-    Supported: every write is one ``copy_`` of a row/column block of a 1-D or
-    2-D checkpoint tensor into a param of the same dtype, with no post-load
-    transform of the weights. Tensors the loader casts (e.g. an fp32 router
-    from a bf16 checkpoint) are only allowed to stay unchanged.
+    Layers quantized online (e.g. ``fp8_per_block``) are recorded in their
+    checkpoint format; their live params must keep that layout (one byte per
+    checkpoint element). Supported: every write is one ``copy_`` of a row/column
+    block of a 1-D or 2-D checkpoint tensor. Tensors the loader casts (e.g. an
+    fp32 router from a bf16 checkpoint) are only allowed to stay unchanged.
     """
 
     def __init__(self, model: Module, tensors: dict[str, tuple[list[int], str]]):
         device = next(model.parameters()).device
-        params: dict[int, torch.Tensor] = {}
-        for param in model.parameters():
-            if param.storage_offset() == 0 and param.is_contiguous():
-                params.setdefault(param.untyped_storage().data_ptr(), param)
-        self.params: list[torch.Tensor] = list(params.values())
-        param_index = {ptr: index for index, ptr in enumerate(params)}
         source_names: dict[int, str] = {}
-        # name -> [(row0, rows, col0, cols, param index, offset, row stride, col stride)]
-        records: dict[str, list[tuple[int, ...]]] = defaultdict(list)
+        # name -> [(row0, rows, col0, cols, param name, offset, row stride, col stride)]
+        records: dict[str, list[tuple]] = defaultdict(list)
+        shapes: dict[str, torch.Size] = {}
         unpatchable: set[str] = set()
 
         def record_copy_(destination: torch.Tensor, source: torch.Tensor, non_blocking: bool = False) -> torch.Tensor:
-            name = source_names.get(source.untyped_storage().data_ptr())
-            index = param_index.get(destination.untyped_storage().data_ptr())
+            name = source_names.get(source.untyped_storage()._cdata)
+            param_name = param_names.get(destination.untyped_storage()._cdata)
             shape = tensors[name][0] if name is not None else None
             num_cols = shape[-1] if shape else 1
             source_2d = source.view(1, -1) if source.dim() == 1 else source
@@ -64,7 +67,7 @@ def record_copy_(destination: torch.Tensor, source: torch.Tensor, non_blocking:
                 return destination
             if (
                 name is None
-                or index is None
+                or param_name is None
                 or len(shape) > 2
                 or source.shape != destination.shape
                 or source.dim() > 2
@@ -74,7 +77,7 @@ def record_copy_(destination: torch.Tensor, source: torch.Tensor, non_blocking:
                 raise NotImplementedError("Sparse weight deltas need plain block copies from checkpoint tensors")
             row0, col0 = divmod(source.storage_offset(), num_cols)
             records[name].append(
-                (row0, source_2d.shape[0], col0, source_2d.shape[1], index, destination.storage_offset())
+                (row0, source_2d.shape[0], col0, source_2d.shape[1], param_name, destination.storage_offset())
                 + tuple(destination_2d.stride())
             )
             return destination
@@ -82,18 +85,35 @@ def record_copy_(destination: torch.Tensor, source: torch.Tensor, non_blocking:
         def placeholders():
             for name, (shape, dtype) in tensors.items():
                 tensor = torch.empty(shape, dtype=getattr(torch, dtype), device=device)
-                source_names[tensor.untyped_storage().data_ptr()] = name
+                source_names[tensor.untyped_storage()._cdata] = name
                 yield name, tensor
 
+        swapped = _restore_checkpoint_format(model)
         original_copy = torch.Tensor.copy_
-        torch.Tensor.copy_ = record_copy_
         try:
+            param_names = {}
+            for name, param in model.named_parameters():
+                if param.storage_offset() == 0 and param.is_contiguous():
+                    param_names.setdefault(param.untyped_storage()._cdata, name)
+                    shapes[name] = param.shape
+            torch.Tensor.copy_ = record_copy_
             model.load_weights(placeholders())  # type: ignore
         finally:
             torch.Tensor.copy_ = original_copy
+            _restore_kernel_format(swapped)
 
+        live = dict(model.named_parameters())
+        used = list(dict.fromkeys(write[4] for writes in records.values() for write in writes))
+        for name in used:
+            if live[name].shape != shapes[name] or not live[name].is_contiguous() or live[name].storage_offset():
+                raise NotImplementedError(f"Sparse weight deltas cannot patch {name}: its layout changes after loading")
+        self.params: list[torch.Tensor] = [live[name] for name in used]
+        param_index = {name: index for index, name in enumerate(used)}
+        self.records = {
+            name: [(*write[:4], param_index[write[4]], *write[5:]) for write in writes]
+            for name, writes in records.items()
+        }
         self.num_cols = {name: (shape[-1] if shape else 1) for name, (shape, _) in tensors.items()}
-        self.records = dict(records)
         self.unpatchable = unpatchable
         # [first, last + 1) flat checkpoint index each tensor's writes cover; unpatchable tensors are
         # read whole so that a change to them is caught in apply_xor
@@ -105,6 +125,33 @@ def placeholders():
             for name, writes in self.records.items()
         } | {name: (0, 2**62) for name in unpatchable}
 
+    def fp8_layout(self, read_tensor) -> dict[str, str]:
+        """The vLLM per-block quantization (``"fp32"`` or ``"ue8m0"`` scales) that reproduces each
+        fp8 param from the checkpoint, per checkpoint tensor stored in one. Checks one checkpoint
+        tensor per param; ``read_tensor(name)`` returns it."""
+        modes: dict[int, str] = {}
+        layout = {}
+        for name, writes in self.records.items():
+            for row0, rows, col0, cols, index, offset, *strides in writes:
+                if self.params[index].dtype != torch.float8_e4m3fn:
+                    continue
+                if index not in modes:
+                    live = as_bits(self.params[index].data).view(-1).as_strided((rows, cols), strides, offset)
+                    weight = read_tensor(name).to(self.params[index].device)
+                    for mode in ("ue8m0", "fp32"):
+                        quantized, _ = quantize_fp8_blocks(weight, mode == "ue8m0")
+                        if torch.equal(quantized[row0 : row0 + rows, col0 : col0 + cols], live):
+                            modes[index] = mode
+                            break
+                    else:
+                        raise NotImplementedError(
+                            f"Sparse weight deltas cannot patch {name}: the engine's fp8 weights are not vLLM's "
+                            "per-block quantization of the checkpoint (e.g. a weight-shuffling kernel)"
+                        )
+                if layout.setdefault(name, modes[index]) != modes[index]:
+                    raise NotImplementedError(f"Sparse weight deltas cannot patch {name}: mixed fp8 layouts")
+        return layout
+
     @torch.no_grad()
     def apply_xor(
         self, entries: list[tuple[str, int]], indices: torch.Tensor, xor: torch.Tensor, chunk: int = 1 << 25
@@ -155,10 +202,52 @@ def _scatter_xor(self, owner: torch.Tensor, target: torch.Tensor, source: torch.
             if not count:
                 continue
             flat = as_bits(self.params[index].data).view(-1)
+            if flat.dtype != source.dtype:
+                raise RuntimeError(
+                    f"Sparse weight delta dtype {source.dtype} does not match the engine param ({flat.dtype})"
+                )
             positions = target[begin:end]
             flat.index_put_((positions,), flat[positions] ^ source[begin:end])
 
 
+def _restore_checkpoint_format(model: Module) -> list[tuple]:
+    """Put the checkpoint-format params (on meta, with their unwrapped loaders) back into
+    every layer whose live params were transformed after loading (online quantization)."""
+    swapped = []
+    for layer in model.modules():
+        info = LAYERWISE_INFO.get(layer)
+        if info is None:
+            continue
+        restore_params, _ = info.restore_metadata
+        live = layer._parameters
+        if all(
+            live.get(name) is not None and (live[name].dtype, live[name].shape) == (param.dtype, param.shape)
+            for name, param in restore_params.items()
+        ):
+            continue
+        kernel_tensors = get_layer_params_buffers(layer)
+        non_persistent = set(layer._non_persistent_buffers_set)
+        restore_layer_on_meta(layer, info)
+        loaders = {}
+        for name, param in get_layer_params_buffers(layer)[0].items():
+            loaders[name] = _get_weight_loader(param)
+            param.weight_loader = _get_original_loader(param)
+        swapped.append((layer, kernel_tensors, non_persistent, loaders))
+    return swapped
+
+
+def _restore_kernel_format(swapped: list[tuple]) -> None:
+    for layer, (params, buffers), non_persistent, loaders in swapped:
+        for name, loader in loaders.items():
+            getattr(layer, name).weight_loader = loader
+        for name in get_layer_tensors(layer):
+            delattr(layer, name)
+        for name, param in params.items():
+            layer.register_parameter(name, param)
+        for name, buffer in buffers.items():
+            layer.register_buffer(name, buffer, persistent=name not in non_persistent)
+
+
 @torch.no_grad()
 def weight_checksums(model: Module, chunk: int = 1 << 24) -> dict[str, int]:
     """Position-weighted sum of every parameter's raw bits (exact, any dtype)."""
diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index 363bb9530ec..77c55566abf 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -16,7 +16,16 @@
 from prime_rl.trainer.lora import get_lora_state, save_lora_config
 from prime_rl.transports.weights.base import FINISHED_MARKER, WeightReceiver, WeightSender
 from prime_rl.utils.pathing import wait_for_path
-from prime_rl.utils.sparse_weights import DELTA_DIR, MANIFEST, as_bits, encode_gaps, write_rank_file
+from prime_rl.utils.sparse_weights import (
+    DELTA_DIR,
+    FP8_BLOCK,
+    MANIFEST,
+    as_bits,
+    encode_gaps,
+    quantize_fp8_blocks,
+    read_engine_layout,
+    write_rank_file,
+)
 from prime_rl.utils.weights import (
     convert_state_dict_to_hf,
     gather_weights_parallel,
@@ -35,9 +44,15 @@ class SparseDeltaEncoder:
     The prime->HF conversion is replayed once on meta tensors: every HF tensor
     must come out as a contiguous view of one trainer tensor (renames, expert
     unstacking), which maps a changed trainer element to its HF name and index.
+
+    Tensors the engines hold in fp8 (``fp8``: HF name -> scale format, from the
+    engine layout) are quantized with vLLM's own per-block quantizer and diffed
+    as fp8 bytes. Their 128-row blocks must not straddle trainer ranks, else the
+    trainer tensor is gathered. A step that changes any block scale returns no
+    delta, so it is published as a full checkpoint.
     """
 
-    def __init__(self, model: nn.Module):
+    def __init__(self, model: nn.Module, fp8: dict[str, str]):
         keep_in_fp32 = getattr(model, "keep_in_fp32_for_weight_transfer", None)
         self.wire_dtypes: dict[str, torch.dtype] = {}
         meta: dict[str, torch.Tensor] = {}
@@ -56,10 +71,34 @@ def __init__(self, model: nn.Module):
             if tensor.numel() >= 2**32:
                 raise NotImplementedError(f"Sparse weight broadcast indexes tensors below 2^32 elements ({name})")
             segments[owner[id(base)]].append((tensor.storage_offset(), name))
-            self.hf_tensors[name] = (list(tensor.shape), str(tensor.dtype).removeprefix("torch."))
+            dtype = "float8_e4m3fn" if name in fp8 else str(tensor.dtype).removeprefix("torch.")
+            self.hf_tensors[name] = (list(tensor.shape), dtype)
         self.segments = {fqn: sorted(segs) for fqn, segs in segments.items()}
+
+        # fqn -> (ue8m0 scales, columns) for trainer tensors held in fp8; quantized as one [rows, columns] matrix
+        self.fp8: dict[str, tuple[bool, int]] = {}
+        for fqn, segs in self.segments.items():
+            modes = {fp8.get(name) for _, name in segs}
+            if modes == {None}:
+                continue
+            shapes = [self.hf_tensors[name][0] for _, name in segs]
+            num_cols = shapes[0][-1]
+            num_rows = meta[fqn].numel() // num_cols
+            boundaries = [start // num_cols for start, _ in segs] + [
+                start // num_cols + shape[0] for (start, _), shape in zip(segs, shapes)
+            ]
+            if (
+                len(modes) > 1
+                or any(len(shape) != 2 or shape[1] != num_cols for shape in shapes)
+                or any(row % FP8_BLOCK[0] and row != num_rows for row in boundaries)
+            ):
+                raise NotImplementedError(f"Sparse weight broadcast cannot quantize {fqn} to fp8 as one matrix")
+            self.fp8[fqn] = (modes.pop() == "ue8m0", num_cols)
+        self.gather: set[str] | None = None
         self.segment_starts: dict[str, torch.Tensor] = {}
         self.snapshot: dict[str, torch.Tensor] = {}
+        self.scales: dict[str, torch.Tensor] = {}
+        self.scales_changed = False
         self.initialized = False
 
     def _local_shard(self, key: str, value: torch.Tensor, is_master: bool) -> tuple[torch.Tensor, int] | None:
@@ -77,22 +116,73 @@ def _local_shard(self, key: str, value: torch.Tensor, is_master: bool) -> tuple[
         local = value.to_local().detach()[: local_shape[0]]
         return local.reshape(-1), global_offset[0] * prod(value.shape[1:])
 
+    def _find_gathered(self, model: nn.Module, is_master: bool) -> set[str]:
+        """fp8 trainer tensors whose 128-row blocks straddle ranks on some rank."""
+        keys, straddles = [], []
+        for key, value in model.state_dict().items():
+            fqn = resolve_fqn(model, key)
+            if fqn not in self.fp8:
+                continue
+            if self.wire_dtypes[key] != torch.bfloat16:
+                raise NotImplementedError(f"Sparse weight broadcast quantizes bf16 weights to fp8, not {key}")
+            shard = self._local_shard(key, value, is_master)
+            num_cols = self.fp8[fqn][1]
+            num_rows = value.numel() // num_cols
+            first = 0 if shard is None else shard[1] // num_cols
+            end = first if shard is None else first + shard[0].numel() // num_cols
+            keys.append(key)
+            straddles.append(first % FP8_BLOCK[0] != 0 or (end % FP8_BLOCK[0] != 0 and end != num_rows))
+        flags = torch.tensor(straddles, dtype=torch.int32, device=next(model.parameters()).device)
+        if dist.is_initialized():
+            dist.all_reduce(flags, op=dist.ReduceOp.MAX)
+        return {key for key, flag in zip(keys, flags.tolist()) if flag}
+
+    def _fp8_bits(
+        self, key: str, fqn: str, value: torch.Tensor, shard: tuple[torch.Tensor, int] | None
+    ) -> tuple[torch.Tensor, int] | None:
+        """fp8 bits of this rank's shard, quantized as the engines do."""
+        ue8m0, num_cols = self.fp8[fqn]
+        if key in self.gather:
+            quantized, scales = quantize_fp8_blocks(value.full_tensor().to(torch.bfloat16).view(-1, num_cols), ue8m0)
+            if shard is None:
+                return None
+            bits = quantized.view(-1)[shard[1] : shard[1] + shard[0].numel()]
+        elif shard is None or shard[0].numel() == 0:
+            return None
+        else:
+            quantized, scales = quantize_fp8_blocks(shard[0].to(torch.bfloat16).view(-1, num_cols), ue8m0)
+            bits = quantized.view(-1)
+        previous = self.scales.get(key)
+        if previous is not None and not torch.equal(previous, scales):
+            self.scales_changed = True
+        self.scales[key] = scales
+        return bits, shard[1]
+
     @torch.no_grad()
     def encode(
         self, model: nn.Module, is_master: bool
     ) -> tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor] | None:
         """Diff the current shards against the snapshot and advance it. Returns
-        ``(entries, gaps, xor value bytes)`` on the host, or None on the first call."""
+        ``(entries, gaps, xor value bytes)`` on the host, or None on the first
+        call and when an fp8 block scale changed."""
         first = not self.initialized
         self.initialized = True
+        if self.gather is None:
+            self.gather = self._find_gathered(model, is_master)
+        self.scales_changed = False
         names, counts, gaps, values = [], [], [], []
         for key, value in model.state_dict().items():
             fqn = resolve_fqn(model, key)
-            shard = self._local_shard(key, value, is_master) if fqn in self.segments else None
+            if fqn not in self.segments:
+                continue
+            shard = self._local_shard(key, value, is_master)
+            if fqn in self.fp8:
+                shard = self._fp8_bits(key, fqn, value, shard)
+            elif shard is not None:
+                shard = (as_bits(shard[0].to(self.wire_dtypes[key])), shard[1])
             if shard is None or shard[0].numel() == 0:
                 continue
-            local, offset = shard
-            bits = as_bits(local.to(self.wire_dtypes[key]))
+            bits, offset = shard
             if first:
                 self.snapshot[key] = bits.clone()
                 continue
@@ -111,7 +201,11 @@ def encode(
             values.append(xor[changed].view(torch.uint8))
             counts.append(torch.bincount(segment, minlength=len(starts)))
             names.extend(name for _, name in self.segments[fqn])
-        if first:
+        if self.fp8 and dist.is_initialized():
+            changed_anywhere = torch.tensor(int(self.scales_changed), device=next(model.parameters()).device)
+            dist.all_reduce(changed_anywhere, op=dist.ReduceOp.MAX)
+            self.scales_changed = bool(changed_anywhere.item())
+        if first or self.scales_changed:
             return None
         counts = torch.cat(counts).tolist() if counts else []
         entries = [(name, count) for name, count in zip(names, counts) if count]
@@ -141,7 +235,7 @@ def __init__(
         self.encoder: SparseDeltaEncoder | None = None
         self.last_step: int | None = None
         self.anchor_step: int | None = None
-        self.num_broadcasts = 0
+        self.num_deltas = 0
         self.logger.debug("Initialized filesystem weight broadcast")
 
     def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
@@ -165,34 +259,42 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
                 )
             return
 
-        # The second broadcast is an anchor too, so engines check the first delta right away.
-        self.num_broadcasts += 1
-        anchor = not self.config.sparse or self.num_broadcasts <= 2 or step % self.config.anchor_interval == 0
+        wrote_delta = False
         if self.config.sparse:
+            if self.world.is_master:
+                # Engines that load a step with a delta dir report their weight layout next to it.
+                (step_dir / DELTA_DIR).mkdir(exist_ok=True)
             if self.encoder is None:
-                self.encoder = SparseDeltaEncoder(model)
-            # Anchors also carry the delta, so engines can check the chain against them.
-            self._save_delta(model, step, step_dir)
+                # Deltas start once every engine worker has reported which tensors it holds in fp8.
+                layout = [read_engine_layout(step_dir.parent) if self.world.is_master else None]
+                dist.broadcast_object_list(layout, src=0)
+                if layout[0] is not None:
+                    self.encoder = SparseDeltaEncoder(model, layout[0])
+            if self.encoder is not None:
+                wrote_delta = self._save_delta(model, step, step_dir)
+                self.num_deltas += wrote_delta
             self.last_step = step
-        if anchor:
+        # The first delta's step is an anchor too, so engines check the delta path right away.
+        if not wrote_delta or self.num_deltas == 1 or step % self.config.anchor_interval == 0:
             state_dict = gather_weights_parallel(model)
             state_dict = convert_state_dict_to_hf(model, state_dict)
             self.logger.debug(f"Saving weights to {step_dir}")
             save_state_dict_parallel(state_dict, step_dir)
             self.anchor_step = step
 
-    def _save_delta(self, model: nn.Module, step: int, step_dir: Path) -> None:
+    def _save_delta(self, model: nn.Module, step: int, step_dir: Path) -> bool:
+        """Write this step's delta (all ranks). False when the step must be a full checkpoint instead."""
         assert self.encoder is not None
         start = time.perf_counter()
         delta = self.encoder.encode(model, self.world.is_master)
+        if self.encoder.scales_changed:
+            self.logger.info(f"fp8 block scales changed at v{step}: publishing a full checkpoint")
         if delta is None:
-            return
+            return False
         entries, gaps, values = delta
         encode_time = time.perf_counter() - start
         file_name = f"rank_{self.world.rank}.safetensors"
         delta_dir = step_dir / DELTA_DIR
-        if self.world.is_master:
-            delta_dir.mkdir()
         dist.barrier()
         num_bytes = write_rank_file(
             delta_dir / file_name, entries, gaps.numpy(), values.numpy(), self.encoder.hf_tensors
@@ -201,7 +303,7 @@ def _save_delta(self, model: nn.Module, step: int, step_dir: Path) -> None:
         all_stats: list = [None] * self.world.world_size
         dist.all_gather_object(all_stats, stats)
         if not self.world.is_master:
-            return
+            return True
         manifest = {
             "step": step,
             "base_step": self.last_step,
@@ -217,6 +319,7 @@ def _save_delta(self, model: nn.Module, step: int, step_dir: Path) -> None:
             f"{sum(stat[1] for stat in all_stats) / 2**20:.1f} MiB, "
             f"encode {max(stat[3] for stat in all_stats):.2f}s, total {max(stat[4] for stat in all_stats):.2f}s"
         )
+        return True
 
     def _clean(self, step: int) -> None:
         # Sparse chains replay from their anchor, so keep everything since it.
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index e5a3579c6a3..a60f1ac8bce 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -21,6 +21,9 @@
 
 DELTA_DIR = "delta"
 MANIFEST = "manifest.json"
+# Dir next to the step dirs where every engine worker writes which checkpoint tensors it holds in fp8.
+ENGINE_LAYOUT = "engine_layout"
+FP8_BLOCK = [128, 128]
 
 _BITS_DTYPE = {1: torch.uint8, 2: torch.int16, 4: torch.int32, 8: torch.int64}
 
@@ -34,6 +37,19 @@ def bits_dtype(dtype: str) -> torch.dtype:
     return _BITS_DTYPE[getattr(torch, dtype).itemsize]
 
 
+def quantize_fp8_blocks(weight: torch.Tensor, ue8m0: bool) -> tuple[torch.Tensor, torch.Tensor]:
+    """vLLM's online ``fp8_per_block`` quantization of a 2-D weight: fp8 bits (uint8) and fp32 block
+    scales. With ``ue8m0``, also DeepGEMM's requantization to power-of-two scales. Uses vLLM's own
+    compiled kernels: a reimplementation rounds differently on a few elements."""
+    from vllm.model_executor.layers.quantization.utils.fp8_utils import requant_weight_ue8m0_inplace
+    from vllm.utils.deep_gemm import per_block_cast_to_fp8
+
+    quantized, scales = per_block_cast_to_fp8(weight, FP8_BLOCK, use_ue8m0=False)
+    if ue8m0:
+        requant_weight_ue8m0_inplace(quantized, scales, FP8_BLOCK)
+    return quantized.view(torch.uint8), scales
+
+
 def encode_gaps(indices: torch.Tensor, entry_starts: torch.Tensor) -> torch.Tensor:
     """Sorted flat indices -> ``index - previous - 1``, restarting at every
     ``entry_starts`` position (where the gap is the index itself)."""
@@ -107,6 +123,20 @@ def _read_rank_file(
     return entries, gaps, values
 
 
+def read_engine_layout(broadcast_dir: Path) -> dict[str, str] | None:
+    """Checkpoint tensor -> fp8 scale format (``"fp32"``/``"ue8m0"``) the engines hold it in,
+    or None until every engine worker has reported."""
+    layouts = [json.loads(path.read_text()) for path in (broadcast_dir / ENGINE_LAYOUT).glob("*.json")]
+    if not layouts or len(layouts) < layouts[0]["world_size"]:
+        return None
+    fp8 = {}
+    for layout in layouts:
+        for name, mode in layout["fp8"].items():
+            if fp8.setdefault(name, mode) != mode:
+                raise NotImplementedError(f"Engine workers hold {name} in different fp8 formats")
+    return fp8
+
+
 def read_manifest(step_dir: Path) -> dict:
     return json.loads((step_dir / DELTA_DIR / MANIFEST).read_text())
 

From ab7f5721b4a0b91097c736a7ad797d0330389a12 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 14 Oct 2026 14:43:44 +0000
Subject: [PATCH 05/88] feat(weights): write periodic sparse-broadcast anchors
 off the trainer's critical path

---
 .../inference/vllm/worker/filesystem.py       |  5 +-
 src/prime_rl/transports/weights/filesystem.py | 74 +++++++++++++------
 src/prime_rl/utils/sparse_weights.py          | 25 +++++++
 3 files changed, 81 insertions(+), 23 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 68a11ce3830..218db8551b9 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -23,6 +23,7 @@
     DELTA_DIR,
     ENGINE_LAYOUT,
     MANIFEST,
+    anchor_complete,
     bits_dtype,
     decode_gaps,
     read_delta,
@@ -88,7 +89,7 @@ def update_weights_from_path(self, weight_path: str) -> None:
             self._load_checkpoint(step_dir)
         elif read_manifest(step_dir)["base_step"] == self.weight_step:
             self._apply_delta(step_dir)
-            if any(step_dir.glob("*.safetensors")):
+            if anchor_complete(step_dir):
                 # Anchor step: check that the delta chain reproduced the full checkpoint bit for bit.
                 chained = weight_checksums(self._model())
                 self._load_checkpoint(step_dir)
@@ -99,7 +100,7 @@ def update_weights_from_path(self, weight_path: str) -> None:
         else:
             self._replay_from_anchor(step_dir)
         self.weight_step = step
-        if (step_dir / DELTA_DIR).is_dir() and not self.wrote_layout and any(step_dir.glob("*.safetensors")):
+        if (step_dir / DELTA_DIR).is_dir() and not self.wrote_layout and anchor_complete(step_dir):
             self._write_engine_layout(step_dir)
             self.wrote_layout = True
 
diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index 77c55566abf..3d79f9e34e3 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -1,4 +1,5 @@
 import json
+import threading
 import time
 from collections import defaultdict
 from math import prod
@@ -20,10 +21,12 @@
     DELTA_DIR,
     FP8_BLOCK,
     MANIFEST,
+    anchor_complete,
     as_bits,
     encode_gaps,
     quantize_fp8_blocks,
     read_engine_layout,
+    write_anchor_shards,
     write_rank_file,
 )
 from prime_rl.utils.weights import (
@@ -236,6 +239,7 @@ def __init__(
         self.last_step: int | None = None
         self.anchor_step: int | None = None
         self.num_deltas = 0
+        self.pending_anchor: int | None = None
         self.logger.debug("Initialized filesystem weight broadcast")
 
     def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
@@ -259,28 +263,56 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
                 )
             return
 
-        wrote_delta = False
-        if self.config.sparse:
-            if self.world.is_master:
-                # Engines that load a step with a delta dir report their weight layout next to it.
-                (step_dir / DELTA_DIR).mkdir(exist_ok=True)
-            if self.encoder is None:
-                # Deltas start once every engine worker has reported which tensors it holds in fp8.
-                layout = [read_engine_layout(step_dir.parent) if self.world.is_master else None]
-                dist.broadcast_object_list(layout, src=0)
-                if layout[0] is not None:
-                    self.encoder = SparseDeltaEncoder(model, layout[0])
-            if self.encoder is not None:
-                wrote_delta = self._save_delta(model, step, step_dir)
-                self.num_deltas += wrote_delta
-            self.last_step = step
-        # The first delta's step is an anchor too, so engines check the delta path right away.
-        if not wrote_delta or self.num_deltas == 1 or step % self.config.anchor_interval == 0:
-            state_dict = gather_weights_parallel(model)
-            state_dict = convert_state_dict_to_hf(model, state_dict)
-            self.logger.debug(f"Saving weights to {step_dir}")
-            save_state_dict_parallel(state_dict, step_dir)
+        if not self.config.sparse:
+            self._save_full(model, step_dir)
+            return
+        if self.world.is_master:
+            # Engines that load a step with a delta dir report their weight layout next to it.
+            (step_dir / DELTA_DIR).mkdir(exist_ok=True)
+        if self.encoder is None:
+            # Deltas start once every engine worker has reported which tensors it holds in fp8.
+            layout = [read_engine_layout(step_dir.parent) if self.world.is_master else None]
+            dist.broadcast_object_list(layout, src=0)
+            if layout[0] is not None:
+                self.encoder = SparseDeltaEncoder(model, layout[0])
+        wrote_delta = self.encoder is not None and self._save_delta(model, step, step_dir)
+        self.num_deltas += wrote_delta
+        self.last_step = step
+        if not wrote_delta or self.num_deltas == 1:
+            # Engines need this full checkpoint now; the first delta's step is one too, so engines
+            # check the delta path against it right away.
+            self._save_anchor(model, step_dir).join()
+            dist.barrier()
             self.anchor_step = step
+            return
+        # Periodic anchors only serve engines that fall off the chain: write them off the critical path.
+        if (
+            self.world.is_master
+            and self.pending_anchor is not None
+            and anchor_complete(self.step_dir(self.pending_anchor))
+        ):
+            self.anchor_step, self.pending_anchor = max(self.anchor_step, self.pending_anchor), None
+        start = [step % self.config.anchor_interval == 0 and self.pending_anchor is None]
+        dist.broadcast_object_list(start, src=0)
+        if start[0]:
+            self._save_anchor(model, step_dir)
+            self.pending_anchor = step
+
+    def _save_full(self, model: nn.Module, step_dir: Path) -> None:
+        state_dict = gather_weights_parallel(model)
+        state_dict = convert_state_dict_to_hf(model, state_dict)
+        self.logger.debug(f"Saving weights to {step_dir}")
+        save_state_dict_parallel(state_dict, step_dir)
+
+    def _save_anchor(self, model: nn.Module, step_dir: Path) -> threading.Thread:
+        """Gather this rank's share of the full checkpoint to CPU, then write it on a thread
+        (no collectives); engines treat the step as an anchor once every rank's marker exists."""
+        state_dict = convert_state_dict_to_hf(model, gather_weights_parallel(model))
+        thread = threading.Thread(
+            target=write_anchor_shards, args=(state_dict, step_dir, self.world.rank, self.world.world_size)
+        )
+        thread.start()
+        return thread
 
     def _save_delta(self, model: nn.Module, step: int, step_dir: Path) -> bool:
         """Write this step's delta (all ranks). False when the step must be a full checkpoint instead."""
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index a60f1ac8bce..8effb618eb8 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -18,6 +18,7 @@
 import torch
 from safetensors import safe_open
 from safetensors.numpy import save_file
+from safetensors.torch import save_file as save_torch_file
 
 DELTA_DIR = "delta"
 MANIFEST = "manifest.json"
@@ -123,6 +124,30 @@ def _read_rank_file(
     return entries, gaps, values
 
 
+def write_anchor_shards(state_dict: dict[str, torch.Tensor], step_dir: Path, rank: int, world_size: int) -> None:
+    """Write one trainer rank's share of a full checkpoint, without collectives; the
+    ``anchor_rank{rank}.done`` marker (holding the world size) is written last."""
+    shards, size = [[]], 0
+    for name, tensor in state_dict.items():
+        if size > 4 * 2**30:
+            shards.append([])
+            size = 0
+        shards[-1].append(name)
+        size += tensor.nbytes
+    for index, names in enumerate(shards):
+        path = step_dir / f"model-rank{rank:05d}-{index:05d}.safetensors"
+        tmp = step_dir / f".{path.name}.tmp"
+        save_torch_file({name: state_dict[name].contiguous() for name in names}, tmp, metadata={"format": "pt"})
+        tmp.rename(path)
+    (step_dir / f"anchor_rank{rank}.done").write_text(str(world_size))
+
+
+def anchor_complete(step_dir: Path) -> bool:
+    """Whether ``step_dir`` holds a complete full checkpoint written by ``write_anchor_shards``."""
+    markers = list(step_dir.glob("anchor_rank*.done"))
+    return bool(markers) and len(markers) == int(markers[0].read_text())
+
+
 def read_engine_layout(broadcast_dir: Path) -> dict[str, str] | None:
     """Checkpoint tensor -> fp8 scale format (``"fp32"``/``"ue8m0"``) the engines hold it in,
     or None until every engine worker has reported."""

From 23123e9af80b1255ae51c316baa1e3bf872aead0 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 14 Oct 2026 16:11:23 +0000
Subject: [PATCH 06/88] perf(weights): keep the sparse snapshot in pinned host
 memory and reuse pinned buffers for anchors and prefetch

---
 .../inference/vllm/worker/filesystem.py       | 35 ++++++++++--
 src/prime_rl/transports/weights/filesystem.py | 57 +++++++++++++------
 src/prime_rl/utils/sparse_weights.py          | 11 ++--
 src/prime_rl/utils/weights.py                 | 34 +++++++++--
 4 files changed, 107 insertions(+), 30 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 218db8551b9..4cb3d22ee64 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -57,6 +57,7 @@ class FileSystemWeightUpdateWorker(Worker):
     prefetched: tuple[str, Future] | None = None
     index_map: CheckpointIndexMap | None = None
     wrote_layout = False
+    host_buffer: torch.Tensor | None = None
 
     def liveness_probe(self) -> None:
         """No-op RPC used by the API server liveness endpoint."""
@@ -72,7 +73,31 @@ def _model(self) -> Module:
 
     def prefetch_weights_from_path(self, weight_path: str) -> None:
         """Start decoding the sparse delta in ``weight_path`` into pinned host memory."""
-        self.prefetched = (weight_path, _prefetch_pool.submit(read_delta, Path(weight_path), self._spans()))
+        self.prefetched = (weight_path, _prefetch_pool.submit(self._read_delta, Path(weight_path)))
+
+    def _read_delta(self, step_dir: Path) -> list[tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]]:
+        """Read this worker's part of a delta into a pinned host buffer that is reused across steps
+        (allocating pinned memory every step is slow on a loaded host)."""
+        start = time.perf_counter()
+        delta = read_delta(step_dir, self._spans())
+        read_time = time.perf_counter() - start
+        num_bytes = sum(-(-gaps.nbytes // 4) * 4 + -(-values.nbytes // 4) * 4 for _, gaps, values in delta)
+        if self.host_buffer is None or self.host_buffer.numel() < num_bytes:
+            self.host_buffer = torch.empty(num_bytes * 5 // 4, dtype=torch.uint8, pin_memory=True)
+        pinned, offset = [], 0
+        for entries, gaps, values in delta:
+            parts = []
+            for array in (gaps, values):
+                part = self.host_buffer[offset : offset + array.nbytes].view(array.dtype)
+                part.copy_(array)
+                parts.append(part)
+                offset += -(-array.nbytes // 4) * 4
+            pinned.append((entries, *parts))
+        logger.info(
+            f"Read sparse weight delta {step_dir.name}: {num_bytes / 2**30:.2f} GiB, "
+            f"read+decompress {read_time:.2f}s, pin {time.perf_counter() - start - read_time:.2f}s"
+        )
+        return pinned
 
     def _spans(self) -> dict[str, tuple[int, int]] | None:
         """The checkpoint index ranges this worker's params hold, once the index map exists."""
@@ -149,10 +174,10 @@ def _replay_from_anchor(self, step_dir: Path) -> None:
     def _apply_delta(self, step_dir: Path) -> None:
         start = time.perf_counter()
         prefetched, self.prefetched = self.prefetched, None
-        if prefetched is not None and prefetched[0] == step_dir.as_posix():
-            delta = prefetched[1].result()
-        else:
-            delta = read_delta(step_dir, self._spans())
+        # Always wait: a running prefetch writes into the same host buffer.
+        delta = prefetched[1].result() if prefetched is not None else None
+        if prefetched is None or prefetched[0] != step_dir.as_posix():
+            delta = self._read_delta(step_dir)
         read_time = time.perf_counter() - start
         tensors = read_manifest(step_dir)["tensors"]
         model = self._model()
diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index 3d79f9e34e3..b12b339a793 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -32,6 +32,7 @@
 from prime_rl.utils.weights import (
     convert_state_dict_to_hf,
     gather_weights_parallel,
+    owned_weight_bytes,
     resolve_fqn,
     resolve_wire_dtype,
     save_state_dict,
@@ -161,6 +162,20 @@ def _fp8_bits(
         self.scales[key] = scales
         return bits, shard[1]
 
+    @staticmethod
+    def _diff(bits: torch.Tensor, previous: torch.Tensor, chunk: int = 1 << 28) -> tuple[torch.Tensor, torch.Tensor]:
+        """Changed flat indices of ``bits`` vs the pinned host ``previous`` and their XOR values,
+        streaming ``previous`` through the GPU in chunks and updating it in place."""
+        changed, changed_values = [], []
+        for start in range(0, bits.numel(), chunk):
+            current = bits[start : start + chunk]
+            xor = current ^ previous[start : start + chunk].to(bits.device, non_blocking=True)
+            previous[start : start + chunk].copy_(current, non_blocking=True)
+            index = xor.nonzero().squeeze(1)
+            changed.append(index + start)
+            changed_values.append(xor[index])
+        return torch.cat(changed), torch.cat(changed_values)
+
     @torch.no_grad()
     def encode(
         self, model: nn.Module, is_master: bool
@@ -187,21 +202,20 @@ def encode(
                 continue
             bits, offset = shard
             if first:
-                self.snapshot[key] = bits.clone()
+                # The previous broadcast's bits live in pinned host memory, so the trainer's GPU memory is unchanged.
+                self.snapshot[key] = torch.empty(bits.shape, dtype=bits.dtype, pin_memory=True)
+                self.snapshot[key].copy_(bits)
                 continue
-            previous = self.snapshot[key]
-            xor = bits ^ previous
-            previous.copy_(bits)
-            changed = xor.nonzero().squeeze(1)
+            changed, changed_values = self._diff(bits, self.snapshot[key])
             if fqn not in self.segment_starts:
-                self.segment_starts[fqn] = torch.tensor([start for start, _ in self.segments[fqn]], device=xor.device)
+                self.segment_starts[fqn] = torch.tensor([start for start, _ in self.segments[fqn]], device=bits.device)
             starts = self.segment_starts[fqn]
             position = changed + offset
             segment = torch.searchsorted(starts, position, right=True) - 1
             entry_starts = torch.ones_like(segment, dtype=torch.bool)
             entry_starts[1:] = segment[1:] != segment[:-1]
-            gaps.append(encode_gaps(position - starts[segment], entry_starts))
-            values.append(xor[changed].view(torch.uint8))
+            gaps.append(encode_gaps(position - starts[segment], entry_starts).to("cpu", non_blocking=True))
+            values.append(changed_values.view(torch.uint8).to("cpu", non_blocking=True))
             counts.append(torch.bincount(segment, minlength=len(starts)))
             names.extend(name for _, name in self.segments[fqn])
         if self.fp8 and dist.is_initialized():
@@ -210,12 +224,11 @@ def encode(
             self.scales_changed = bool(changed_anywhere.item())
         if first or self.scales_changed:
             return None
-        counts = torch.cat(counts).tolist() if counts else []
+        counts = torch.cat(counts).tolist() if counts else []  # also waits for the copies to the host
+        torch.cuda.empty_cache()
         entries = [(name, count) for name, count in zip(names, counts) if count]
         empty = torch.empty(0, dtype=torch.uint8)
-        gaps_host = torch.cat(gaps).cpu() if gaps else empty
-        values_host = torch.cat(values).cpu() if values else empty
-        return entries, gaps_host, values_host
+        return entries, torch.cat(gaps) if gaps else empty, torch.cat(values) if values else empty
 
 
 class FileSystemWeightSender(WeightSender):
@@ -240,6 +253,8 @@ def __init__(
         self.anchor_step: int | None = None
         self.num_deltas = 0
         self.pending_anchor: int | None = None
+        # Reused across anchors: a new one starts only after the previous one was written
+        self.anchor_buffer: torch.Tensor | None = None
         self.logger.debug("Initialized filesystem weight broadcast")
 
     def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
@@ -307,10 +322,20 @@ def _save_full(self, model: nn.Module, step_dir: Path) -> None:
     def _save_anchor(self, model: nn.Module, step_dir: Path) -> threading.Thread:
         """Gather this rank's share of the full checkpoint to CPU, then write it on a thread
         (no collectives); engines treat the step as an anchor once every rank's marker exists."""
-        state_dict = convert_state_dict_to_hf(model, gather_weights_parallel(model))
-        thread = threading.Thread(
-            target=write_anchor_shards, args=(state_dict, step_dir, self.world.rank, self.world.world_size)
-        )
+        start = time.perf_counter()
+        if self.anchor_buffer is None:
+            self.anchor_buffer = torch.empty(owned_weight_bytes(model), dtype=torch.uint8, pin_memory=True)
+        state_dict = convert_state_dict_to_hf(model, gather_weights_parallel(model, host_buffer=self.anchor_buffer))
+        torch.cuda.empty_cache()
+        gather_time = time.perf_counter() - start
+
+        def write() -> None:
+            write_anchor_shards(state_dict, step_dir, self.world.rank, self.world.world_size)
+            self.logger.debug(
+                f"Wrote the {step_dir.name} anchor: gather {gather_time:.1f}s, write {time.perf_counter() - start - gather_time:.1f}s"
+            )
+
+        thread = threading.Thread(target=write)
         thread.start()
         return thread
 
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index 8effb618eb8..ef5aa34a524 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -11,6 +11,7 @@
 
 import json
 from concurrent.futures import ThreadPoolExecutor
+from math import prod
 from pathlib import Path
 
 import numpy as np
@@ -104,9 +105,11 @@ def _read_rank_file(
                 entry_values = np.frombuffer(
                     pyzstd.decompress(values_bytes[values_offset:values_end].tobytes()), np.uint8
                 )
-                if spans is not None:
+                # Trim to this reader's index range unless it holds the whole tensor
+                span_start, span_end = spans[name] if spans is not None else (0, 2**62)
+                if span_start > 0 or span_end < prod(tensors[name][0]):
                     indices = np.cumsum(entry_gaps, dtype=np.int64) + np.arange(1, count + 1) - 1
-                    begin, end = np.searchsorted(indices, spans[name])
+                    begin, end = np.searchsorted(indices, (span_start, span_end))
                     itemsize = bits_dtype(tensors[name][1]).itemsize
                     entry_values = entry_values[begin * itemsize : end * itemsize]
                     entry_gaps = entry_gaps[begin:end].copy()
@@ -119,8 +122,6 @@ def _read_rank_file(
             gaps_offset, values_offset = gaps_end, values_end
     gaps = torch.from_numpy(np.concatenate(gaps) if gaps else np.empty(0, np.int32))
     values = torch.from_numpy(np.concatenate(values) if values else np.empty(0, np.uint8))
-    if torch.cuda.is_available():
-        gaps, values = gaps.pin_memory(), values.pin_memory()
     return entries, gaps, values
 
 
@@ -170,7 +171,7 @@ def read_delta(
     step_dir: Path, spans: dict[str, tuple[int, int]] | None = None
 ) -> list[tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]]:
     """Read and decompress every rank file of a delta step dir into ``(entries,
-    gaps, xor value bytes)`` in (pinned) host memory; ``decode_gaps`` turns the
+    gaps, xor value bytes)`` in host memory; ``decode_gaps`` turns the
     gaps into flat indices on the device that applies them. With ``spans``, keep
     only the listed tensors and, within each, the flat indices in ``[start, end)``."""
     manifest = read_manifest(step_dir)
diff --git a/src/prime_rl/utils/weights.py b/src/prime_rl/utils/weights.py
index ac43509c4b1..e50e8ee637e 100644
--- a/src/prime_rl/utils/weights.py
+++ b/src/prime_rl/utils/weights.py
@@ -159,14 +159,17 @@ def resolve_wire_dtype(keep_in_fp32: Callable[[str], bool] | None, key: str, def
     return torch.float32 if keep_in_fp32 is not None and keep_in_fp32(key) else default
 
 
-def gather_weights_parallel(model: nn.Module, dtype: torch.dtype = torch.bfloat16) -> dict[str, Tensor]:
+def gather_weights_parallel(
+    model: nn.Module, dtype: torch.dtype = torch.bfloat16, host_buffer: Tensor | None = None
+) -> dict[str, Tensor]:
     """Gather distributed weights cooperatively, each rank keeping a slice on CPU.
 
     Every rank participates in the per-tensor all-gathers (a ``full_tensor`` call is
     collective and materializes the full tensor on all ranks anyway), but instead of
     only the master keeping everything, each rank copies its owned slice to CPU —
-    so the D2H traffic and the shard writes are split across ranks. The copies are
-    plain blocking transfers; no stream-ordering assumptions.
+    so the D2H traffic and the shard writes are split across ranks. With
+    ``host_buffer`` (pinned uint8 of at least ``owned_weight_bytes``), the slice is
+    copied into it instead of fresh pageable memory, which is several times faster.
 
     Parameters the model reports via ``keep_in_fp32_for_weight_transfer`` skip the
     downcast. Those lists exist because the inference engine's kernels require fp32
@@ -179,6 +182,7 @@ def gather_weights_parallel(model: nn.Module, dtype: torch.dtype = torch.bfloat1
     world = get_world()
     owners = partition_weights(model.state_dict(), world.world_size, dtype)
     partial: dict[str, Tensor] = {}
+    offset = 0
     with warnings.catch_warnings():
         warnings.filterwarnings("ignore", category=FutureWarning, module="torch.distributed")
         warnings.filterwarnings("ignore", category=UserWarning, module="torch.distributed.*")
@@ -190,7 +194,15 @@ def gather_weights_parallel(model: nn.Module, dtype: torch.dtype = torch.bfloat1
                 value = cast(DTensor, value.to(target_dtype)).full_tensor()
             if owners[key] != world.rank:
                 continue
-            partial[resolve_fqn(model, key)] = value.to("cpu")
+            if host_buffer is None:
+                partial[resolve_fqn(model, key)] = value.to("cpu")
+                continue
+            out = host_buffer[offset : offset + value.nbytes].view(value.dtype).view(value.shape)
+            out.copy_(value, non_blocking=True)
+            partial[resolve_fqn(model, key)] = out
+            offset += -(-value.nbytes // 8) * 8
+        if host_buffer is not None:
+            torch.cuda.synchronize()
         dist.barrier()
 
     if any(".base_layer." in key or "lora_A" in key or "lora_B" in key for key in partial.keys()):
@@ -199,6 +211,20 @@ def gather_weights_parallel(model: nn.Module, dtype: torch.dtype = torch.bfloat1
     return partial
 
 
+def owned_weight_bytes(model: nn.Module, dtype: torch.dtype = torch.bfloat16) -> int:
+    """Host bytes of this rank's share of ``gather_weights_parallel`` (8-byte aligned per tensor)."""
+    keep_in_fp32 = getattr(model, "keep_in_fp32_for_weight_transfer", None)
+    world = get_world()
+    owners = partition_weights(model.state_dict(), world.world_size, dtype)
+    total = 0
+    for key, value in model.state_dict().items():
+        if owners[key] == world.rank:
+            is_dtensor = isinstance(value, DTensor)
+            itemsize = resolve_wire_dtype(keep_in_fp32, key, dtype).itemsize if is_dtensor else value.element_size()
+            total += -(-value.numel() * itemsize // 8) * 8
+    return total
+
+
 def save_state_dict_parallel(state_dict: dict[str, Tensor], save_dir: Path) -> None:
     """Cooperatively save a rank-partitioned state dict as sharded safetensors.
 

From 0f7fad666267c294db62545b95dc206ebd6b15ae Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 14 Oct 2026 18:04:23 +0000
Subject: [PATCH 07/88] feat(weights): verify every sparse delta against
 trainer checksums; larger zstd frames read in parallel

---
 .../inference/vllm/worker/filesystem.py       |  33 ++++-
 .../inference/vllm/worker/weight_transfer.py  |  25 ++++
 src/prime_rl/transports/weights/filesystem.py |  22 ++-
 src/prime_rl/utils/sparse_weights.py          | 135 +++++++++++++-----
 tests/unit/utils/test_sparse_weights.py       |  14 +-
 5 files changed, 175 insertions(+), 54 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 4cb3d22ee64..7a98ce63b01 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -9,7 +9,7 @@
 import torch
 from safetensors import safe_open
 from torch.nn import Module
-from vllm.distributed.parallel_state import get_ep_group, get_world_group
+from vllm.distributed.parallel_state import get_ep_group, get_tp_group, get_world_group
 from vllm.logger import init_logger
 from vllm.model_executor.model_loader import DefaultModelLoader, get_model_loader
 
@@ -75,11 +75,13 @@ def prefetch_weights_from_path(self, weight_path: str) -> None:
         """Start decoding the sparse delta in ``weight_path`` into pinned host memory."""
         self.prefetched = (weight_path, _prefetch_pool.submit(self._read_delta, Path(weight_path)))
 
-    def _read_delta(self, step_dir: Path) -> list[tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]]:
+    def _read_delta(
+        self, step_dir: Path
+    ) -> tuple[list[tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]], dict[str, int]]:
         """Read this worker's part of a delta into a pinned host buffer that is reused across steps
         (allocating pinned memory every step is slow on a loaded host)."""
         start = time.perf_counter()
-        delta = read_delta(step_dir, self._spans())
+        delta, checksums = read_delta(step_dir, self._spans())
         read_time = time.perf_counter() - start
         num_bytes = sum(-(-gaps.nbytes // 4) * 4 + -(-values.nbytes // 4) * 4 for _, gaps, values in delta)
         if self.host_buffer is None or self.host_buffer.numel() < num_bytes:
@@ -97,7 +99,7 @@ def _read_delta(self, step_dir: Path) -> list[tuple[list[tuple[str, int]], torch
             f"Read sparse weight delta {step_dir.name}: {num_bytes / 2**30:.2f} GiB, "
             f"read+decompress {read_time:.2f}s, pin {time.perf_counter() - start - read_time:.2f}s"
         )
-        return pinned
+        return pinned, checksums
 
     def _spans(self) -> dict[str, tuple[int, int]] | None:
         """The checkpoint index ranges this worker's params hold, once the index map exists."""
@@ -178,6 +180,7 @@ def _apply_delta(self, step_dir: Path) -> None:
         delta = prefetched[1].result() if prefetched is not None else None
         if prefetched is None or prefetched[0] != step_dir.as_posix():
             delta = self._read_delta(step_dir)
+        delta, checksums = delta
         read_time = time.perf_counter() - start
         tensors = read_manifest(step_dir)["tensors"]
         model = self._model()
@@ -202,11 +205,31 @@ def _apply_delta(self, step_dir: Path) -> None:
                 num_bytes += num_changes * itemsize
         update_mla_absorbed_weights(model)
         torch.cuda.synchronize()
+        apply_time = time.perf_counter() - start - read_time
+        self._verify(step_dir, checksums)
         logger.info(
             f"Applied sparse weight delta {step_dir.name}: {sum(len(gaps) for _, gaps, _ in delta)} changes, "
-            f"read {read_time:.2f}s (0 if prefetched), apply {time.perf_counter() - start - read_time:.2f}s"
+            f"read {read_time:.2f}s (0 if prefetched), apply {apply_time:.2f}s, "
+            f"verify {time.perf_counter() - start - read_time - apply_time:.2f}s"
         )
 
+    def _verify(self, step_dir: Path, checksums: dict[str, int]) -> None:
+        """Check every tensor this worker holds against the trainer's checksum of the new weights.
+        Tensors split across TP workers are summed over the TP group first."""
+        assert self.index_map is not None
+        sums, partial = self.index_map.held_checksums(checksums)
+        tp_group = get_tp_group()
+        partial_names = sorted(partial) if tp_group.world_size > 1 else []
+        if partial_names:
+            reduced = torch.stack([sums[name] for name in partial_names])
+            torch.distributed.all_reduce(reduced, group=tp_group.device_group)
+            sums.update(zip(partial_names, reduced))
+        names = [name for name in sums if name not in partial or partial_names]
+        values = torch.stack([sums[name] for name in names]).tolist() if names else []
+        mismatched = [name for name, value in zip(names, values) if value != checksums[name]]
+        if mismatched:
+            raise RuntimeError(f"Sparse weight delta {step_dir.name} diverged from the trainer: {mismatched[:8]}")
+
     def _load_checkpoint(self, step_dir: Path) -> None:
         """Reload every weight from the HF-compatible checkpoint in ``step_dir``."""
         model = self._model()
diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py
index 566da796473..1d8e503dccd 100644
--- a/src/prime_rl/inference/vllm/worker/weight_transfer.py
+++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py
@@ -1,4 +1,5 @@
 from collections import defaultdict
+from math import prod
 from typing import Iterable
 
 import torch
@@ -114,6 +115,7 @@ def placeholders():
             for name, writes in records.items()
         }
         self.num_cols = {name: (shape[-1] if shape else 1) for name, (shape, _) in tensors.items()}
+        self.numel = {name: prod(shape) for name, (shape, _) in tensors.items()}
         self.unpatchable = unpatchable
         # [first, last + 1) flat checkpoint index each tensor's writes cover; unpatchable tensors are
         # read whole so that a change to them is caught in apply_xor
@@ -125,6 +127,29 @@ def placeholders():
             for name, writes in self.records.items()
         } | {name: (0, 2**62) for name in unpatchable}
 
+    def held_checksums(self, names) -> tuple[dict[str, torch.Tensor], set[str]]:
+        """``bits_checksum`` of the live param elements each checkpoint tensor was written to, for the
+        ``names`` this worker writes, and the subset it holds only part of (TP slices)."""
+        sums, partial = {}, set()
+        for name in names:
+            writes = self.records.get(name)
+            if writes is None or name in self.unpatchable:
+                continue
+            held = sum(rows * cols for _, rows, _, cols, *_ in writes)
+            if held > self.numel[name]:
+                continue
+            parts = [
+                as_bits(self.params[index].data)
+                .view(-1)
+                .as_strided((rows, cols), strides, offset)
+                .sum(dtype=torch.int64)
+                for _, rows, _, cols, index, offset, *strides in writes
+            ]
+            sums[name] = torch.stack(parts).sum()
+            if held < self.numel[name]:
+                partial.add(name)
+        return sums, partial
+
     def fp8_layout(self, read_tensor) -> dict[str, str]:
         """The vLLM per-block quantization (``"fp32"`` or ``"ue8m0"`` scales) that reproduces each
         fp8 param from the checkpoint, per checkpoint tensor stored in one. Checks one checkpoint
diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index b12b339a793..85b8dbfad5b 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -68,6 +68,7 @@ def __init__(self, model: nn.Module, fp8: dict[str, str]):
 
         segments = defaultdict(list)
         self.hf_tensors: dict[str, tuple[list[int], str]] = {}
+        self.hf_numel: dict[str, int] = {}
         for name, tensor in convert_state_dict_to_hf(model, dict(meta)).items():
             base = tensor if tensor._base is None else tensor._base
             if id(base) not in owner or not tensor.is_contiguous():
@@ -77,6 +78,7 @@ def __init__(self, model: nn.Module, fp8: dict[str, str]):
             segments[owner[id(base)]].append((tensor.storage_offset(), name))
             dtype = "float8_e4m3fn" if name in fp8 else str(tensor.dtype).removeprefix("torch.")
             self.hf_tensors[name] = (list(tensor.shape), dtype)
+            self.hf_numel[name] = tensor.numel()
         self.segments = {fqn: sorted(segs) for fqn, segs in segments.items()}
 
         # fqn -> (ue8m0 scales, columns) for trainer tensors held in fp8; quantized as one [rows, columns] matrix
@@ -179,16 +181,17 @@ def _diff(bits: torch.Tensor, previous: torch.Tensor, chunk: int = 1 << 28) -> t
     @torch.no_grad()
     def encode(
         self, model: nn.Module, is_master: bool
-    ) -> tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor] | None:
+    ) -> tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor, dict[str, int]] | None:
         """Diff the current shards against the snapshot and advance it. Returns
-        ``(entries, gaps, xor value bytes)`` on the host, or None on the first
-        call and when an fp8 block scale changed."""
+        ``(entries, gaps, xor value bytes, new-weight checksum per HF tensor part)``
+        on the host, or None on the first call and when an fp8 block scale changed."""
         first = not self.initialized
         self.initialized = True
         if self.gather is None:
             self.gather = self._find_gathered(model, is_master)
         self.scales_changed = False
         names, counts, gaps, values = [], [], [], []
+        checksum_names, checksums = [], []
         for key, value in model.state_dict().items():
             fqn = resolve_fqn(model, key)
             if fqn not in self.segments:
@@ -201,6 +204,12 @@ def encode(
             if shard is None or shard[0].numel() == 0:
                 continue
             bits, offset = shard
+            if not first:
+                for start, name in self.segments[fqn]:
+                    begin, end = max(start, offset), min(start + self.hf_numel[name], offset + bits.numel())
+                    if begin < end:
+                        checksum_names.append(name)
+                        checksums.append(bits[begin - offset : end - offset].sum(dtype=torch.int64))
             if first:
                 # The previous broadcast's bits live in pinned host memory, so the trainer's GPU memory is unchanged.
                 self.snapshot[key] = torch.empty(bits.shape, dtype=bits.dtype, pin_memory=True)
@@ -225,10 +234,11 @@ def encode(
         if first or self.scales_changed:
             return None
         counts = torch.cat(counts).tolist() if counts else []  # also waits for the copies to the host
+        checksums = dict(zip(checksum_names, torch.stack(checksums).tolist())) if checksums else {}
         torch.cuda.empty_cache()
         entries = [(name, count) for name, count in zip(names, counts) if count]
         empty = torch.empty(0, dtype=torch.uint8)
-        return entries, torch.cat(gaps) if gaps else empty, torch.cat(values) if values else empty
+        return entries, torch.cat(gaps) if gaps else empty, torch.cat(values) if values else empty, checksums
 
 
 class FileSystemWeightSender(WeightSender):
@@ -348,13 +358,13 @@ def _save_delta(self, model: nn.Module, step: int, step_dir: Path) -> bool:
             self.logger.info(f"fp8 block scales changed at v{step}: publishing a full checkpoint")
         if delta is None:
             return False
-        entries, gaps, values = delta
+        entries, gaps, values, checksums = delta
         encode_time = time.perf_counter() - start
         file_name = f"rank_{self.world.rank}.safetensors"
         delta_dir = step_dir / DELTA_DIR
         dist.barrier()
         num_bytes = write_rank_file(
-            delta_dir / file_name, entries, gaps.numpy(), values.numpy(), self.encoder.hf_tensors
+            delta_dir / file_name, entries, gaps.numpy(), values.numpy(), self.encoder.hf_tensors, checksums
         )
         stats = (file_name, num_bytes, len(gaps), encode_time, time.perf_counter() - start)
         all_stats: list = [None] * self.world.world_size
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index ef5aa34a524..01720d8c729 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -10,6 +10,7 @@
 """
 
 import json
+import os
 from concurrent.futures import ThreadPoolExecutor
 from math import prod
 from pathlib import Path
@@ -17,7 +18,6 @@
 import numpy as np
 import pyzstd
 import torch
-from safetensors import safe_open
 from safetensors.numpy import save_file
 from safetensors.torch import save_file as save_torch_file
 
@@ -62,23 +62,49 @@ def encode_gaps(indices: torch.Tensor, entry_starts: torch.Tensor) -> torch.Tens
 
 
 def write_rank_file(
-    path: Path, entries: list[tuple[str, int]], gaps: np.ndarray, values: np.ndarray, tensors: dict
+    path: Path,
+    entries: list[tuple[str, int]],
+    gaps: np.ndarray,
+    values: np.ndarray,
+    tensors: dict,
+    checksums: dict[str, int],
+    frame_size: int = 1 << 20,
 ) -> int:
-    """Write every entry's gaps and values as separate zstd frames, so readers can skip entries."""
+    """Write a rank's delta as zstd frames of about ``frame_size`` changes each (whole entries), so
+    readers can skip frames without entries they hold. ``checksums`` (HF name -> ``bits_checksum`` of
+    the new weights this rank holds) let engines verify every step."""
     options = {pyzstd.CParameter.compressionLevel: 1, pyzstd.CParameter.checksumFlag: 1}
-    counts = [count for _, count in entries]
-    value_counts = [count * bits_dtype(tensors[name][1]).itemsize for name, count in entries]
-    gap_parts = np.split(gaps, np.cumsum(counts)[:-1])
-    value_parts = np.split(values, np.cumsum(value_counts)[:-1])
+    frames, size = [[]], 0
+    for index, (_, count) in enumerate(entries):
+        if size >= frame_size:
+            frames.append([])
+            size = 0
+        frames[-1].append(index)
+        size += count
+    gap_ends = np.cumsum([count for _, count in entries], dtype=np.int64)
+    value_ends = np.cumsum([count * bits_dtype(tensors[name][1]).itemsize for name, count in entries], dtype=np.int64)
+
+    def compress(frame: list[int]) -> tuple[bytes, bytes]:
+        first, last = frame[0], frame[-1]
+        gap_start = gap_ends[first - 1] if first else 0
+        value_start = value_ends[first - 1] if first else 0
+        return (
+            pyzstd.compress(gaps[gap_start : gap_ends[last]].tobytes(), options),
+            pyzstd.compress(values[value_start : value_ends[last]].tobytes(), options),
+        )
+
     with ThreadPoolExecutor(8) as pool:
-        gap_frames = list(pool.map(lambda part: pyzstd.compress(part.tobytes(), options), gap_parts))
-        value_frames = list(pool.map(lambda part: pyzstd.compress(part.tobytes(), options), value_parts))
+        compressed = list(pool.map(compress, frames if entries else []))
     arrays = {
-        "gaps": np.frombuffer(b"".join(gap_frames), np.uint8),
-        "values": np.frombuffer(b"".join(value_frames), np.uint8),
+        "gaps": np.frombuffer(b"".join(g for g, _ in compressed), np.uint8),
+        "values": np.frombuffer(b"".join(v for _, v in compressed), np.uint8),
     }
-    frames = [[name, count, len(g), len(v)] for (name, count), g, v in zip(entries, gap_frames, value_frames)]
-    save_file(arrays, str(path), metadata={"entries": json.dumps(frames)})
+    metadata = {
+        "entries": json.dumps(entries),
+        "frames": json.dumps([[len(frame), len(g), len(v)] for frame, (g, v) in zip(frames, compressed)]),
+        "checksums": json.dumps(checksums),
+    }
+    save_file(arrays, str(path), metadata=metadata)
     return sum(array.nbytes for array in arrays.values())
 
 
@@ -92,37 +118,54 @@ def decode_gaps(gaps: torch.Tensor, counts: torch.Tensor) -> torch.Tensor:
 
 def _read_rank_file(
     path: Path, tensors: dict, spans: dict[str, tuple[int, int]] | None
-) -> tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]:
-    with safe_open(str(path), framework="np") as f:
-        frames = json.loads(f.metadata()["entries"])
-        gaps_bytes, values_bytes = f.get_slice("gaps"), f.get_slice("values")
-        entries, gaps, values = [], [], []
-        gaps_offset = values_offset = 0
-        for name, count, num_gap_bytes, num_value_bytes in frames:
-            gaps_end, values_end = gaps_offset + num_gap_bytes, values_offset + num_value_bytes
-            if spans is None or name in spans:
-                entry_gaps = np.frombuffer(pyzstd.decompress(gaps_bytes[gaps_offset:gaps_end].tobytes()), np.int32)
-                entry_values = np.frombuffer(
-                    pyzstd.decompress(values_bytes[values_offset:values_end].tobytes()), np.uint8
-                )
+) -> tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor, dict[str, int]]:
+    with open(path, "rb") as file:
+        header_size = int.from_bytes(file.read(8), "little")
+        header = json.loads(file.read(header_size))
+        metadata = header["__metadata__"]
+        gaps_start = 8 + header_size + header["gaps"]["data_offsets"][0]
+        values_start = 8 + header_size + header["values"]["data_offsets"][0]
+        entries = [tuple(entry) for entry in json.loads(metadata["entries"])]
+        frames, start, gaps_offset, values_offset = [], 0, gaps_start, values_start
+        for num_entries, num_gap_bytes, num_value_bytes in json.loads(metadata["frames"]):
+            frame_entries = entries[start : start + num_entries]
+            if spans is None or any(name in spans for name, _ in frame_entries):
+                frames.append((frame_entries, (gaps_offset, num_gap_bytes), (values_offset, num_value_bytes)))
+            start += num_entries
+            gaps_offset += num_gap_bytes
+            values_offset += num_value_bytes
+
+        def read_frame(frame) -> list[tuple[str, np.ndarray, np.ndarray]]:
+            frame_entries, (gap_at, gap_size), (value_at, value_size) = frame
+            frame_gaps = np.frombuffer(pyzstd.decompress(os.pread(file.fileno(), gap_size, gap_at)), np.int32)
+            frame_values = np.frombuffer(pyzstd.decompress(os.pread(file.fileno(), value_size, value_at)), np.uint8)
+            out, gap_offset, value_offset = [], 0, 0
+            for name, count in frame_entries:
+                itemsize = bits_dtype(tensors[name][1]).itemsize
+                entry_gaps = frame_gaps[gap_offset : gap_offset + count]
+                entry_values = frame_values[value_offset : value_offset + count * itemsize]
+                gap_offset += count
+                value_offset += count * itemsize
+                if spans is not None and name not in spans:
+                    continue
                 # Trim to this reader's index range unless it holds the whole tensor
                 span_start, span_end = spans[name] if spans is not None else (0, 2**62)
                 if span_start > 0 or span_end < prod(tensors[name][0]):
-                    indices = np.cumsum(entry_gaps, dtype=np.int64) + np.arange(1, count + 1) - 1
+                    indices = np.cumsum(entry_gaps, dtype=np.int64) + np.arange(count)
                     begin, end = np.searchsorted(indices, (span_start, span_end))
-                    itemsize = bits_dtype(tensors[name][1]).itemsize
                     entry_values = entry_values[begin * itemsize : end * itemsize]
                     entry_gaps = entry_gaps[begin:end].copy()
                     if end > begin:
                         entry_gaps[0] = indices[begin]
                 if len(entry_gaps):
-                    entries.append((name, len(entry_gaps)))
-                    gaps.append(entry_gaps)
-                    values.append(entry_values)
-            gaps_offset, values_offset = gaps_end, values_end
-    gaps = torch.from_numpy(np.concatenate(gaps) if gaps else np.empty(0, np.int32))
-    values = torch.from_numpy(np.concatenate(values) if values else np.empty(0, np.uint8))
-    return entries, gaps, values
+                    out.append((name, entry_gaps, entry_values))
+            return out
+
+        with ThreadPoolExecutor(4) as pool:
+            kept = [entry for frame in pool.map(read_frame, frames) for entry in frame]
+    gaps = torch.from_numpy(np.concatenate([g for _, g, _ in kept]) if kept else np.empty(0, np.int32))
+    values = torch.from_numpy(np.concatenate([v for _, _, v in kept]) if kept else np.empty(0, np.uint8))
+    return [(name, len(g)) for name, g, _ in kept], gaps, values, json.loads(metadata["checksums"])
 
 
 def write_anchor_shards(state_dict: dict[str, torch.Tensor], step_dir: Path, rank: int, world_size: int) -> None:
@@ -169,12 +212,28 @@ def read_manifest(step_dir: Path) -> dict:
 
 def read_delta(
     step_dir: Path, spans: dict[str, tuple[int, int]] | None = None
-) -> list[tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]]:
+) -> tuple[list[tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]], dict[str, int]]:
     """Read and decompress every rank file of a delta step dir into ``(entries,
     gaps, xor value bytes)`` in host memory; ``decode_gaps`` turns the
     gaps into flat indices on the device that applies them. With ``spans``, keep
-    only the listed tensors and, within each, the flat indices in ``[start, end)``."""
+    only the listed tensors and, within each, the flat indices in ``[start, end)``.
+    Also returns the new weights' ``bits_checksum`` per HF tensor, summed over ranks."""
     manifest = read_manifest(step_dir)
     paths = [step_dir / DELTA_DIR / name for name in manifest["files"]]
     with ThreadPoolExecutor(len(paths)) as pool:
-        return list(pool.map(lambda path: _read_rank_file(path, manifest["tensors"], spans), paths))
+        files = list(pool.map(lambda path: _read_rank_file(path, manifest["tensors"], spans), paths))
+    checksums: dict[str, int] = {}
+    for *_, partial in files:
+        for name, value in partial.items():
+            checksums[name] = add_checksums(checksums.get(name, 0), value)
+    return [delta for *delta, _ in files], checksums
+
+
+def bits_checksum(bits: torch.Tensor) -> int:
+    """Wrapping int64 sum of raw bits: additive over any split of a tensor, so trainer
+    ranks and engine workers can each sum the part they hold."""
+    return int(bits.sum(dtype=torch.int64))
+
+
+def add_checksums(a: int, b: int) -> int:
+    return (a + b + 2**63) % 2**64 - 2**63
diff --git a/tests/unit/utils/test_sparse_weights.py b/tests/unit/utils/test_sparse_weights.py
index 9841f607575..70daf045f52 100644
--- a/tests/unit/utils/test_sparse_weights.py
+++ b/tests/unit/utils/test_sparse_weights.py
@@ -6,6 +6,7 @@
     DELTA_DIR,
     MANIFEST,
     as_bits,
+    bits_checksum,
     decode_gaps,
     encode_gaps,
     read_delta,
@@ -25,8 +26,9 @@ def test_sparse_delta_roundtrip_is_exact(tmp_path):
     delta_dir = tmp_path / DELTA_DIR
     delta_dir.mkdir()
     for rank, names in enumerate([[("a", slice(0, 32))], [("a", slice(32, 64)), ("b", slice(None))]]):
-        entries, gaps, values = [], [], []
+        entries, gaps, values, checksums = [], [], [], {}
         for name, rows in names:
+            checksums[name] = bits_checksum(as_bits(current[name][rows]))
             offset = 0 if rows.start is None else rows.start * 32
             xor = as_bits(current[name][rows]).reshape(-1) ^ as_bits(previous[name][rows]).reshape(-1)
             changed = xor.nonzero().squeeze(1)
@@ -35,15 +37,17 @@ def test_sparse_delta_roundtrip_is_exact(tmp_path):
             entries.append((name, len(changed)))
             gaps.append(encode_gaps(changed + offset, starts))
             values.append(xor[changed].view(torch.uint8))
-        write_rank_file(
-            delta_dir / f"rank_{rank}.safetensors", entries, torch.cat(gaps).numpy(), torch.cat(values).numpy(), tensors
-        )
+        gaps, values = torch.cat(gaps).numpy(), torch.cat(values).numpy()
+        # Small frames, so a frame boundary falls inside rank 1's entries
+        write_rank_file(delta_dir / f"rank_{rank}.safetensors", entries, gaps, values, tensors, checksums, frame_size=8)
     files = ["rank_0.safetensors", "rank_1.safetensors"]
     (delta_dir / MANIFEST).write_text(json.dumps({"files": files, "tensors": tensors}))
 
     def apply(spans):
         patched = {name: as_bits(tensor.clone()).view(-1) for name, tensor in previous.items()}
-        for entries, gaps, values in read_delta(tmp_path, spans):
+        delta, checksums = read_delta(tmp_path, spans)
+        assert checksums == {name: bits_checksum(as_bits(tensor)) for name, tensor in current.items()}
+        for entries, gaps, values in delta:
             indices = decode_gaps(gaps, torch.tensor([count for _, count in entries]))
             i = b = 0
             for name, count in entries:

From 66917af82a84c8e58beaae779c548fd2671cce82 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 14 Oct 2026 18:20:14 +0000
Subject: [PATCH 08/88] perf(weights): drop sparse deltas and anchors from the
 engine host's page cache after reading

---
 src/prime_rl/inference/vllm/worker/filesystem.py | 5 +++++
 src/prime_rl/utils/sparse_weights.py             | 6 ++++++
 2 files changed, 11 insertions(+)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 7a98ce63b01..55dc927733c 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -26,6 +26,7 @@
     anchor_complete,
     bits_dtype,
     decode_gaps,
+    drop_page_cache,
     read_delta,
     read_manifest,
 )
@@ -249,3 +250,7 @@ def _load_checkpoint(self, step_dir: Path) -> None:
             self.model_runner.model_config,
             self.vllm_config,
         )
+        if anchor_complete(step_dir):
+            for path in step_dir.glob("*.safetensors"):
+                with open(path, "rb") as file:
+                    drop_page_cache(file.fileno())
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index 01720d8c729..acb6c3e241d 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -163,6 +163,7 @@ def read_frame(frame) -> list[tuple[str, np.ndarray, np.ndarray]]:
 
         with ThreadPoolExecutor(4) as pool:
             kept = [entry for frame in pool.map(read_frame, frames) for entry in frame]
+        drop_page_cache(file.fileno())
     gaps = torch.from_numpy(np.concatenate([g for _, g, _ in kept]) if kept else np.empty(0, np.int32))
     values = torch.from_numpy(np.concatenate([v for _, _, v in kept]) if kept else np.empty(0, np.uint8))
     return [(name, len(g)) for name, g, _ in kept], gaps, values, json.loads(metadata["checksums"])
@@ -229,6 +230,11 @@ def read_delta(
     return [delta for *delta, _ in files], checksums
 
 
+def drop_page_cache(fd: int) -> None:
+    """Deltas and anchors are read once; don't let them fill the inference host's page cache."""
+    os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)
+
+
 def bits_checksum(bits: torch.Tensor) -> int:
     """Wrapping int64 sum of raw bits: additive over any split of a tensor, so trainer
     ranks and engine workers can each sum the part they hold."""

From 5bcb2b038224ac9efb3c924dec0374f3c8598c6a Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 14 Oct 2026 18:22:04 +0000
Subject: [PATCH 09/88] perf(weights): pin the anchor host buffer on a thread
 ahead of each anchor and release it after the write

---
 src/prime_rl/transports/weights/filesystem.py | 31 ++++++++++++++-----
 1 file changed, 24 insertions(+), 7 deletions(-)

diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index 85b8dbfad5b..81435295ed4 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -2,6 +2,7 @@
 import threading
 import time
 from collections import defaultdict
+from concurrent.futures import Future, ThreadPoolExecutor
 from math import prod
 from pathlib import Path
 
@@ -263,8 +264,10 @@ def __init__(
         self.anchor_step: int | None = None
         self.num_deltas = 0
         self.pending_anchor: int | None = None
-        # Reused across anchors: a new one starts only after the previous one was written
-        self.anchor_buffer: torch.Tensor | None = None
+        # Pinned host buffer for the next anchor, allocated on a thread ahead of it
+        self.anchor_buffer: Future | None = None
+        self.anchor_bytes: int | None = None
+        self.allocator = ThreadPoolExecutor(1)
         self.logger.debug("Initialized filesystem weight broadcast")
 
     def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
@@ -311,6 +314,10 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
             self.anchor_step = step
             return
         # Periodic anchors only serve engines that fall off the chain: write them off the critical path.
+        # Pinning their host buffer is slow, so it starts on a thread two steps ahead.
+        if 0 < -step % self.config.anchor_interval <= 2 and self.anchor_buffer is None:
+            self.anchor_bytes = self.anchor_bytes or owned_weight_bytes(model)
+            self.anchor_buffer = self.allocator.submit(self._allocate_anchor_buffer, model)
         if (
             self.world.is_master
             and self.pending_anchor is not None
@@ -330,25 +337,35 @@ def _save_full(self, model: nn.Module, step_dir: Path) -> None:
         save_state_dict_parallel(state_dict, step_dir)
 
     def _save_anchor(self, model: nn.Module, step_dir: Path) -> threading.Thread:
-        """Gather this rank's share of the full checkpoint to CPU, then write it on a thread
-        (no collectives); engines treat the step as an anchor once every rank's marker exists."""
+        """Gather this rank's share of the full checkpoint into a pinned host buffer, then write it
+        on a thread (no collectives) and release the buffer; engines treat the step as an anchor
+        once every rank's marker exists."""
         start = time.perf_counter()
-        if self.anchor_buffer is None:
-            self.anchor_buffer = torch.empty(owned_weight_bytes(model), dtype=torch.uint8, pin_memory=True)
-        state_dict = convert_state_dict_to_hf(model, gather_weights_parallel(model, host_buffer=self.anchor_buffer))
+        buffer = self.anchor_buffer.result() if self.anchor_buffer is not None else self._allocate_anchor_buffer(model)
+        self.anchor_buffer = None
+        state_dict = convert_state_dict_to_hf(model, gather_weights_parallel(model, host_buffer=buffer))
         torch.cuda.empty_cache()
         gather_time = time.perf_counter() - start
 
         def write() -> None:
+            nonlocal buffer
             write_anchor_shards(state_dict, step_dir, self.world.rank, self.world.world_size)
             self.logger.debug(
                 f"Wrote the {step_dir.name} anchor: gather {gather_time:.1f}s, write {time.perf_counter() - start - gather_time:.1f}s"
             )
+            state_dict.clear()
+            buffer = None
+            torch._C._host_emptyCache()  # return the pinned buffer to the OS
 
         thread = threading.Thread(target=write)
         thread.start()
         return thread
 
+    def _allocate_anchor_buffer(self, model: nn.Module) -> torch.Tensor:
+        if self.anchor_bytes is None:
+            self.anchor_bytes = owned_weight_bytes(model)
+        return torch.empty(self.anchor_bytes, dtype=torch.uint8, pin_memory=True)
+
     def _save_delta(self, model: nn.Module, step: int, step_dir: Path) -> bool:
         """Write this step's delta (all ranks). False when the step must be a full checkpoint instead."""
         assert self.encoder is not None

From 1615bbe8142cfbf83a0d95dba5f70b3a272b8e9b Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 14 Oct 2026 19:47:52 +0000
Subject: [PATCH 10/88] perf(weights): checksum held params with one cumsum per
 param; filter prefetch by tensor only

---
 .../inference/vllm/worker/filesystem.py       | 26 +++----
 .../inference/vllm/worker/weight_transfer.py  | 77 +++++++++++--------
 src/prime_rl/utils/sparse_weights.py          | 24 ++----
 tests/unit/utils/test_sparse_weights.py       | 12 ++-
 4 files changed, 71 insertions(+), 68 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 55dc927733c..3f1680502ae 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -82,7 +82,7 @@ def _read_delta(
         """Read this worker's part of a delta into a pinned host buffer that is reused across steps
         (allocating pinned memory every step is slow on a loaded host)."""
         start = time.perf_counter()
-        delta, checksums = read_delta(step_dir, self._spans())
+        delta, checksums = read_delta(step_dir, None if self.index_map is None else self.index_map.names)
         read_time = time.perf_counter() - start
         num_bytes = sum(-(-gaps.nbytes // 4) * 4 + -(-values.nbytes // 4) * 4 for _, gaps, values in delta)
         if self.host_buffer is None or self.host_buffer.numel() < num_bytes:
@@ -102,10 +102,6 @@ def _read_delta(
         )
         return pinned, checksums
 
-    def _spans(self) -> dict[str, tuple[int, int]] | None:
-        """The checkpoint index ranges this worker's params hold, once the index map exists."""
-        return None if self.index_map is None else self.index_map.spans
-
     def prefetch_weights_done(self) -> bool:
         return self.prefetched is None or self.prefetched[1].done()
 
@@ -218,16 +214,20 @@ def _verify(self, step_dir: Path, checksums: dict[str, int]) -> None:
         """Check every tensor this worker holds against the trainer's checksum of the new weights.
         Tensors split across TP workers are summed over the TP group first."""
         assert self.index_map is not None
-        sums, partial = self.index_map.held_checksums(checksums)
+        names, totals, partial = self.index_map.held_checksums()
         tp_group = get_tp_group()
-        partial_names = sorted(partial) if tp_group.world_size > 1 else []
-        if partial_names:
-            reduced = torch.stack([sums[name] for name in partial_names])
+        if tp_group.world_size > 1 and partial.any():
+            reduced = totals[partial]
             torch.distributed.all_reduce(reduced, group=tp_group.device_group)
-            sums.update(zip(partial_names, reduced))
-        names = [name for name in sums if name not in partial or partial_names]
-        values = torch.stack([sums[name] for name in names]).tolist() if names else []
-        mismatched = [name for name, value in zip(names, values) if value != checksums[name]]
+            totals[partial] = reduced
+        elif partial.any():
+            totals[partial] = 0
+            names = [name if not is_partial else None for name, is_partial in zip(names, partial.tolist())]
+        mismatched = [
+            name
+            for name, value in zip(names, totals.tolist())
+            if name is not None and name in checksums and value != checksums[name]
+        ]
         if mismatched:
             raise RuntimeError(f"Sparse weight delta {step_dir.name} diverged from the trainer: {mismatched[:8]}")
 
diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py
index 1d8e503dccd..07d3a37a83a 100644
--- a/src/prime_rl/inference/vllm/worker/weight_transfer.py
+++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py
@@ -116,39 +116,54 @@ def placeholders():
         }
         self.num_cols = {name: (shape[-1] if shape else 1) for name, (shape, _) in tensors.items()}
         self.numel = {name: prod(shape) for name, (shape, _) in tensors.items()}
+        self.checksum_plan = None
         self.unpatchable = unpatchable
-        # [first, last + 1) flat checkpoint index each tensor's writes cover; unpatchable tensors are
-        # read whole so that a change to them is caught in apply_xor
-        self.spans = {
-            name: (
-                min(row0 * self.num_cols[name] + col0 for row0, _, col0, *_ in writes),
-                max((row0 + rows - 1) * self.num_cols[name] + col0 + cols for row0, rows, col0, cols, *_ in writes),
-            )
-            for name, writes in self.records.items()
-        } | {name: (0, 2**62) for name in unpatchable}
+        # Checkpoint tensors this worker holds; unpatchable ones too, so a change to them is caught in apply_xor
+        self.names = set(self.records) | unpatchable
 
-    def held_checksums(self, names) -> tuple[dict[str, torch.Tensor], set[str]]:
-        """``bits_checksum`` of the live param elements each checkpoint tensor was written to, for the
-        ``names`` this worker writes, and the subset it holds only part of (TP slices)."""
-        sums, partial = {}, set()
-        for name in names:
-            writes = self.records.get(name)
-            if writes is None or name in self.unpatchable:
-                continue
-            held = sum(rows * cols for _, rows, _, cols, *_ in writes)
-            if held > self.numel[name]:
-                continue
-            parts = [
-                as_bits(self.params[index].data)
-                .view(-1)
-                .as_strided((rows, cols), strides, offset)
-                .sum(dtype=torch.int64)
-                for _, rows, _, cols, index, offset, *strides in writes
-            ]
-            sums[name] = torch.stack(parts).sum()
-            if held < self.numel[name]:
-                partial.add(name)
-        return sums, partial
+    def held_checksums(self) -> tuple[list[str], torch.Tensor, torch.Tensor]:
+        """``bits_checksum`` of the live param elements each checkpoint tensor was written to: the
+        names, their sums, and a mask of the ones this worker holds only part of (TP slices)."""
+        if self.checksum_plan is None:
+            self.checksum_plan = self._plan_checksums()
+        names, partial, row_ranges, other_writes = self.checksum_plan
+        totals = torch.zeros(len(names), dtype=torch.int64, device=partial.device)
+        for index, (starts, ends, ids) in row_ranges.items():
+            param = self.params[index]
+            row_sums = as_bits(param.data).reshape(-1, param.shape[-1]).sum(dim=1, dtype=torch.int64)
+            cumulative = torch.cat([row_sums.new_zeros(1), row_sums.cumsum(0)])
+            totals.index_add_(0, ids, cumulative[ends] - cumulative[starts])
+        for name_id, rows, cols, index, offset, *strides in other_writes:
+            view = as_bits(self.params[index].data).view(-1).as_strided((rows, cols), strides, offset)
+            totals[name_id] += view.sum(dtype=torch.int64)
+        return names, totals, partial
+
+    def _plan_checksums(self):
+        """Writes that cover whole rows of a param become row ranges, summed with one cumsum per param."""
+        names = [
+            name
+            for name, writes in self.records.items()
+            if name not in self.unpatchable and sum(rows * cols for _, rows, _, cols, *_ in writes) <= self.numel[name]
+        ]
+        device = self.params[0].device
+        ranges: dict[int, list[list[int]]] = defaultdict(lambda: [[], [], []])
+        other_writes = []
+        for name_id, name in enumerate(names):
+            for _, rows, _, cols, index, offset, row_stride, col_stride in self.records[name]:
+                width = self.params[index].shape[-1]
+                if col_stride == 1 and row_stride == width and cols == width and offset % width == 0:
+                    starts, ends, ids = ranges[index]
+                    starts.append(offset // width)
+                    ends.append(offset // width + rows)
+                    ids.append(name_id)
+                else:
+                    other_writes.append((name_id, rows, cols, index, offset, row_stride, col_stride))
+        row_ranges = {
+            index: tuple(torch.tensor(values, device=device) for values in lists) for index, lists in ranges.items()
+        }
+        held = [sum(rows * cols for _, rows, _, cols, *_ in self.records[name]) for name in names]
+        partial = torch.tensor([h < self.numel[name] for h, name in zip(held, names)], device=device)
+        return names, partial, row_ranges, other_writes
 
     def fp8_layout(self, read_tensor) -> dict[str, str]:
         """The vLLM per-block quantization (``"fp32"`` or ``"ue8m0"`` scales) that reproduces each
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index acb6c3e241d..2a2842ee270 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -12,7 +12,6 @@
 import json
 import os
 from concurrent.futures import ThreadPoolExecutor
-from math import prod
 from pathlib import Path
 
 import numpy as np
@@ -117,7 +116,7 @@ def decode_gaps(gaps: torch.Tensor, counts: torch.Tensor) -> torch.Tensor:
 
 
 def _read_rank_file(
-    path: Path, tensors: dict, spans: dict[str, tuple[int, int]] | None
+    path: Path, tensors: dict, names: set[str] | None
 ) -> tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor, dict[str, int]]:
     with open(path, "rb") as file:
         header_size = int.from_bytes(file.read(8), "little")
@@ -129,7 +128,7 @@ def _read_rank_file(
         frames, start, gaps_offset, values_offset = [], 0, gaps_start, values_start
         for num_entries, num_gap_bytes, num_value_bytes in json.loads(metadata["frames"]):
             frame_entries = entries[start : start + num_entries]
-            if spans is None or any(name in spans for name, _ in frame_entries):
+            if names is None or any(name in names for name, _ in frame_entries):
                 frames.append((frame_entries, (gaps_offset, num_gap_bytes), (values_offset, num_value_bytes)))
             start += num_entries
             gaps_offset += num_gap_bytes
@@ -146,17 +145,8 @@ def read_frame(frame) -> list[tuple[str, np.ndarray, np.ndarray]]:
                 entry_values = frame_values[value_offset : value_offset + count * itemsize]
                 gap_offset += count
                 value_offset += count * itemsize
-                if spans is not None and name not in spans:
+                if names is not None and name not in names:
                     continue
-                # Trim to this reader's index range unless it holds the whole tensor
-                span_start, span_end = spans[name] if spans is not None else (0, 2**62)
-                if span_start > 0 or span_end < prod(tensors[name][0]):
-                    indices = np.cumsum(entry_gaps, dtype=np.int64) + np.arange(count)
-                    begin, end = np.searchsorted(indices, (span_start, span_end))
-                    entry_values = entry_values[begin * itemsize : end * itemsize]
-                    entry_gaps = entry_gaps[begin:end].copy()
-                    if end > begin:
-                        entry_gaps[0] = indices[begin]
                 if len(entry_gaps):
                     out.append((name, entry_gaps, entry_values))
             return out
@@ -212,17 +202,17 @@ def read_manifest(step_dir: Path) -> dict:
 
 
 def read_delta(
-    step_dir: Path, spans: dict[str, tuple[int, int]] | None = None
+    step_dir: Path, names: set[str] | None = None
 ) -> tuple[list[tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]], dict[str, int]]:
     """Read and decompress every rank file of a delta step dir into ``(entries,
     gaps, xor value bytes)`` in host memory; ``decode_gaps`` turns the
-    gaps into flat indices on the device that applies them. With ``spans``, keep
-    only the listed tensors and, within each, the flat indices in ``[start, end)``.
+    gaps into flat indices on the device that applies them. With ``names``, keep
+    only those tensors (frames without any are not read).
     Also returns the new weights' ``bits_checksum`` per HF tensor, summed over ranks."""
     manifest = read_manifest(step_dir)
     paths = [step_dir / DELTA_DIR / name for name in manifest["files"]]
     with ThreadPoolExecutor(len(paths)) as pool:
-        files = list(pool.map(lambda path: _read_rank_file(path, manifest["tensors"], spans), paths))
+        files = list(pool.map(lambda path: _read_rank_file(path, manifest["tensors"], names), paths))
     checksums: dict[str, int] = {}
     for *_, partial in files:
         for name, value in partial.items():
diff --git a/tests/unit/utils/test_sparse_weights.py b/tests/unit/utils/test_sparse_weights.py
index 70daf045f52..fad2c4fc915 100644
--- a/tests/unit/utils/test_sparse_weights.py
+++ b/tests/unit/utils/test_sparse_weights.py
@@ -43,9 +43,9 @@ def test_sparse_delta_roundtrip_is_exact(tmp_path):
     files = ["rank_0.safetensors", "rank_1.safetensors"]
     (delta_dir / MANIFEST).write_text(json.dumps({"files": files, "tensors": tensors}))
 
-    def apply(spans):
+    def apply(names):
         patched = {name: as_bits(tensor.clone()).view(-1) for name, tensor in previous.items()}
-        delta, checksums = read_delta(tmp_path, spans)
+        delta, checksums = read_delta(tmp_path, names)
         assert checksums == {name: bits_checksum(as_bits(tensor)) for name, tensor in current.items()}
         for entries, gaps, values in delta:
             indices = decode_gaps(gaps, torch.tensor([count for _, count in entries]))
@@ -63,9 +63,7 @@ def apply(spans):
     for name, tensor in current.items():
         assert torch.equal(patched[name], as_bits(tensor).view(-1))
 
-    # A reader that holds rows 16-47 of "a" (spanning both rank files) and none of "b".
-    patched = apply({"a": (16 * 32, 48 * 32)})
-    expected = as_bits(previous["a"].clone())
-    expected[16:48] = as_bits(current["a"][16:48])
-    assert torch.equal(patched["a"], expected.view(-1))
+    # A reader that holds only "a" (split over both rank files) skips the frame holding "b".
+    patched = apply({"a"})
+    assert torch.equal(patched["a"], as_bits(current["a"]).view(-1))
     assert torch.equal(patched["b"], as_bits(previous["b"]).view(-1))

From 8d4d5e3182273bfbf9c5b8905a22c2aa714faf15 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 14 Oct 2026 20:46:02 +0000
Subject: [PATCH 11/88] fix(weights): engines off the delta chain replay from
 the latest complete anchor

---
 src/prime_rl/inference/vllm/worker/filesystem.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 3f1680502ae..5cb4b4ae4fb 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -160,13 +160,13 @@ def read_tensor(name: str) -> torch.Tensor:
         logger.info(f"Wrote the engine layout: {len(layout['fp8'])} checkpoint tensors held in fp8")
 
     def _replay_from_anchor(self, step_dir: Path) -> None:
-        chain = [step_dir]
-        manifest = read_manifest(step_dir)
-        while manifest["base_step"] != manifest["anchor_step"]:
-            chain.append(step_dir.parent / f"step_{manifest['base_step']}")
-            manifest = read_manifest(chain[-1])
-        logger.info(f"Replaying {len(chain)} sparse weight deltas on top of the v{manifest['anchor_step']} anchor")
-        self._load_checkpoint(step_dir.parent / f"step_{manifest['anchor_step']}")
+        """Load the latest complete anchor at or before ``step_dir`` and replay the deltas since."""
+        chain = []
+        while not anchor_complete(step_dir):
+            chain.append(step_dir)
+            step_dir = step_dir.parent / f"step_{read_manifest(step_dir)['base_step']}"
+        logger.info(f"Replaying {len(chain)} sparse weight deltas on top of the {step_dir.name} anchor")
+        self._load_checkpoint(step_dir)
         for delta_dir in reversed(chain):
             self._apply_delta(delta_dir)
 

From a5b97348cd5b1006d304b5081adc54efa388b4ff Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 14 Oct 2026 23:45:31 +0000
Subject: [PATCH 12/88] feat(weights): boot weights as the first sparse step,
 background delta writes, uint8 gap frames

---
 .../inference/vllm/worker/filesystem.py       |  62 +++--
 src/prime_rl/transports/weights/base.py       |   8 +-
 src/prime_rl/transports/weights/filesystem.py | 249 ++++++++++--------
 src/prime_rl/utils/sparse_weights.py          | 140 +++++-----
 4 files changed, 265 insertions(+), 194 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 5cb4b4ae4fb..eb87c554854 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -111,34 +111,54 @@ def update_weights_from_path(self, weight_path: str) -> None:
         step = int(step_dir.name.removeprefix("step_"))
         if not (step_dir / DELTA_DIR / MANIFEST).exists():
             self._load_checkpoint(step_dir)
-        elif read_manifest(step_dir)["base_step"] == self.weight_step:
-            self._apply_delta(step_dir)
-            if anchor_complete(step_dir):
-                # Anchor step: check that the delta chain reproduced the full checkpoint bit for bit.
-                chained = weight_checksums(self._model())
-                self._load_checkpoint(step_dir)
-                mismatched = [name for name, value in weight_checksums(self._model()).items() if chained[name] != value]
-                if mismatched:
-                    raise RuntimeError(f"Sparse weight deltas diverged from the v{step} anchor: {mismatched[:8]}")
-                logger.info(f"Sparse weight delta chain matches the v{step} anchor")
         else:
-            self._replay_from_anchor(step_dir)
+            base_step = read_manifest(step_dir)["base_step"]
+            if base_step is None:
+                # The run's first broadcast: the weights the engine booted with.
+                if self.weight_step is not None:
+                    self._load_checkpoint(self._boot_dir())
+            elif base_step == self.weight_step:
+                self._apply_delta(step_dir)
+                if anchor_complete(step_dir):
+                    # Anchor step: check that the delta chain reproduced the full checkpoint bit for bit.
+                    chained = weight_checksums(self._model())
+                    self._load_checkpoint(step_dir)
+                    current = weight_checksums(self._model())
+                    mismatched = [name for name, value in current.items() if chained[name] != value]
+                    if mismatched:
+                        raise RuntimeError(f"Sparse weight deltas diverged from the v{step} anchor: {mismatched[:8]}")
+                    logger.info(f"Sparse weight delta chain matches the v{step} anchor")
+            else:
+                self._replay(step_dir)
         self.weight_step = step
-        if (step_dir / DELTA_DIR).is_dir() and not self.wrote_layout and anchor_complete(step_dir):
+        if (step_dir / DELTA_DIR).is_dir() and not self.wrote_layout:
             self._write_engine_layout(step_dir)
             self.wrote_layout = True
 
+    def _boot_dir(self) -> Path:
+        """The local dir of the checkpoint the engine booted from."""
+        model_loader = get_model_loader(self.load_config)
+        assert isinstance(model_loader, DefaultModelLoader)
+        folder, *_ = model_loader._prepare_weights(
+            self.model_config.model, None, self.model_config.revision, True, None
+        )
+        return Path(folder)
+
     def _write_engine_layout(self, step_dir: Path) -> None:
         """Tell the trainer which checkpoint tensors this engine holds in fp8 (and with which
-        scales), so it can send their deltas in the fp8 domain. Needs a full checkpoint."""
+        scales), so it can send their deltas in the fp8 domain. Reads one checkpoint tensor per fp8
+        param from the step's full checkpoint, or from the boot checkpoint."""
+        checkpoint_dir = step_dir if anchor_complete(step_dir) else self._boot_dir()
         files = {}
         tensors = {}
-        for path in step_dir.glob("*.safetensors"):
+        for path in checkpoint_dir.glob("*.safetensors"):
             with safe_open(path, framework="pt") as f:
                 for name in f.keys():
                     tensor_slice = f.get_slice(name)
                     files[name] = path
                     tensors[name] = (tensor_slice.get_shape(), _SAFETENSORS_DTYPES[tensor_slice.get_dtype()])
+        if (step_dir / DELTA_DIR / MANIFEST).exists():
+            tensors = read_manifest(step_dir)["tensors"]
         if self.index_map is None:
             self.index_map = CheckpointIndexMap(self._model(), tensors)
 
@@ -159,14 +179,18 @@ def read_tensor(name: str) -> torch.Tensor:
         os.replace(tmp, path)
         logger.info(f"Wrote the engine layout: {len(layout['fp8'])} checkpoint tensors held in fp8")
 
-    def _replay_from_anchor(self, step_dir: Path) -> None:
-        """Load the latest complete anchor at or before ``step_dir`` and replay the deltas since."""
+    def _replay(self, step_dir: Path) -> None:
+        """Load the latest complete anchor at or before ``step_dir`` (or the boot checkpoint, for a
+        chain from the run's first broadcast) and replay the deltas since."""
         chain = []
         while not anchor_complete(step_dir):
+            base_step = read_manifest(step_dir)["base_step"]
+            if base_step is None:
+                break
             chain.append(step_dir)
-            step_dir = step_dir.parent / f"step_{read_manifest(step_dir)['base_step']}"
-        logger.info(f"Replaying {len(chain)} sparse weight deltas on top of the {step_dir.name} anchor")
-        self._load_checkpoint(step_dir)
+            step_dir = step_dir.parent / f"step_{base_step}"
+        logger.info(f"Replaying {len(chain)} sparse weight deltas on top of {step_dir.name}")
+        self._load_checkpoint(step_dir if anchor_complete(step_dir) else self._boot_dir())
         for delta_dir in reversed(chain):
             self._apply_delta(delta_dir)
 
diff --git a/src/prime_rl/transports/weights/base.py b/src/prime_rl/transports/weights/base.py
index e9da73300d9..bb51f23f837 100644
--- a/src/prime_rl/transports/weights/base.py
+++ b/src/prime_rl/transports/weights/base.py
@@ -66,8 +66,7 @@ def broadcast(self, model: nn.Module, step: int) -> None:
         wait_for_master(f"weight_broadcast/{step}", self.timeout)
         self._broadcast(model, step, step_dir)
         if self.world.is_master:
-            (step_dir / FINISHED_MARKER).touch()
-            self._clean(step)
+            self._finish(step, step_dir)
             self.logger.debug(f"Broadcasted weights for step {step} in {time.perf_counter() - start_time:.2f}s")
 
     def step_dir(self, step: int) -> Path:
@@ -98,6 +97,11 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
         """Move v{step}'s weights to the consumer. All ranks enter only after
         the master finished the handshake."""
 
+    def _finish(self, step: int, step_dir: Path) -> None:
+        """Mark v{step} complete for the consumer (master only)."""
+        (step_dir / FINISHED_MARKER).touch()
+        self._clean(step)
+
     def _clean(self, step: int) -> None:
         """Remove old broadcast dirs, keeping ``step`` and ``step - 1`` (a
         lagging consumer may still be reading it). Broadcasts are purely
diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index 81435295ed4..d6449ea38c3 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -45,19 +45,21 @@ class SparseDeltaEncoder:
     """Rank-local XOR deltas of the trainer's weight shards, in HF coordinates.
 
     Every rank keeps the wire-dtype bits of its own shards from the previous
-    broadcast and diffs against them on device, so no weights are gathered.
-    The prime->HF conversion is replayed once on meta tensors: every HF tensor
-    must come out as a contiguous view of one trainer tensor (renames, expert
-    unstacking), which maps a changed trainer element to its HF name and index.
+    broadcast in pinned host memory and diffs against them on device, so no
+    weights are gathered. The prime->HF conversion is replayed once on meta
+    tensors: every HF tensor must come out as a contiguous view of one trainer
+    tensor (renames, expert unstacking), which maps a changed trainer element to
+    its HF name and index.
 
-    Tensors the engines hold in fp8 (``fp8``: HF name -> scale format, from the
-    engine layout) are quantized with vLLM's own per-block quantizer and diffed
-    as fp8 bytes. Their 128-row blocks must not straddle trainer ranks, else the
-    trainer tensor is gathered. A step that changes any block scale returns no
-    delta, so it is published as a full checkpoint.
+    Tensors the engines hold in fp8 (``set_fp8_layout``: HF name -> scale format,
+    from the engine layout) are quantized with vLLM's own per-block quantizer,
+    previous and current weights alike, and diffed as fp8 bytes. Their 128-row
+    blocks must not straddle trainer ranks, else the trainer tensor is gathered.
+    A step that changes any block scale returns no delta, so it is published as
+    a full checkpoint.
     """
 
-    def __init__(self, model: nn.Module, fp8: dict[str, str]):
+    def __init__(self, model: nn.Module):
         keep_in_fp32 = getattr(model, "keep_in_fp32_for_weight_transfer", None)
         self.wire_dtypes: dict[str, torch.dtype] = {}
         meta: dict[str, torch.Tensor] = {}
@@ -77,20 +79,29 @@ def __init__(self, model: nn.Module, fp8: dict[str, str]):
             if tensor.numel() >= 2**32:
                 raise NotImplementedError(f"Sparse weight broadcast indexes tensors below 2^32 elements ({name})")
             segments[owner[id(base)]].append((tensor.storage_offset(), name))
-            dtype = "float8_e4m3fn" if name in fp8 else str(tensor.dtype).removeprefix("torch.")
-            self.hf_tensors[name] = (list(tensor.shape), dtype)
+            self.hf_tensors[name] = (list(tensor.shape), str(tensor.dtype).removeprefix("torch."))
             self.hf_numel[name] = tensor.numel()
         self.segments = {fqn: sorted(segs) for fqn, segs in segments.items()}
-
+        self.numel = {fqn: tensor.numel() for fqn, tensor in meta.items()}
         # fqn -> (ue8m0 scales, columns) for trainer tensors held in fp8; quantized as one [rows, columns] matrix
         self.fp8: dict[str, tuple[bool, int]] = {}
+        self.fp8_layout: dict[str, str] | None = None
+        self.gather: set[str] = set()
+        self.segment_starts: dict[str, torch.Tensor] = {}
+        self.snapshot: dict[str, torch.Tensor] = {}
+        self.scales_changed = False
+        self.initialized = False
+
+    def set_fp8_layout(self, model: nn.Module, fp8: dict[str, str], is_master: bool) -> None:
+        """Diff the tensors the engines hold in fp8 (HF name -> scale format) as fp8 bytes from now on."""
+        self.fp8_layout = fp8
         for fqn, segs in self.segments.items():
             modes = {fp8.get(name) for _, name in segs}
             if modes == {None}:
                 continue
             shapes = [self.hf_tensors[name][0] for _, name in segs]
             num_cols = shapes[0][-1]
-            num_rows = meta[fqn].numel() // num_cols
+            num_rows = self.numel[fqn] // num_cols
             boundaries = [start // num_cols for start, _ in segs] + [
                 start // num_cols + shape[0] for (start, _), shape in zip(segs, shapes)
             ]
@@ -101,12 +112,9 @@ def __init__(self, model: nn.Module, fp8: dict[str, str]):
             ):
                 raise NotImplementedError(f"Sparse weight broadcast cannot quantize {fqn} to fp8 as one matrix")
             self.fp8[fqn] = (modes.pop() == "ue8m0", num_cols)
-        self.gather: set[str] | None = None
-        self.segment_starts: dict[str, torch.Tensor] = {}
-        self.snapshot: dict[str, torch.Tensor] = {}
-        self.scales: dict[str, torch.Tensor] = {}
-        self.scales_changed = False
-        self.initialized = False
+            for _, name in segs:
+                self.hf_tensors[name] = (self.hf_tensors[name][0], "float8_e4m3fn")
+        self.gather = self._find_gathered(model, is_master)
 
     def _local_shard(self, key: str, value: torch.Tensor, is_master: bool) -> tuple[torch.Tensor, int] | None:
         """This rank's contiguous flat slice of ``value`` and its global flat offset,
@@ -144,26 +152,37 @@ def _find_gathered(self, model: nn.Module, is_master: bool) -> set[str]:
             dist.all_reduce(flags, op=dist.ReduceOp.MAX)
         return {key for key, flag in zip(keys, flags.tolist()) if flag}
 
-    def _fp8_bits(
-        self, key: str, fqn: str, value: torch.Tensor, shard: tuple[torch.Tensor, int] | None
-    ) -> tuple[torch.Tensor, int] | None:
-        """fp8 bits of this rank's shard, quantized as the engines do."""
+    def _fp8_diff(
+        self, key: str, fqn: str, value: torch.Tensor, bits: torch.Tensor, offset: int
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Changed flat indices, XOR values and new fp8 bits of this rank's shard, with previous and
+        current weights quantized as the engines do; advances the snapshot."""
         ue8m0, num_cols = self.fp8[fqn]
+        previous = self.snapshot[key].to(bits.device, non_blocking=True).view(torch.bfloat16)
+        current = bits.view(torch.bfloat16)
         if key in self.gather:
-            quantized, scales = quantize_fp8_blocks(value.full_tensor().to(torch.bfloat16).view(-1, num_cols), ue8m0)
-            if shard is None:
-                return None
-            bits = quantized.view(-1)[shard[1] : shard[1] + shard[0].numel()]
-        elif shard is None or shard[0].numel() == 0:
-            return None
-        else:
-            quantized, scales = quantize_fp8_blocks(shard[0].to(torch.bfloat16).view(-1, num_cols), ue8m0)
-            bits = quantized.view(-1)
-        previous = self.scales.get(key)
-        if previous is not None and not torch.equal(previous, scales):
+            local_shape = (previous.numel() // prod(value.shape[1:]), *value.shape[1:])
+            previous = DTensor.from_local(
+                previous.view(local_shape),
+                value.device_mesh,
+                value.placements,
+                shape=value.shape,
+                stride=value.stride(),
+            ).full_tensor()
+            current = value.full_tensor().to(torch.bfloat16)
+        (previous_bits, previous_scales), (current_bits, current_scales) = (
+            quantize_fp8_blocks(weights.reshape(-1, num_cols), ue8m0) for weights in (previous, current)
+        )
+        if not torch.equal(previous_scales, current_scales):
             self.scales_changed = True
-        self.scales[key] = scales
-        return bits, shard[1]
+        if key in self.gather:
+            start = offset
+            previous_bits = previous_bits.view(-1)[start : start + bits.numel()]
+            current_bits = current_bits.view(-1)[start : start + bits.numel()]
+        xor = previous_bits.view(-1) ^ current_bits.view(-1)
+        self.snapshot[key].copy_(bits, non_blocking=True)
+        changed = xor.nonzero().squeeze(1)
+        return changed, xor[changed], current_bits.view(-1)
 
     @staticmethod
     def _diff(bits: torch.Tensor, previous: torch.Tensor, chunk: int = 1 << 28) -> tuple[torch.Tensor, torch.Tensor]:
@@ -188,8 +207,6 @@ def encode(
         on the host, or None on the first call and when an fp8 block scale changed."""
         first = not self.initialized
         self.initialized = True
-        if self.gather is None:
-            self.gather = self._find_gathered(model, is_master)
         self.scales_changed = False
         names, counts, gaps, values = [], [], [], []
         checksum_names, checksums = [], []
@@ -198,25 +215,26 @@ def encode(
             if fqn not in self.segments:
                 continue
             shard = self._local_shard(key, value, is_master)
-            if fqn in self.fp8:
-                shard = self._fp8_bits(key, fqn, value, shard)
-            elif shard is not None:
-                shard = (as_bits(shard[0].to(self.wire_dtypes[key])), shard[1])
             if shard is None or shard[0].numel() == 0:
+                if key in self.gather and not first:
+                    value.full_tensor()  # every rank joins the all-gathers of a gathered fp8 tensor
+                    value.full_tensor()
                 continue
-            bits, offset = shard
-            if not first:
-                for start, name in self.segments[fqn]:
-                    begin, end = max(start, offset), min(start + self.hf_numel[name], offset + bits.numel())
-                    if begin < end:
-                        checksum_names.append(name)
-                        checksums.append(bits[begin - offset : end - offset].sum(dtype=torch.int64))
+            bits, offset = as_bits(shard[0].to(self.wire_dtypes[key])), shard[1]
             if first:
                 # The previous broadcast's bits live in pinned host memory, so the trainer's GPU memory is unchanged.
                 self.snapshot[key] = torch.empty(bits.shape, dtype=bits.dtype, pin_memory=True)
                 self.snapshot[key].copy_(bits)
                 continue
-            changed, changed_values = self._diff(bits, self.snapshot[key])
+            if fqn in self.fp8:
+                changed, changed_values, bits = self._fp8_diff(key, fqn, value, bits, offset)
+            else:
+                changed, changed_values = self._diff(bits, self.snapshot[key])
+            for start, name in self.segments[fqn]:
+                begin, end = max(start, offset), min(start + self.hf_numel[name], offset + bits.numel())
+                if begin < end:
+                    checksum_names.append(name)
+                    checksums.append(bits[begin - offset : end - offset].sum(dtype=torch.int64))
             if fqn not in self.segment_starts:
                 self.segment_starts[fqn] = torch.tensor([start for start, _ in self.segments[fqn]], device=bits.device)
             starts = self.segment_starts[fqn]
@@ -232,9 +250,10 @@ def encode(
             changed_anywhere = torch.tensor(int(self.scales_changed), device=next(model.parameters()).device)
             dist.all_reduce(changed_anywhere, op=dist.ReduceOp.MAX)
             self.scales_changed = bool(changed_anywhere.item())
+        torch.cuda.synchronize()  # the snapshot and host copies are written asynchronously
         if first or self.scales_changed:
             return None
-        counts = torch.cat(counts).tolist() if counts else []  # also waits for the copies to the host
+        counts = torch.cat(counts).tolist() if counts else []
         checksums = dict(zip(checksum_names, torch.stack(checksums).tolist())) if checksums else {}
         torch.cuda.empty_cache()
         entries = [(name, count) for name, count in zip(names, counts) if count]
@@ -245,8 +264,10 @@ def encode(
 class FileSystemWeightSender(WeightSender):
     """Broadcast weights by saving a HF-compatible checkpoint (or, for LoRA
     runs, the PEFT-shaped adapter) to a shared filesystem. With ``sparse``,
-    steps between anchors only carry the XOR delta against the previous
-    broadcast (see ``prime_rl.utils.sparse_weights``)."""
+    steps only carry the XOR delta against the previous broadcast (see
+    ``prime_rl.utils.sparse_weights``), written on a background thread while
+    training continues; full checkpoints (anchors) serve engines that fall off
+    the delta chain."""
 
     def __init__(
         self,
@@ -262,12 +283,13 @@ def __init__(
         self.encoder: SparseDeltaEncoder | None = None
         self.last_step: int | None = None
         self.anchor_step: int | None = None
-        self.num_deltas = 0
         self.pending_anchor: int | None = None
         # Pinned host buffer for the next anchor, allocated on a thread ahead of it
         self.anchor_buffer: Future | None = None
         self.anchor_bytes: int | None = None
         self.allocator = ThreadPoolExecutor(1)
+        # At most one delta is being written while the next step trains
+        self.writer: threading.Thread | None = None
         self.logger.debug("Initialized filesystem weight broadcast")
 
     def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
@@ -294,25 +316,41 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
         if not self.config.sparse:
             self._save_full(model, step_dir)
             return
+        if self.writer is not None:
+            self.writer.join()
+            self.writer = None
         if self.world.is_master:
             # Engines that load a step with a delta dir report their weight layout next to it.
             (step_dir / DELTA_DIR).mkdir(exist_ok=True)
         if self.encoder is None:
-            # Deltas start once every engine worker has reported which tensors it holds in fp8.
+            self.encoder = SparseDeltaEncoder(model)
+        if self.encoder.fp8_layout is None:
+            # Until every engine worker has reported which tensors it holds in fp8, steps are full checkpoints.
             layout = [read_engine_layout(step_dir.parent) if self.world.is_master else None]
             dist.broadcast_object_list(layout, src=0)
             if layout[0] is not None:
-                self.encoder = SparseDeltaEncoder(model, layout[0])
-        wrote_delta = self.encoder is not None and self._save_delta(model, step, step_dir)
-        self.num_deltas += wrote_delta
-        self.last_step = step
-        if not wrote_delta or self.num_deltas == 1:
-            # Engines need this full checkpoint now; the first delta's step is one too, so engines
-            # check the delta path against it right away.
+                self.encoder.set_fp8_layout(model, layout[0], self.world.is_master)
+        first = not self.encoder.initialized
+        start = time.perf_counter()
+        delta = self.encoder.encode(model, self.world.is_master)
+        encode_time = time.perf_counter() - start
+        if self.encoder.scales_changed:
+            self.logger.info(f"fp8 block scales changed at v{step}: publishing a full checkpoint")
+        base_step, self.last_step = self.last_step, step
+        if first and step == 0:
+            # The run's first broadcast is the checkpoint the engines booted from.
+            if self.world.is_master:
+                manifest = {"step": step, "base_step": None, "files": [], "tensors": self.encoder.hf_tensors}
+                (step_dir / DELTA_DIR / MANIFEST).write_text(json.dumps(manifest))
+            self.anchor_step = step
+            return
+        if delta is None or self.encoder.fp8_layout is None:
             self._save_anchor(model, step_dir).join()
             dist.barrier()
             self.anchor_step = step
             return
+        self.writer = threading.Thread(target=self._write_delta, args=(step, base_step, step_dir, delta, encode_time))
+        self.writer.start()
         # Periodic anchors only serve engines that fall off the chain: write them off the critical path.
         # Pinning their host buffer is slow, so it starts on a thread two steps ahead.
         if 0 < -step % self.config.anchor_interval <= 2 and self.anchor_buffer is None:
@@ -324,12 +362,52 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
             and anchor_complete(self.step_dir(self.pending_anchor))
         ):
             self.anchor_step, self.pending_anchor = max(self.anchor_step, self.pending_anchor), None
-        start = [step % self.config.anchor_interval == 0 and self.pending_anchor is None]
-        dist.broadcast_object_list(start, src=0)
-        if start[0]:
+        start_anchor = [step % self.config.anchor_interval == 0 and self.pending_anchor is None]
+        dist.broadcast_object_list(start_anchor, src=0)
+        if start_anchor[0]:
             self._save_anchor(model, step_dir)
             self.pending_anchor = step
 
+    def _write_delta(self, step: int, base_step: int, step_dir: Path, delta: tuple, encode_time: float) -> None:
+        """Write this rank's delta file; the master then waits for every rank's file and commits
+        the step (manifest, then the finished marker). Runs on a thread, without collectives."""
+        assert self.encoder is not None
+        start = time.perf_counter()
+        entries, gaps, values, checksums = delta
+        delta_dir = step_dir / DELTA_DIR
+        files = [f"rank_{rank}.safetensors" for rank in range(self.world.world_size)]
+        write_rank_file(
+            delta_dir / files[self.world.rank],
+            entries,
+            gaps.numpy(),
+            values.numpy(),
+            self.encoder.hf_tensors,
+            checksums,
+        )
+        if not self.world.is_master:
+            return
+        while not all((delta_dir / name).exists() for name in files):
+            time.sleep(0.05)
+        manifest = {
+            "step": step,
+            "base_step": base_step,
+            "anchor_step": self.anchor_step,
+            "files": files,
+            "tensors": self.encoder.hf_tensors,
+        }
+        (delta_dir / MANIFEST).write_text(json.dumps(manifest))
+        super()._finish(step, step_dir)
+        num_bytes = sum((delta_dir / name).stat().st_size for name in files)
+        self.logger.info(
+            f"Sparse weight delta v{step} (base v{base_step}): {num_bytes / 2**20:.1f} MiB, "
+            f"encode {encode_time:.2f}s, write {time.perf_counter() - start:.2f}s (in the background)"
+        )
+
+    def _finish(self, step: int, step_dir: Path) -> None:
+        # A delta step is committed by its writer thread once every rank's file is on disk.
+        if self.writer is None:
+            super()._finish(step, step_dir)
+
     def _save_full(self, model: nn.Module, step_dir: Path) -> None:
         state_dict = gather_weights_parallel(model)
         state_dict = convert_state_dict_to_hf(model, state_dict)
@@ -366,45 +444,6 @@ def _allocate_anchor_buffer(self, model: nn.Module) -> torch.Tensor:
             self.anchor_bytes = owned_weight_bytes(model)
         return torch.empty(self.anchor_bytes, dtype=torch.uint8, pin_memory=True)
 
-    def _save_delta(self, model: nn.Module, step: int, step_dir: Path) -> bool:
-        """Write this step's delta (all ranks). False when the step must be a full checkpoint instead."""
-        assert self.encoder is not None
-        start = time.perf_counter()
-        delta = self.encoder.encode(model, self.world.is_master)
-        if self.encoder.scales_changed:
-            self.logger.info(f"fp8 block scales changed at v{step}: publishing a full checkpoint")
-        if delta is None:
-            return False
-        entries, gaps, values, checksums = delta
-        encode_time = time.perf_counter() - start
-        file_name = f"rank_{self.world.rank}.safetensors"
-        delta_dir = step_dir / DELTA_DIR
-        dist.barrier()
-        num_bytes = write_rank_file(
-            delta_dir / file_name, entries, gaps.numpy(), values.numpy(), self.encoder.hf_tensors, checksums
-        )
-        stats = (file_name, num_bytes, len(gaps), encode_time, time.perf_counter() - start)
-        all_stats: list = [None] * self.world.world_size
-        dist.all_gather_object(all_stats, stats)
-        if not self.world.is_master:
-            return True
-        manifest = {
-            "step": step,
-            "base_step": self.last_step,
-            "anchor_step": self.anchor_step,
-            "files": [file_name for file_name, *_ in all_stats],
-            "tensors": self.encoder.hf_tensors,
-        }
-        (delta_dir / MANIFEST).write_text(json.dumps(manifest))
-        num_changes = sum(stat[2] for stat in all_stats)
-        num_elements = sum(prod(shape) for shape, _ in self.encoder.hf_tensors.values())
-        self.logger.info(
-            f"Sparse weight delta v{step} (base v{self.last_step}): {num_changes / num_elements:.3%} changed, "
-            f"{sum(stat[1] for stat in all_stats) / 2**20:.1f} MiB, "
-            f"encode {max(stat[3] for stat in all_stats):.2f}s, total {max(stat[4] for stat in all_stats):.2f}s"
-        )
-        return True
-
     def _clean(self, step: int) -> None:
         # Sparse chains replay from their anchor, so keep everything since it.
         super()._clean(step if self.anchor_step is None else min(step, self.anchor_step + 1))
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index 2a2842ee270..0977b71e26f 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -2,17 +2,19 @@
 
 A delta lives in the ``delta/`` subdir of a broadcast step dir (next to the
 full checkpoint on anchor steps): one ``rank_{r}.safetensors`` per trainer rank
-and a ``manifest.json``, written last. A rank file stores two zstd streams over its
-entries ``(hf_name, count)``: ``gaps`` (uint32 gaps between sorted flat indices,
-restarted per entry) and ``values`` (the XOR of the new and the previous raw
-bits). XOR is exact for any dtype, but applying it twice reverts it, so a delta
-must be applied on top of exactly its ``base_step``.
+and a ``manifest.json``, written last. A rank file stores zstd frames over its
+entries ``(hf_name, count)``: gaps between sorted flat indices (restarted per
+entry) and the XOR of the new and the previous raw bits. XOR is exact for any
+dtype, but applying it twice reverts it, so a delta must be applied on top of
+exactly its ``base_step``; per-tensor checksums of the new weights catch it if not.
 """
 
 import json
 import os
 from concurrent.futures import ThreadPoolExecutor
+from contextlib import ExitStack
 from pathlib import Path
+from typing import BinaryIO
 
 import numpy as np
 import pyzstd
@@ -68,10 +70,11 @@ def write_rank_file(
     tensors: dict,
     checksums: dict[str, int],
     frame_size: int = 1 << 20,
-) -> int:
+) -> None:
     """Write a rank's delta as zstd frames of about ``frame_size`` changes each (whole entries), so
-    readers can skip frames without entries they hold. ``checksums`` (HF name -> ``bits_checksum`` of
-    the new weights this rank holds) let engines verify every step."""
+    readers can skip frames without entries they hold. Gaps are stored as uint8, with gaps of 255 or
+    more in a separate int32 stream. ``checksums`` (HF name -> ``bits_checksum`` of the new weights
+    this rank holds) let engines verify every step. Written under a temporary name, then renamed."""
     options = {pyzstd.CParameter.compressionLevel: 1, pyzstd.CParameter.checksumFlag: 1}
     frames, size = [[]], 0
     for index, (_, count) in enumerate(entries):
@@ -83,28 +86,30 @@ def write_rank_file(
     gap_ends = np.cumsum([count for _, count in entries], dtype=np.int64)
     value_ends = np.cumsum([count * bits_dtype(tensors[name][1]).itemsize for name, count in entries], dtype=np.int64)
 
-    def compress(frame: list[int]) -> tuple[bytes, bytes]:
+    def compress(frame: list[int]) -> tuple[bytes, bytes, bytes]:
         first, last = frame[0], frame[-1]
-        gap_start = gap_ends[first - 1] if first else 0
-        value_start = value_ends[first - 1] if first else 0
+        frame_gaps = gaps[(gap_ends[first - 1] if first else 0) : gap_ends[last]]
+        frame_values = values[(value_ends[first - 1] if first else 0) : value_ends[last]]
         return (
-            pyzstd.compress(gaps[gap_start : gap_ends[last]].tobytes(), options),
-            pyzstd.compress(values[value_start : value_ends[last]].tobytes(), options),
+            pyzstd.compress(np.minimum(frame_gaps, 255).astype(np.uint8).tobytes(), options),
+            pyzstd.compress(frame_gaps[frame_gaps >= 255].tobytes(), options),
+            pyzstd.compress(frame_values.tobytes(), options),
         )
 
     with ThreadPoolExecutor(8) as pool:
         compressed = list(pool.map(compress, frames if entries else []))
     arrays = {
-        "gaps": np.frombuffer(b"".join(g for g, _ in compressed), np.uint8),
-        "values": np.frombuffer(b"".join(v for _, v in compressed), np.uint8),
+        stream: np.frombuffer(b"".join(parts[index] for parts in compressed), np.uint8)
+        for index, stream in enumerate(("gaps", "escapes", "values"))
     }
     metadata = {
         "entries": json.dumps(entries),
-        "frames": json.dumps([[len(frame), len(g), len(v)] for frame, (g, v) in zip(frames, compressed)]),
+        "frames": json.dumps([[len(frame), *map(len, parts)] for frame, parts in zip(frames, compressed)]),
         "checksums": json.dumps(checksums),
     }
-    save_file(arrays, str(path), metadata=metadata)
-    return sum(array.nbytes for array in arrays.values())
+    tmp = path.with_name(f".{path.name}.tmp")
+    save_file(arrays, str(tmp), metadata=metadata)
+    tmp.rename(path)
 
 
 def decode_gaps(gaps: torch.Tensor, counts: torch.Tensor) -> torch.Tensor:
@@ -115,48 +120,39 @@ def decode_gaps(gaps: torch.Tensor, counts: torch.Tensor) -> torch.Tensor:
     return indices - torch.repeat_interleave(restart, counts)
 
 
-def _read_rank_file(
-    path: Path, tensors: dict, names: set[str] | None
-) -> tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor, dict[str, int]]:
-    with open(path, "rb") as file:
-        header_size = int.from_bytes(file.read(8), "little")
-        header = json.loads(file.read(header_size))
-        metadata = header["__metadata__"]
-        gaps_start = 8 + header_size + header["gaps"]["data_offsets"][0]
-        values_start = 8 + header_size + header["values"]["data_offsets"][0]
-        entries = [tuple(entry) for entry in json.loads(metadata["entries"])]
-        frames, start, gaps_offset, values_offset = [], 0, gaps_start, values_start
-        for num_entries, num_gap_bytes, num_value_bytes in json.loads(metadata["frames"]):
-            frame_entries = entries[start : start + num_entries]
-            if names is None or any(name in names for name, _ in frame_entries):
-                frames.append((frame_entries, (gaps_offset, num_gap_bytes), (values_offset, num_value_bytes)))
-            start += num_entries
-            gaps_offset += num_gap_bytes
-            values_offset += num_value_bytes
-
-        def read_frame(frame) -> list[tuple[str, np.ndarray, np.ndarray]]:
-            frame_entries, (gap_at, gap_size), (value_at, value_size) = frame
-            frame_gaps = np.frombuffer(pyzstd.decompress(os.pread(file.fileno(), gap_size, gap_at)), np.int32)
-            frame_values = np.frombuffer(pyzstd.decompress(os.pread(file.fileno(), value_size, value_at)), np.uint8)
-            out, gap_offset, value_offset = [], 0, 0
-            for name, count in frame_entries:
-                itemsize = bits_dtype(tensors[name][1]).itemsize
-                entry_gaps = frame_gaps[gap_offset : gap_offset + count]
-                entry_values = frame_values[value_offset : value_offset + count * itemsize]
-                gap_offset += count
-                value_offset += count * itemsize
-                if names is not None and name not in names:
-                    continue
-                if len(entry_gaps):
-                    out.append((name, entry_gaps, entry_values))
-            return out
-
-        with ThreadPoolExecutor(4) as pool:
-            kept = [entry for frame in pool.map(read_frame, frames) for entry in frame]
-        drop_page_cache(file.fileno())
-    gaps = torch.from_numpy(np.concatenate([g for _, g, _ in kept]) if kept else np.empty(0, np.int32))
-    values = torch.from_numpy(np.concatenate([v for _, _, v in kept]) if kept else np.empty(0, np.uint8))
-    return [(name, len(g)) for name, g, _ in kept], gaps, values, json.loads(metadata["checksums"])
+def _rank_file_frames(
+    file: BinaryIO, tensors: dict, names: set[str] | None
+) -> tuple[list[tuple[list[tuple[str, int]], list[tuple[int, int]]]], dict[str, int]]:
+    """The frames of an open rank file with entries in ``names``: ``(entries, [(offset, size)] of the
+    gaps, escapes and values streams)``, and the file's checksums."""
+    header_size = int.from_bytes(file.read(8), "little")
+    header = json.loads(file.read(header_size))
+    metadata = header["__metadata__"]
+    offsets = [8 + header_size + header[stream]["data_offsets"][0] for stream in ("gaps", "escapes", "values")]
+    entries = [tuple(entry) for entry in json.loads(metadata["entries"])]
+    frames, start = [], 0
+    for num_entries, *sizes in json.loads(metadata["frames"]):
+        frame_entries = entries[start : start + num_entries]
+        if names is None or any(name in names for name, _ in frame_entries):
+            frames.append((frame_entries, list(zip(offsets, sizes))))
+        start += num_entries
+        offsets = [offset + size for offset, size in zip(offsets, sizes)]
+    return frames, json.loads(metadata["checksums"])
+
+
+def _read_frame(fd: int, frame, tensors: dict, names: set[str] | None) -> list[tuple[str, np.ndarray, np.ndarray]]:
+    frame_entries, ((gaps_at, gaps_size), (escapes_at, escapes_size), (values_at, values_size)) = frame
+    gaps = np.frombuffer(pyzstd.decompress(os.pread(fd, gaps_size, gaps_at)), np.uint8).astype(np.int32)
+    gaps[gaps == 255] = np.frombuffer(pyzstd.decompress(os.pread(fd, escapes_size, escapes_at)), np.int32)
+    values = np.frombuffer(pyzstd.decompress(os.pread(fd, values_size, values_at)), np.uint8)
+    out, gap_offset, value_offset = [], 0, 0
+    for name, count in frame_entries:
+        num_bytes = count * bits_dtype(tensors[name][1]).itemsize
+        if names is None or name in names:
+            out.append((name, gaps[gap_offset : gap_offset + count], values[value_offset : value_offset + num_bytes]))
+        gap_offset += count
+        value_offset += num_bytes
+    return out
 
 
 def write_anchor_shards(state_dict: dict[str, torch.Tensor], step_dir: Path, rank: int, world_size: int) -> None:
@@ -202,7 +198,7 @@ def read_manifest(step_dir: Path) -> dict:
 
 
 def read_delta(
-    step_dir: Path, names: set[str] | None = None
+    step_dir: Path, names: set[str] | None = None, num_threads: int = 16
 ) -> tuple[list[tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]], dict[str, int]]:
     """Read and decompress every rank file of a delta step dir into ``(entries,
     gaps, xor value bytes)`` in host memory; ``decode_gaps`` turns the
@@ -210,14 +206,22 @@ def read_delta(
     only those tensors (frames without any are not read).
     Also returns the new weights' ``bits_checksum`` per HF tensor, summed over ranks."""
     manifest = read_manifest(step_dir)
-    paths = [step_dir / DELTA_DIR / name for name in manifest["files"]]
-    with ThreadPoolExecutor(len(paths)) as pool:
-        files = list(pool.map(lambda path: _read_rank_file(path, manifest["tensors"], names), paths))
+    tensors = manifest["tensors"]
     checksums: dict[str, int] = {}
-    for *_, partial in files:
-        for name, value in partial.items():
-            checksums[name] = add_checksums(checksums.get(name, 0), value)
-    return [delta for *delta, _ in files], checksums
+    with ExitStack() as stack, ThreadPoolExecutor(num_threads) as pool:
+        files = [stack.enter_context(open(step_dir / DELTA_DIR / name, "rb")) for name in manifest["files"]]
+        jobs = []
+        for file in files:
+            frames, partial = _rank_file_frames(file, tensors, names)
+            jobs.extend((file.fileno(), frame) for frame in frames)
+            for tensor_name, value in partial.items():
+                checksums[tensor_name] = add_checksums(checksums.get(tensor_name, 0), value)
+        kept = [entry for frame in pool.map(lambda job: _read_frame(*job, tensors, names), jobs) for entry in frame]
+        for file in files:
+            drop_page_cache(file.fileno())
+    gaps = torch.from_numpy(np.concatenate([g for _, g, _ in kept]) if kept else np.empty(0, np.int32))
+    values = torch.from_numpy(np.concatenate([v for _, _, v in kept]) if kept else np.empty(0, np.uint8))
+    return [([(name, len(g)) for name, g, _ in kept], gaps, values)], checksums
 
 
 def drop_page_cache(fd: int) -> None:

From 2b49427cd860c1b48f38c0f1517fba2c4ed1d240 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Thu, 15 Oct 2026 00:41:44 +0000
Subject: [PATCH 13/88] perf(weights): register exact-size pinned buffers; drop
 written sparse files from the page cache; map fp8 tensors from their bf16
 checkpoint

---
 .../inference/vllm/worker/filesystem.py       |  6 +++-
 .../inference/vllm/worker/weight_transfer.py  |  4 ++-
 src/prime_rl/transports/weights/filesystem.py | 31 ++++++++++++++-----
 src/prime_rl/utils/sparse_weights.py          | 19 ++++++++++++
 4 files changed, 51 insertions(+), 9 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index eb87c554854..09d7985fec7 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -27,8 +27,10 @@
     bits_dtype,
     decode_gaps,
     drop_page_cache,
+    pin,
     read_delta,
     read_manifest,
+    unpin,
 )
 
 # This is to get type hints for the Worker class but not actually extend it at runtime as this is required by vLLM worker extension
@@ -86,7 +88,9 @@ def _read_delta(
         read_time = time.perf_counter() - start
         num_bytes = sum(-(-gaps.nbytes // 4) * 4 + -(-values.nbytes // 4) * 4 for _, gaps, values in delta)
         if self.host_buffer is None or self.host_buffer.numel() < num_bytes:
-            self.host_buffer = torch.empty(num_bytes * 5 // 4, dtype=torch.uint8, pin_memory=True)
+            if self.host_buffer is not None:
+                unpin(self.host_buffer)
+            self.host_buffer = pin(num_bytes * 5 // 4)
         pinned, offset = [], 0
         for entries, gaps, values in delta:
             parts = []
diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py
index 07d3a37a83a..a7e7cb369e8 100644
--- a/src/prime_rl/inference/vllm/worker/weight_transfer.py
+++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py
@@ -85,7 +85,9 @@ def record_copy_(destination: torch.Tensor, source: torch.Tensor, non_blocking:
 
         def placeholders():
             for name, (shape, dtype) in tensors.items():
-                tensor = torch.empty(shape, dtype=getattr(torch, dtype), device=device)
+                # fp8 tensors of a delta are quantized from bf16 checkpoint tensors
+                checkpoint_dtype = torch.bfloat16 if dtype == "float8_e4m3fn" else getattr(torch, dtype)
+                tensor = torch.empty(shape, dtype=checkpoint_dtype, device=device)
                 source_names[tensor.untyped_storage()._cdata] = name
                 yield name, tensor
 
diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index d6449ea38c3..fe7fba9b111 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -24,9 +24,12 @@
     MANIFEST,
     anchor_complete,
     as_bits,
+    bits_dtype,
     encode_gaps,
+    pin,
     quantize_fp8_blocks,
     read_engine_layout,
+    unpin,
     write_anchor_shards,
     write_rank_file,
 )
@@ -184,6 +187,20 @@ def _fp8_diff(
         changed = xor.nonzero().squeeze(1)
         return changed, xor[changed], current_bits.view(-1)
 
+    def _allocate_snapshot(self, model: nn.Module, is_master: bool) -> None:
+        """The previous broadcast's bits live in one pinned host buffer, so trainer GPU memory is unchanged."""
+        sizes = {}
+        for key, value in model.state_dict().items():
+            shard = self._local_shard(key, value, is_master) if resolve_fqn(model, key) in self.segments else None
+            if shard is not None and shard[0].numel():
+                sizes[key] = (shard[0].numel(), self.wire_dtypes[key])
+        aligned = {key: -(-numel * dtype.itemsize // 8) * 8 for key, (numel, dtype) in sizes.items()}
+        buffer, offset = pin(sum(aligned.values())), 0
+        for key, (numel, dtype) in sizes.items():
+            bits = bits_dtype(str(dtype).removeprefix("torch."))
+            self.snapshot[key] = buffer[offset : offset + numel * dtype.itemsize].view(bits)
+            offset += aligned[key]
+
     @staticmethod
     def _diff(bits: torch.Tensor, previous: torch.Tensor, chunk: int = 1 << 28) -> tuple[torch.Tensor, torch.Tensor]:
         """Changed flat indices of ``bits`` vs the pinned host ``previous`` and their XOR values,
@@ -208,6 +225,8 @@ def encode(
         first = not self.initialized
         self.initialized = True
         self.scales_changed = False
+        if first:
+            self._allocate_snapshot(model, is_master)
         names, counts, gaps, values = [], [], [], []
         checksum_names, checksums = [], []
         for key, value in model.state_dict().items():
@@ -222,8 +241,6 @@ def encode(
                 continue
             bits, offset = as_bits(shard[0].to(self.wire_dtypes[key])), shard[1]
             if first:
-                # The previous broadcast's bits live in pinned host memory, so the trainer's GPU memory is unchanged.
-                self.snapshot[key] = torch.empty(bits.shape, dtype=bits.dtype, pin_memory=True)
                 self.snapshot[key].copy_(bits)
                 continue
             if fqn in self.fp8:
@@ -242,15 +259,15 @@ def encode(
             segment = torch.searchsorted(starts, position, right=True) - 1
             entry_starts = torch.ones_like(segment, dtype=torch.bool)
             entry_starts[1:] = segment[1:] != segment[:-1]
-            gaps.append(encode_gaps(position - starts[segment], entry_starts).to("cpu", non_blocking=True))
-            values.append(changed_values.view(torch.uint8).to("cpu", non_blocking=True))
+            gaps.append(encode_gaps(position - starts[segment], entry_starts).cpu())
+            values.append(changed_values.view(torch.uint8).cpu())
             counts.append(torch.bincount(segment, minlength=len(starts)))
             names.extend(name for _, name in self.segments[fqn])
         if self.fp8 and dist.is_initialized():
             changed_anywhere = torch.tensor(int(self.scales_changed), device=next(model.parameters()).device)
             dist.all_reduce(changed_anywhere, op=dist.ReduceOp.MAX)
             self.scales_changed = bool(changed_anywhere.item())
-        torch.cuda.synchronize()  # the snapshot and host copies are written asynchronously
+        torch.cuda.synchronize()  # the snapshot is written asynchronously
         if first or self.scales_changed:
             return None
         counts = torch.cat(counts).tolist() if counts else []
@@ -432,8 +449,8 @@ def write() -> None:
                 f"Wrote the {step_dir.name} anchor: gather {gather_time:.1f}s, write {time.perf_counter() - start - gather_time:.1f}s"
             )
             state_dict.clear()
+            unpin(buffer)
             buffer = None
-            torch._C._host_emptyCache()  # return the pinned buffer to the OS
 
         thread = threading.Thread(target=write)
         thread.start()
@@ -442,7 +459,7 @@ def write() -> None:
     def _allocate_anchor_buffer(self, model: nn.Module) -> torch.Tensor:
         if self.anchor_bytes is None:
             self.anchor_bytes = owned_weight_bytes(model)
-        return torch.empty(self.anchor_bytes, dtype=torch.uint8, pin_memory=True)
+        return pin(self.anchor_bytes)
 
     def _clean(self, step: int) -> None:
         # Sparse chains replay from their anchor, so keep everything since it.
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index 0977b71e26f..39d9edc3c29 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -110,6 +110,8 @@ def compress(frame: list[int]) -> tuple[bytes, bytes, bytes]:
     tmp = path.with_name(f".{path.name}.tmp")
     save_file(arrays, str(tmp), metadata=metadata)
     tmp.rename(path)
+    with open(path, "rb") as file:
+        drop_page_cache(file.fileno())
 
 
 def decode_gaps(gaps: torch.Tensor, counts: torch.Tensor) -> torch.Tensor:
@@ -170,6 +172,8 @@ def write_anchor_shards(state_dict: dict[str, torch.Tensor], step_dir: Path, ran
         tmp = step_dir / f".{path.name}.tmp"
         save_torch_file({name: state_dict[name].contiguous() for name in names}, tmp, metadata={"format": "pt"})
         tmp.rename(path)
+        with open(path, "rb") as file:
+            drop_page_cache(file.fileno())
     (step_dir / f"anchor_rank{rank}.done").write_text(str(world_size))
 
 
@@ -224,6 +228,21 @@ def read_delta(
     return [([(name, len(g)) for name, g, _ in kept], gaps, values)], checksums
 
 
+def pin(num_bytes: int) -> torch.Tensor:
+    """Page-locked host memory of exactly ``num_bytes``, outside torch's caching host allocator
+    (which rounds sizes up and keeps freed blocks). Release it with ``unpin``."""
+    tensor = torch.from_numpy(np.empty(num_bytes, np.uint8))
+    if num_bytes:
+        error = torch.cuda.cudart().cudaHostRegister(tensor.data_ptr(), num_bytes, 0)
+        assert error == 0, f"cudaHostRegister failed: {error}"
+    return tensor
+
+
+def unpin(tensor: torch.Tensor) -> None:
+    if tensor.numel():
+        torch.cuda.cudart().cudaHostUnregister(tensor.data_ptr())
+
+
 def drop_page_cache(fd: int) -> None:
     """Deltas and anchors are read once; don't let them fill the inference host's page cache."""
     os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)

From faab85839f0e1eec946142397dcce4b7ce235301 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Thu, 15 Oct 2026 02:04:02 +0000
Subject: [PATCH 14/88] fix(weights): decode and apply sparse deltas in chunks
 of whole tensors to bound GPU memory

---
 .../inference/vllm/worker/filesystem.py       | 50 +++++++++++++------
 1 file changed, 36 insertions(+), 14 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 09d7985fec7..e80ca11ad5c 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -148,6 +148,25 @@ def _boot_dir(self) -> Path:
         )
         return Path(folder)
 
+    def _apply_chunk(
+        self, entries: list[tuple[str, int]], gaps: torch.Tensor, values: torch.Tensor, tensors: dict, device
+    ) -> None:
+        assert self.index_map is not None
+        counts = torch.tensor([count for _, count in entries], device=device)
+        indices = decode_gaps(gaps.to(device, non_blocking=True), counts)
+        values = values.to(device, non_blocking=True)
+        offset = num_bytes = 0
+        for bits, group in itertools.groupby(entries, key=lambda entry: bits_dtype(tensors[entry[0]][1])):
+            run = list(group)
+            num_changes = sum(count for _, count in run)
+            itemsize = bits.itemsize
+            xor = values[num_bytes : num_bytes + num_changes * itemsize]
+            # A run after an odd number of narrower values is misaligned for its dtype view.
+            xor = xor.view(bits) if num_bytes % itemsize == 0 else xor.clone().view(bits)
+            self.index_map.apply_xor(run, indices[offset : offset + num_changes], xor)
+            offset += num_changes
+            num_bytes += num_changes * itemsize
+
     def _write_engine_layout(self, step_dir: Path) -> None:
         """Tell the trainer which checkpoint tensors this engine holds in fp8 (and with which
         scales), so it can send their deltas in the fp8 domain. Reads one checkpoint tensor per fp8
@@ -214,20 +233,23 @@ def _apply_delta(self, step_dir: Path) -> None:
             self.index_map = CheckpointIndexMap(model, tensors)
             logger.info(f"Built the checkpoint index map in {time.perf_counter() - start - read_time:.2f}s")
         for entries, gaps, values in delta:
-            counts = torch.tensor([count for _, count in entries], device=device)
-            indices = decode_gaps(gaps.to(device, non_blocking=True), counts)
-            values = values.to(device, non_blocking=True)
-            offset = num_bytes = 0
-            for bits, group in itertools.groupby(entries, key=lambda entry: bits_dtype(tensors[entry[0]][1])):
-                run = list(group)
-                num_changes = sum(count for _, count in run)
-                itemsize = bits.itemsize
-                xor = values[num_bytes : num_bytes + num_changes * itemsize]
-                # A run after an odd number of narrower values is misaligned for its dtype view.
-                xor = xor.view(bits) if num_bytes % itemsize == 0 else xor.clone().view(bits)
-                self.index_map.apply_xor(run, indices[offset : offset + num_changes], xor)
-                offset += num_changes
-                num_bytes += num_changes * itemsize
+            # Decode in chunks of whole entries: the int64 indices take 8 B per change on the GPU.
+            first = change_offset = byte_offset = num_changes = 0
+            for last in range(len(entries)):
+                num_changes += entries[last][1]
+                if num_changes < 1 << 26 and last + 1 < len(entries):
+                    continue
+                chunk = entries[first : last + 1]
+                num_bytes = sum(count * bits_dtype(tensors[name][1]).itemsize for name, count in chunk)
+                self._apply_chunk(
+                    chunk,
+                    gaps[change_offset : change_offset + num_changes],
+                    values[byte_offset : byte_offset + num_bytes],
+                    tensors,
+                    device,
+                )
+                first, change_offset, byte_offset = last + 1, change_offset + num_changes, byte_offset + num_bytes
+                num_changes = 0
         update_mla_absorbed_weights(model)
         torch.cuda.synchronize()
         apply_time = time.perf_counter() - start - read_time

From af4dbea61835e1922efbf6dc644c9a7613f96b8e Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Thu, 15 Oct 2026 04:01:09 +0000
Subject: [PATCH 15/88] perf(weights): background sparse writes release the
 GIL; fewer compression threads

---
 src/prime_rl/utils/sparse_weights.py | 31 +++++++++++++++-------------
 1 file changed, 17 insertions(+), 14 deletions(-)

diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index 39d9edc3c29..d32c225effa 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -19,8 +19,8 @@
 import numpy as np
 import pyzstd
 import torch
-from safetensors.numpy import save_file
-from safetensors.torch import save_file as save_torch_file
+from safetensors.numpy import save
+from safetensors.torch import save as save_torch
 
 DELTA_DIR = "delta"
 MANIFEST = "manifest.json"
@@ -96,7 +96,8 @@ def compress(frame: list[int]) -> tuple[bytes, bytes, bytes]:
             pyzstd.compress(frame_values.tobytes(), options),
         )
 
-    with ThreadPoolExecutor(8) as pool:
+    # Few threads: this runs in the background next to training (e.g. a CPU optimizer step)
+    with ThreadPoolExecutor(4) as pool:
         compressed = list(pool.map(compress, frames if entries else []))
     arrays = {
         stream: np.frombuffer(b"".join(parts[index] for parts in compressed), np.uint8)
@@ -107,11 +108,7 @@ def compress(frame: list[int]) -> tuple[bytes, bytes, bytes]:
         "frames": json.dumps([[len(frame), *map(len, parts)] for frame, parts in zip(frames, compressed)]),
         "checksums": json.dumps(checksums),
     }
-    tmp = path.with_name(f".{path.name}.tmp")
-    save_file(arrays, str(tmp), metadata=metadata)
-    tmp.rename(path)
-    with open(path, "rb") as file:
-        drop_page_cache(file.fileno())
+    _write_atomically(path, save(arrays, metadata=metadata))
 
 
 def decode_gaps(gaps: torch.Tensor, counts: torch.Tensor) -> torch.Tensor:
@@ -157,23 +154,29 @@ def _read_frame(fd: int, frame, tensors: dict, names: set[str] | None) -> list[t
     return out
 
 
+def _write_atomically(path: Path, data: bytes) -> None:
+    """Write under a temporary name, then rename. Serializing in memory and writing with plain file
+    I/O releases the GIL during the (possibly slow) write, so background writes don't stall training."""
+    tmp = path.with_name(f".{path.name}.tmp")
+    with open(tmp, "wb") as file:
+        file.write(data)
+        drop_page_cache(file.fileno())
+    tmp.rename(path)
+
+
 def write_anchor_shards(state_dict: dict[str, torch.Tensor], step_dir: Path, rank: int, world_size: int) -> None:
     """Write one trainer rank's share of a full checkpoint, without collectives; the
     ``anchor_rank{rank}.done`` marker (holding the world size) is written last."""
     shards, size = [[]], 0
     for name, tensor in state_dict.items():
-        if size > 4 * 2**30:
+        if size > 2**30:
             shards.append([])
             size = 0
         shards[-1].append(name)
         size += tensor.nbytes
     for index, names in enumerate(shards):
         path = step_dir / f"model-rank{rank:05d}-{index:05d}.safetensors"
-        tmp = step_dir / f".{path.name}.tmp"
-        save_torch_file({name: state_dict[name].contiguous() for name in names}, tmp, metadata={"format": "pt"})
-        tmp.rename(path)
-        with open(path, "rb") as file:
-            drop_page_cache(file.fileno())
+        _write_atomically(path, save_torch({name: state_dict[name] for name in names}, metadata={"format": "pt"}))
     (step_dir / f"anchor_rank{rank}.done").write_text(str(world_size))
 
 

From 7d08f096b5449125d28d9e82990d067e1e2ed90c Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Thu, 15 Oct 2026 05:05:30 +0000
Subject: [PATCH 16/88] chore(weights): log the sparse snapshot initialization
 time

---
 src/prime_rl/transports/weights/filesystem.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index fe7fba9b111..c0e59045287 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -195,7 +195,9 @@ def _allocate_snapshot(self, model: nn.Module, is_master: bool) -> None:
             if shard is not None and shard[0].numel():
                 sizes[key] = (shard[0].numel(), self.wire_dtypes[key])
         aligned = {key: -(-numel * dtype.itemsize // 8) * 8 for key, (numel, dtype) in sizes.items()}
+        start = time.perf_counter()
         buffer, offset = pin(sum(aligned.values())), 0
+        self.pin_time = time.perf_counter() - start
         for key, (numel, dtype) in sizes.items():
             bits = bits_dtype(str(dtype).removeprefix("torch."))
             self.snapshot[key] = buffer[offset : offset + numel * dtype.itemsize].view(bits)
@@ -354,6 +356,11 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
         if self.encoder.scales_changed:
             self.logger.info(f"fp8 block scales changed at v{step}: publishing a full checkpoint")
         base_step, self.last_step = self.last_step, step
+        if first:
+            self.logger.info(
+                f"Initialized the sparse weight snapshot in {encode_time:.1f}s "
+                f"(pinning {sum(s.nbytes for s in self.encoder.snapshot.values()) / 2**30:.1f} GiB: {self.encoder.pin_time:.1f}s)"
+            )
         if first and step == 0:
             # The run's first broadcast is the checkpoint the engines booted from.
             if self.world.is_master:

From 7554b1859d8414e7fb44cb2d0ddbfde9a567c428 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Thu, 15 Oct 2026 05:33:44 +0000
Subject: [PATCH 17/88] feat(weights): write sparse anchors from the trainer
 snapshot; page-lock the snapshot in the background

---
 .../inference/vllm/worker/filesystem.py       |  66 ++++++-----
 src/prime_rl/transports/weights/filesystem.py | 108 +++++++++---------
 src/prime_rl/utils/sparse_weights.py          | 100 ++++++++++++++--
 src/prime_rl/utils/weights.py                 |  34 +-----
 4 files changed, 186 insertions(+), 122 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index e80ca11ad5c..39f5ba15c86 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -24,9 +24,10 @@
     ENGINE_LAYOUT,
     MANIFEST,
     anchor_complete,
+    anchor_index,
+    anchor_tensors,
     bits_dtype,
     decode_gaps,
-    drop_page_cache,
     pin,
     read_delta,
     read_manifest,
@@ -171,24 +172,32 @@ def _write_engine_layout(self, step_dir: Path) -> None:
         """Tell the trainer which checkpoint tensors this engine holds in fp8 (and with which
         scales), so it can send their deltas in the fp8 domain. Reads one checkpoint tensor per fp8
         param from the step's full checkpoint, or from the boot checkpoint."""
-        checkpoint_dir = step_dir if anchor_complete(step_dir) else self._boot_dir()
-        files = {}
-        tensors = {}
-        for path in checkpoint_dir.glob("*.safetensors"):
-            with safe_open(path, framework="pt") as f:
-                for name in f.keys():
-                    tensor_slice = f.get_slice(name)
-                    files[name] = path
-                    tensors[name] = (tensor_slice.get_shape(), _SAFETENSORS_DTYPES[tensor_slice.get_dtype()])
+        if anchor_complete(step_dir):
+            index = anchor_index(step_dir)
+            tensors = index[1]
+
+            def read_tensor(name: str) -> torch.Tensor:
+                return next(anchor_tensors(step_dir, [name], index))[1]
+
+        else:
+            files = {}
+            tensors = {}
+            for path in self._boot_dir().glob("*.safetensors"):
+                with safe_open(path, framework="pt") as f:
+                    for name in f.keys():
+                        tensor_slice = f.get_slice(name)
+                        files[name] = path
+                        tensors[name] = (tensor_slice.get_shape(), _SAFETENSORS_DTYPES[tensor_slice.get_dtype()])
+
+            def read_tensor(name: str) -> torch.Tensor:
+                with safe_open(files[name], framework="pt") as f:
+                    return f.get_tensor(name)
+
         if (step_dir / DELTA_DIR / MANIFEST).exists():
             tensors = read_manifest(step_dir)["tensors"]
         if self.index_map is None:
             self.index_map = CheckpointIndexMap(self._model(), tensors)
 
-        def read_tensor(name: str) -> torch.Tensor:
-            with safe_open(files[name], framework="pt") as f:
-                return f.get_tensor(name)
-
         # Each worker reports the tensors it holds; the trainer waits for all of them.
         if self.vllm_config.parallel_config.enable_expert_parallel:
             group = get_ep_group()
@@ -282,25 +291,24 @@ def _verify(self, step_dir: Path, checksums: dict[str, int]) -> None:
             raise RuntimeError(f"Sparse weight delta {step_dir.name} diverged from the trainer: {mismatched[:8]}")
 
     def _load_checkpoint(self, step_dir: Path) -> None:
-        """Reload every weight from the HF-compatible checkpoint in ``step_dir``."""
+        """Reload every weight from the HF-compatible checkpoint (or sparse-run anchor) in ``step_dir``."""
         model = self._model()
-        model_loader = get_model_loader(self.load_config)
-        assert isinstance(model_loader, DefaultModelLoader)
-        local_source = DefaultModelLoader.Source(
-            step_dir.as_posix(),
-            revision=None,  # TODO: Check that this is correct or if we should use the default (model_config.revision)
-            prefix="",
-            fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True),
-            allow_patterns_overrides=getattr(model, "allow_patterns_overrides", None),
-        )
-        weights_iterator = model_loader._get_weights_iterator(local_source)
+        if anchor_complete(step_dir):
+            weights_iterator = anchor_tensors(step_dir)
+        else:
+            model_loader = get_model_loader(self.load_config)
+            assert isinstance(model_loader, DefaultModelLoader)
+            local_source = DefaultModelLoader.Source(
+                step_dir.as_posix(),
+                revision=None,  # TODO: Check that this is correct or if we should use the default (model_config.revision)
+                prefix="",
+                fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True),
+                allow_patterns_overrides=getattr(model, "allow_patterns_overrides", None),
+            )
+            weights_iterator = model_loader._get_weights_iterator(local_source)
         load_weights_checkpoint_layerwise(
             model,
             weights_iterator,
             self.model_runner.model_config,
             self.vllm_config,
         )
-        if anchor_complete(step_dir):
-            for path in step_dir.glob("*.safetensors"):
-                with open(path, "rb") as file:
-                    drop_page_cache(file.fileno())
diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index c0e59045287..1e894e6f237 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -2,7 +2,6 @@
 import threading
 import time
 from collections import defaultdict
-from concurrent.futures import Future, ThreadPoolExecutor
 from math import prod
 from pathlib import Path
 
@@ -17,6 +16,7 @@
 from prime_rl.orchestrator.clients import load_lora_adapter
 from prime_rl.trainer.lora import get_lora_state, save_lora_config
 from prime_rl.transports.weights.base import FINISHED_MARKER, WeightReceiver, WeightSender
+from prime_rl.utils.logger import get_logger
 from prime_rl.utils.pathing import wait_for_path
 from prime_rl.utils.sparse_weights import (
     DELTA_DIR,
@@ -26,17 +26,15 @@
     as_bits,
     bits_dtype,
     encode_gaps,
-    pin,
+    pin_in_background,
     quantize_fp8_blocks,
     read_engine_layout,
-    unpin,
     write_anchor_shards,
     write_rank_file,
 )
 from prime_rl.utils.weights import (
     convert_state_dict_to_hf,
     gather_weights_parallel,
-    owned_weight_bytes,
     resolve_fqn,
     resolve_wire_dtype,
     save_state_dict,
@@ -92,6 +90,8 @@ def __init__(self, model: nn.Module):
         self.gather: set[str] = set()
         self.segment_starts: dict[str, torch.Tensor] = {}
         self.snapshot: dict[str, torch.Tensor] = {}
+        self.offsets: dict[str, int] = {}
+        self.pinning: threading.Thread | None = None
         self.scales_changed = False
         self.initialized = False
 
@@ -188,21 +188,36 @@ def _fp8_diff(
         return changed, xor[changed], current_bits.view(-1)
 
     def _allocate_snapshot(self, model: nn.Module, is_master: bool) -> None:
-        """The previous broadcast's bits live in one pinned host buffer, so trainer GPU memory is unchanged."""
+        """The previous broadcast's bits live in one host buffer, so trainer GPU memory is unchanged. It
+        is page-locked on a thread while the next step trains."""
         sizes = {}
         for key, value in model.state_dict().items():
             shard = self._local_shard(key, value, is_master) if resolve_fqn(model, key) in self.segments else None
             if shard is not None and shard[0].numel():
-                sizes[key] = (shard[0].numel(), self.wire_dtypes[key])
-        aligned = {key: -(-numel * dtype.itemsize // 8) * 8 for key, (numel, dtype) in sizes.items()}
-        start = time.perf_counter()
-        buffer, offset = pin(sum(aligned.values())), 0
-        self.pin_time = time.perf_counter() - start
-        for key, (numel, dtype) in sizes.items():
+                sizes[key] = (shard[0].numel(), shard[1], self.wire_dtypes[key])
+        aligned = {key: -(-numel * dtype.itemsize // 8) * 8 for key, (numel, _, dtype) in sizes.items()}
+        buffer, self.pinning = pin_in_background(sum(aligned.values()))
+        self.pin_start = time.perf_counter()
+        offset = 0
+        for key, (numel, global_offset, dtype) in sizes.items():
             bits = bits_dtype(str(dtype).removeprefix("torch."))
             self.snapshot[key] = buffer[offset : offset + numel * dtype.itemsize].view(bits)
+            self.offsets[key] = global_offset
             offset += aligned[key]
 
+    def anchor_parts(self, model: nn.Module) -> list[tuple[str, list[int], int, torch.Tensor]]:
+        """This rank's slices of the full HF checkpoint in the snapshot (the last broadcast's wire weights)."""
+        parts = []
+        for key, bits in self.snapshot.items():
+            fqn = resolve_fqn(model, key)
+            values, offset = bits.view(self.wire_dtypes[key]), self.offsets[key]
+            for start, name in self.segments[fqn]:
+                begin, end = max(start, offset), min(start + self.hf_numel[name], offset + bits.numel())
+                if begin < end:
+                    shape = self.hf_tensors[name][0]
+                    parts.append((name, shape, begin - start, values[begin - offset : end - offset]))
+        return parts
+
     @staticmethod
     def _diff(bits: torch.Tensor, previous: torch.Tensor, chunk: int = 1 << 28) -> tuple[torch.Tensor, torch.Tensor]:
         """Changed flat indices of ``bits`` vs the pinned host ``previous`` and their XOR values,
@@ -229,6 +244,14 @@ def encode(
         self.scales_changed = False
         if first:
             self._allocate_snapshot(model, is_master)
+        elif self.pinning is not None:
+            wait = time.perf_counter()
+            self.pinning.join()
+            self.pinning = None
+            get_logger().info(
+                f"Page-locked the sparse weight snapshot in {time.perf_counter() - self.pin_start:.1f}s "
+                f"(waited {time.perf_counter() - wait:.1f}s)"
+            )
         names, counts, gaps, values = [], [], [], []
         checksum_names, checksums = [], []
         for key, value in model.state_dict().items():
@@ -303,12 +326,9 @@ def __init__(
         self.last_step: int | None = None
         self.anchor_step: int | None = None
         self.pending_anchor: int | None = None
-        # Pinned host buffer for the next anchor, allocated on a thread ahead of it
-        self.anchor_buffer: Future | None = None
-        self.anchor_bytes: int | None = None
-        self.allocator = ThreadPoolExecutor(1)
-        # At most one delta is being written while the next step trains
-        self.writer: threading.Thread | None = None
+        # Background writes of the last step (its delta, and its anchor on anchor steps)
+        self.writers: list[threading.Thread] = []
+        self.deferred_finish = False
         self.logger.debug("Initialized filesystem weight broadcast")
 
     def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
@@ -335,9 +355,10 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
         if not self.config.sparse:
             self._save_full(model, step_dir)
             return
-        if self.writer is not None:
-            self.writer.join()
-            self.writer = None
+        # A step's background writes read the snapshot, which the next encode updates.
+        for writer in self.writers:
+            writer.join()
+        self.writers, self.deferred_finish = [], False
         if self.world.is_master:
             # Engines that load a step with a delta dir report their weight layout next to it.
             (step_dir / DELTA_DIR).mkdir(exist_ok=True)
@@ -357,9 +378,9 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
             self.logger.info(f"fp8 block scales changed at v{step}: publishing a full checkpoint")
         base_step, self.last_step = self.last_step, step
         if first:
+            num_bytes = sum(snapshot.nbytes for snapshot in self.encoder.snapshot.values())
             self.logger.info(
-                f"Initialized the sparse weight snapshot in {encode_time:.1f}s "
-                f"(pinning {sum(s.nbytes for s in self.encoder.snapshot.values()) / 2**30:.1f} GiB: {self.encoder.pin_time:.1f}s)"
+                f"Initialized the {num_bytes / 2**30:.1f} GiB sparse weight snapshot in {encode_time:.1f}s"
             )
         if first and step == 0:
             # The run's first broadcast is the checkpoint the engines booted from.
@@ -373,13 +394,11 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
             dist.barrier()
             self.anchor_step = step
             return
-        self.writer = threading.Thread(target=self._write_delta, args=(step, base_step, step_dir, delta, encode_time))
-        self.writer.start()
-        # Periodic anchors only serve engines that fall off the chain: write them off the critical path.
-        # Pinning their host buffer is slow, so it starts on a thread two steps ahead.
-        if 0 < -step % self.config.anchor_interval <= 2 and self.anchor_buffer is None:
-            self.anchor_bytes = self.anchor_bytes or owned_weight_bytes(model)
-            self.anchor_buffer = self.allocator.submit(self._allocate_anchor_buffer, model)
+        writer = threading.Thread(target=self._write_delta, args=(step, base_step, step_dir, delta, encode_time))
+        writer.start()
+        self.writers.append(writer)
+        self.deferred_finish = True
+        # Periodic anchors only serve engines that fall off the chain: written in the background too.
         if (
             self.world.is_master
             and self.pending_anchor is not None
@@ -389,7 +408,7 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
         start_anchor = [step % self.config.anchor_interval == 0 and self.pending_anchor is None]
         dist.broadcast_object_list(start_anchor, src=0)
         if start_anchor[0]:
-            self._save_anchor(model, step_dir)
+            self.writers.append(self._save_anchor(model, step_dir))
             self.pending_anchor = step
 
     def _write_delta(self, step: int, base_step: int, step_dir: Path, delta: tuple, encode_time: float) -> None:
@@ -429,7 +448,7 @@ def _write_delta(self, step: int, base_step: int, step_dir: Path, delta: tuple,
 
     def _finish(self, step: int, step_dir: Path) -> None:
         # A delta step is committed by its writer thread once every rank's file is on disk.
-        if self.writer is None:
+        if not self.deferred_finish:
             super()._finish(step, step_dir)
 
     def _save_full(self, model: nn.Module, step_dir: Path) -> None:
@@ -439,35 +458,20 @@ def _save_full(self, model: nn.Module, step_dir: Path) -> None:
         save_state_dict_parallel(state_dict, step_dir)
 
     def _save_anchor(self, model: nn.Module, step_dir: Path) -> threading.Thread:
-        """Gather this rank's share of the full checkpoint into a pinned host buffer, then write it
-        on a thread (no collectives) and release the buffer; engines treat the step as an anchor
-        once every rank's marker exists."""
+        """Write this rank's slices of the full checkpoint from the snapshot on a thread; engines treat
+        the step as an anchor once every rank's marker exists."""
+        assert self.encoder is not None
+        parts = self.encoder.anchor_parts(model)
         start = time.perf_counter()
-        buffer = self.anchor_buffer.result() if self.anchor_buffer is not None else self._allocate_anchor_buffer(model)
-        self.anchor_buffer = None
-        state_dict = convert_state_dict_to_hf(model, gather_weights_parallel(model, host_buffer=buffer))
-        torch.cuda.empty_cache()
-        gather_time = time.perf_counter() - start
 
         def write() -> None:
-            nonlocal buffer
-            write_anchor_shards(state_dict, step_dir, self.world.rank, self.world.world_size)
-            self.logger.debug(
-                f"Wrote the {step_dir.name} anchor: gather {gather_time:.1f}s, write {time.perf_counter() - start - gather_time:.1f}s"
-            )
-            state_dict.clear()
-            unpin(buffer)
-            buffer = None
+            write_anchor_shards(parts, step_dir, self.world.rank, self.world.world_size)
+            self.logger.debug(f"Wrote the {step_dir.name} anchor in {time.perf_counter() - start:.1f}s")
 
         thread = threading.Thread(target=write)
         thread.start()
         return thread
 
-    def _allocate_anchor_buffer(self, model: nn.Module) -> torch.Tensor:
-        if self.anchor_bytes is None:
-            self.anchor_bytes = owned_weight_bytes(model)
-        return pin(self.anchor_bytes)
-
     def _clean(self, step: int) -> None:
         # Sparse chains replay from their anchor, so keep everything since it.
         super()._clean(step if self.anchor_step is None else min(step, self.anchor_step + 1))
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index d32c225effa..0a3b29631bb 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -11,16 +11,20 @@
 
 import json
 import os
+import threading
+from collections import defaultdict
+from collections.abc import Iterable, Iterator
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import ExitStack
+from math import prod
 from pathlib import Path
 from typing import BinaryIO
 
 import numpy as np
 import pyzstd
 import torch
+from safetensors import safe_open
 from safetensors.numpy import save
-from safetensors.torch import save as save_torch
 
 DELTA_DIR = "delta"
 MANIFEST = "manifest.json"
@@ -164,22 +168,79 @@ def _write_atomically(path: Path, data: bytes) -> None:
     tmp.rename(path)
 
 
-def write_anchor_shards(state_dict: dict[str, torch.Tensor], step_dir: Path, rank: int, world_size: int) -> None:
-    """Write one trainer rank's share of a full checkpoint, without collectives; the
-    ``anchor_rank{rank}.done`` marker (holding the world size) is written last."""
-    shards, size = [[]], 0
-    for name, tensor in state_dict.items():
+def write_anchor_shards(
+    parts: list[tuple[str, list[int], int, torch.Tensor]], step_dir: Path, rank: int, world_size: int
+) -> None:
+    """Write one trainer rank's share of a full checkpoint: ``(hf_name, shape, flat start, flat values)``
+    slices, stored as tensors ``"{hf_name}@{start}"`` in ~1 GiB safetensors files, straight from host
+    memory and without collectives. The ``anchor_rank{rank}.done`` marker (the world size) comes last."""
+    files, size = [[]], 0
+    for part in parts:
         if size > 2**30:
-            shards.append([])
+            files.append([])
             size = 0
-        shards[-1].append(name)
-        size += tensor.nbytes
-    for index, names in enumerate(shards):
+        files[-1].append(part)
+        size += part[3].nbytes
+    for index, file_parts in enumerate(files):
+        header, offset = {"__metadata__": {"shapes": json.dumps({name: shape for name, shape, _, _ in file_parts})}}, 0
+        for name, _, start, values in file_parts:
+            dtype = _SAFETENSORS_DTYPES[values.dtype]
+            header[f"{name}@{start}"] = {
+                "dtype": dtype,
+                "shape": [values.numel()],
+                "data_offsets": [offset, offset + values.nbytes],
+            }
+            offset += values.nbytes
+        encoded = json.dumps(header).encode()
+        encoded += b" " * (-len(encoded) % 8)
         path = step_dir / f"model-rank{rank:05d}-{index:05d}.safetensors"
-        _write_atomically(path, save_torch({name: state_dict[name] for name in names}, metadata={"format": "pt"}))
+        tmp = path.with_name(f".{path.name}.tmp")
+        with open(tmp, "wb") as file:
+            file.write(len(encoded).to_bytes(8, "little") + encoded)
+            for *_, values in file_parts:
+                file.write(memoryview(values.view(torch.uint8).numpy()))
+            drop_page_cache(file.fileno())
+        tmp.rename(path)
     (step_dir / f"anchor_rank{rank}.done").write_text(str(world_size))
 
 
+_SAFETENSORS_DTYPES = {torch.bfloat16: "BF16", torch.float16: "F16", torch.float32: "F32"}
+
+
+def anchor_index(step_dir: Path) -> tuple[dict[str, list[tuple[Path, str, int]]], dict[str, tuple[list[int], str]]]:
+    """Where each tensor's slices are in an anchor written by ``write_anchor_shards``, and its shape and dtype."""
+    slices: dict[str, list[tuple[Path, str, int]]] = defaultdict(list)
+    tensors: dict[str, tuple[list[int], str]] = {}
+    for path in sorted(step_dir.glob("model-rank*.safetensors")):
+        with safe_open(path, framework="pt") as file:
+            shapes = json.loads(file.metadata()["shapes"])
+            for key in file.keys():
+                name, start = key.rsplit("@", 1)
+                slices[name].append((path, key, int(start)))
+                dtype = str(file.get_slice(key).get_dtype())
+                tensors[name] = (shapes[name], {"BF16": "bfloat16", "F16": "float16", "F32": "float32"}[dtype])
+    return dict(slices), tensors
+
+
+def anchor_tensors(
+    step_dir: Path, names: Iterable[str] | None = None, index=None
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """Assemble the full tensors of an anchor (all, or ``names``), in the order the trainer wrote them."""
+    slices, tensors = index or anchor_index(step_dir)
+    for name in slices if names is None else names:
+        shape, dtype = tensors[name]
+        tensor = torch.empty(prod(shape), dtype=getattr(torch, dtype))
+        for path, key, start in slices[name]:
+            with safe_open(path, framework="pt") as file:
+                values = file.get_tensor(key)
+            tensor[start : start + values.numel()] = values
+        yield name, tensor.view(shape)
+    if names is None:
+        for path in step_dir.glob("model-rank*.safetensors"):
+            with open(path, "rb") as file:
+                drop_page_cache(file.fileno())
+
+
 def anchor_complete(step_dir: Path) -> bool:
     """Whether ``step_dir`` holds a complete full checkpoint written by ``write_anchor_shards``."""
     markers = list(step_dir.glob("anchor_rank*.done"))
@@ -241,6 +302,23 @@ def pin(num_bytes: int) -> torch.Tensor:
     return tensor
 
 
+def pin_in_background(num_bytes: int, chunk: int = 1 << 30) -> tuple[torch.Tensor, threading.Thread]:
+    """Host memory of ``num_bytes`` that a thread page-locks chunk by chunk (``cudaHostRegister``);
+    it is usable as pageable memory right away. Page-locking a large buffer is slow on a loaded host."""
+    raw = np.empty(num_bytes + 4096, np.uint8)
+    offset = -raw.ctypes.data % 4096
+    tensor = torch.from_numpy(raw[offset : offset + num_bytes])
+
+    def register() -> None:
+        for start in range(0, num_bytes, chunk):
+            error = torch.cuda.cudart().cudaHostRegister(tensor.data_ptr() + start, min(chunk, num_bytes - start), 0)
+            assert error == 0, f"cudaHostRegister failed: {error}"
+
+    thread = threading.Thread(target=register)
+    thread.start()
+    return tensor, thread
+
+
 def unpin(tensor: torch.Tensor) -> None:
     if tensor.numel():
         torch.cuda.cudart().cudaHostUnregister(tensor.data_ptr())
diff --git a/src/prime_rl/utils/weights.py b/src/prime_rl/utils/weights.py
index e50e8ee637e..ac43509c4b1 100644
--- a/src/prime_rl/utils/weights.py
+++ b/src/prime_rl/utils/weights.py
@@ -159,17 +159,14 @@ def resolve_wire_dtype(keep_in_fp32: Callable[[str], bool] | None, key: str, def
     return torch.float32 if keep_in_fp32 is not None and keep_in_fp32(key) else default
 
 
-def gather_weights_parallel(
-    model: nn.Module, dtype: torch.dtype = torch.bfloat16, host_buffer: Tensor | None = None
-) -> dict[str, Tensor]:
+def gather_weights_parallel(model: nn.Module, dtype: torch.dtype = torch.bfloat16) -> dict[str, Tensor]:
     """Gather distributed weights cooperatively, each rank keeping a slice on CPU.
 
     Every rank participates in the per-tensor all-gathers (a ``full_tensor`` call is
     collective and materializes the full tensor on all ranks anyway), but instead of
     only the master keeping everything, each rank copies its owned slice to CPU —
-    so the D2H traffic and the shard writes are split across ranks. With
-    ``host_buffer`` (pinned uint8 of at least ``owned_weight_bytes``), the slice is
-    copied into it instead of fresh pageable memory, which is several times faster.
+    so the D2H traffic and the shard writes are split across ranks. The copies are
+    plain blocking transfers; no stream-ordering assumptions.
 
     Parameters the model reports via ``keep_in_fp32_for_weight_transfer`` skip the
     downcast. Those lists exist because the inference engine's kernels require fp32
@@ -182,7 +179,6 @@ def gather_weights_parallel(
     world = get_world()
     owners = partition_weights(model.state_dict(), world.world_size, dtype)
     partial: dict[str, Tensor] = {}
-    offset = 0
     with warnings.catch_warnings():
         warnings.filterwarnings("ignore", category=FutureWarning, module="torch.distributed")
         warnings.filterwarnings("ignore", category=UserWarning, module="torch.distributed.*")
@@ -194,15 +190,7 @@ def gather_weights_parallel(
                 value = cast(DTensor, value.to(target_dtype)).full_tensor()
             if owners[key] != world.rank:
                 continue
-            if host_buffer is None:
-                partial[resolve_fqn(model, key)] = value.to("cpu")
-                continue
-            out = host_buffer[offset : offset + value.nbytes].view(value.dtype).view(value.shape)
-            out.copy_(value, non_blocking=True)
-            partial[resolve_fqn(model, key)] = out
-            offset += -(-value.nbytes // 8) * 8
-        if host_buffer is not None:
-            torch.cuda.synchronize()
+            partial[resolve_fqn(model, key)] = value.to("cpu")
         dist.barrier()
 
     if any(".base_layer." in key or "lora_A" in key or "lora_B" in key for key in partial.keys()):
@@ -211,20 +199,6 @@ def gather_weights_parallel(
     return partial
 
 
-def owned_weight_bytes(model: nn.Module, dtype: torch.dtype = torch.bfloat16) -> int:
-    """Host bytes of this rank's share of ``gather_weights_parallel`` (8-byte aligned per tensor)."""
-    keep_in_fp32 = getattr(model, "keep_in_fp32_for_weight_transfer", None)
-    world = get_world()
-    owners = partition_weights(model.state_dict(), world.world_size, dtype)
-    total = 0
-    for key, value in model.state_dict().items():
-        if owners[key] == world.rank:
-            is_dtensor = isinstance(value, DTensor)
-            itemsize = resolve_wire_dtype(keep_in_fp32, key, dtype).itemsize if is_dtensor else value.element_size()
-            total += -(-value.numel() * itemsize // 8) * 8
-    return total
-
-
 def save_state_dict_parallel(state_dict: dict[str, Tensor], save_dir: Path) -> None:
     """Cooperatively save a rank-partitioned state dict as sharded safetensors.
 

From 8a488797637d3fe2ae872abaa604cea4a6bae846 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Thu, 15 Oct 2026 06:03:09 +0000
Subject: [PATCH 18/88] fix(weights): page-lock the snapshot only after it is
 filled

---
 src/prime_rl/transports/weights/filesystem.py | 4 +++-
 src/prime_rl/utils/sparse_weights.py          | 9 ++++-----
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index 1e894e6f237..61a4129f05a 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -197,7 +197,6 @@ def _allocate_snapshot(self, model: nn.Module, is_master: bool) -> None:
                 sizes[key] = (shard[0].numel(), shard[1], self.wire_dtypes[key])
         aligned = {key: -(-numel * dtype.itemsize // 8) * 8 for key, (numel, _, dtype) in sizes.items()}
         buffer, self.pinning = pin_in_background(sum(aligned.values()))
-        self.pin_start = time.perf_counter()
         offset = 0
         for key, (numel, global_offset, dtype) in sizes.items():
             bits = bits_dtype(str(dtype).removeprefix("torch."))
@@ -288,6 +287,9 @@ def encode(
             values.append(changed_values.view(torch.uint8).cpu())
             counts.append(torch.bincount(segment, minlength=len(starts)))
             names.extend(name for _, name in self.segments[fqn])
+        if first:
+            self.pin_start = time.perf_counter()
+            self.pinning.start()
         if self.fp8 and dist.is_initialized():
             changed_anywhere = torch.tensor(int(self.scales_changed), device=next(model.parameters()).device)
             dist.all_reduce(changed_anywhere, op=dist.ReduceOp.MAX)
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index 0a3b29631bb..52c9275c74c 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -303,8 +303,9 @@ def pin(num_bytes: int) -> torch.Tensor:
 
 
 def pin_in_background(num_bytes: int, chunk: int = 1 << 30) -> tuple[torch.Tensor, threading.Thread]:
-    """Host memory of ``num_bytes`` that a thread page-locks chunk by chunk (``cudaHostRegister``);
-    it is usable as pageable memory right away. Page-locking a large buffer is slow on a loaded host."""
+    """Host memory of ``num_bytes`` and a (not yet started) thread that page-locks it chunk by chunk
+    (``cudaHostRegister``). Page-locking a large buffer is slow on a loaded host. No CUDA copies may
+    touch the memory while the thread runs."""
     raw = np.empty(num_bytes + 4096, np.uint8)
     offset = -raw.ctypes.data % 4096
     tensor = torch.from_numpy(raw[offset : offset + num_bytes])
@@ -314,9 +315,7 @@ def register() -> None:
             error = torch.cuda.cudart().cudaHostRegister(tensor.data_ptr() + start, min(chunk, num_bytes - start), 0)
             assert error == 0, f"cudaHostRegister failed: {error}"
 
-    thread = threading.Thread(target=register)
-    thread.start()
-    return tensor, thread
+    return tensor, threading.Thread(target=register)
 
 
 def unpin(tensor: torch.Tensor) -> None:

From 939c14e75898fd6cb75858bc5cef010ff38f9936 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Thu, 15 Oct 2026 06:43:03 +0000
Subject: [PATCH 19/88] fix(weights): page-lock the snapshot with one
 registration

---
 src/prime_rl/utils/sparse_weights.py | 16 +++++++---------
 1 file changed, 7 insertions(+), 9 deletions(-)

diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index 52c9275c74c..65066132016 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -302,17 +302,15 @@ def pin(num_bytes: int) -> torch.Tensor:
     return tensor
 
 
-def pin_in_background(num_bytes: int, chunk: int = 1 << 30) -> tuple[torch.Tensor, threading.Thread]:
-    """Host memory of ``num_bytes`` and a (not yet started) thread that page-locks it chunk by chunk
-    (``cudaHostRegister``). Page-locking a large buffer is slow on a loaded host. No CUDA copies may
-    touch the memory while the thread runs."""
-    raw = np.empty(num_bytes + 4096, np.uint8)
-    offset = -raw.ctypes.data % 4096
-    tensor = torch.from_numpy(raw[offset : offset + num_bytes])
+def pin_in_background(num_bytes: int) -> tuple[torch.Tensor, threading.Thread]:
+    """Host memory of ``num_bytes`` and a (not yet started) thread that page-locks it with one
+    ``cudaHostRegister`` (a copy may not span two registrations). Page-locking a large buffer is slow
+    on a loaded host. No CUDA copies may touch the memory while the thread runs."""
+    tensor = torch.from_numpy(np.empty(num_bytes, np.uint8))
 
     def register() -> None:
-        for start in range(0, num_bytes, chunk):
-            error = torch.cuda.cudart().cudaHostRegister(tensor.data_ptr() + start, min(chunk, num_bytes - start), 0)
+        if num_bytes:
+            error = torch.cuda.cudart().cudaHostRegister(tensor.data_ptr(), num_bytes, 0)
             assert error == 0, f"cudaHostRegister failed: {error}"
 
     return tensor, threading.Thread(target=register)

From e07626299ddd1caec8d36d29831dddecf0e1d5f5 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Thu, 15 Oct 2026 07:23:06 +0000
Subject: [PATCH 20/88] fix(weights): load anchors layer by layer

---
 src/prime_rl/utils/sparse_weights.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index 65066132016..4d7775025cc 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -11,6 +11,7 @@
 
 import json
 import os
+import re
 import threading
 from collections import defaultdict
 from collections.abc import Iterable, Iterator
@@ -222,12 +223,17 @@ def anchor_index(step_dir: Path) -> tuple[dict[str, list[tuple[Path, str, int]]]
     return dict(slices), tensors
 
 
+def _numeric_order(name: str) -> list:
+    return [int(part) if part.isdigit() else part for part in re.split(r"(\d+)", name)]
+
+
 def anchor_tensors(
     step_dir: Path, names: Iterable[str] | None = None, index=None
 ) -> Iterator[tuple[str, torch.Tensor]]:
-    """Assemble the full tensors of an anchor (all, or ``names``), in the order the trainer wrote them."""
+    """Assemble the full tensors of an anchor (all, or ``names``)."""
     slices, tensors = index or anchor_index(step_dir)
-    for name in slices if names is None else names:
+    # Layer by layer (numeric order), so the engine's layerwise reload holds few layers at a time
+    for name in sorted(slices, key=_numeric_order) if names is None else names:
         shape, dtype = tensors[name]
         tensor = torch.empty(prod(shape), dtype=getattr(torch, dtype))
         for path, key, start in slices[name]:

From 3c6f97ac7ad34c73d50a2e9f3051602bb8d64767 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Thu, 15 Oct 2026 12:06:50 +0000
Subject: [PATCH 21/88] feat(weights): encode sparse deltas in the background,
 overlapped with the next step; GPU snapshot when it fits

---
 .../src/prime_rl/configs/shared.py            |   5 +
 src/prime_rl/trainer/rl/train.py              |   4 +
 src/prime_rl/trainer/sft/train.py             |   3 +
 src/prime_rl/transports/weights/base.py       |   7 +
 src/prime_rl/transports/weights/filesystem.py | 575 ++++++++++++------
 src/prime_rl/utils/sparse_weights.py          |  15 -
 6 files changed, 400 insertions(+), 209 deletions(-)

diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
index bd23b47b955..7c425a012a0 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
@@ -57,6 +57,11 @@ class FileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig):
     """With ``sparse``, publish a full checkpoint every this many steps (and at the first broadcast).
     Engines that fall off the delta chain reload the last anchor and replay the deltas since."""
 
+    snapshot: Literal["auto", "gpu", "cpu"] = "auto"
+    """With ``sparse``, where each trainer rank keeps the previous broadcast's weights to diff against.
+    ``auto`` keeps them on the GPU if they fit next to the first training step's peak memory, else in
+    page-locked host memory."""
+
 
 class InMemoryWeightBroadcastConfig(BaseWeightBroadcastConfig):
     host: str = "localhost"
diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py
index 5f132da55dc..4d4c7891835 100644
--- a/src/prime_rl/trainer/rl/train.py
+++ b/src/prime_rl/trainer/rl/train.py
@@ -574,6 +574,8 @@ def train(config: TrainerConfig):
             grad_norm = clip_grad_norm_(gradient_manager, model, config.optim.max_norm)
 
         # Update the model parameters
+        if weight_sender is not None:
+            weight_sender.wait_for_weight_reads()
         optimizer.step()
         optimizer.zero_grad()
 
@@ -732,6 +734,8 @@ def train(config: TrainerConfig):
     if gradient_manager is not None:
         gradient_manager.close()
     micro_batch_workers.close()
+    if weight_sender is not None:
+        weight_sender.close()
 
     logger.info(f"Peak memory: {max_peak_memory:.1f} GiB")
     logger.success(f"RL trainer {'paused' if paused else 'finished'} at step {progress.step}")
diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py
index 4008e381854..bdc90d2fae9 100644
--- a/src/prime_rl/trainer/sft/train.py
+++ b/src/prime_rl/trainer/sft/train.py
@@ -519,6 +519,8 @@ def is_online_eval_step(step: int) -> bool:
             logger.debug(f"Clipping gradients with max norm {config.optim.max_norm}")
             grad_norm = clip_grad_norm_(gradient_manager, model, config.optim.max_norm)
         logger.debug("Optimizer step")
+        if weight_sender is not None:
+            weight_sender.wait_for_weight_reads()
         optimizer.step()
         optimizer.zero_grad()
 
@@ -683,6 +685,7 @@ def is_online_eval_step(step: int) -> bool:
     if weight_sender is not None and not paused:
         logger.info("Broadcasting final weights")
         weight_sender.broadcast(model, step=progress.step)
+        weight_sender.close()
 
     if gradient_manager is not None:
         gradient_manager.close()
diff --git a/src/prime_rl/transports/weights/base.py b/src/prime_rl/transports/weights/base.py
index bb51f23f837..13182ea0d63 100644
--- a/src/prime_rl/transports/weights/base.py
+++ b/src/prime_rl/transports/weights/base.py
@@ -97,6 +97,13 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
         """Move v{step}'s weights to the consumer. All ranks enter only after
         the master finished the handshake."""
 
+    def wait_for_weight_reads(self) -> None:
+        """Block until background work of the last broadcast no longer reads the trainer weights. Called
+        before every optimizer step."""
+
+    def close(self) -> None:
+        """Finish the background work of the last broadcast."""
+
     def _finish(self, step: int, step_dir: Path) -> None:
         """Mark v{step} complete for the consumer (master only)."""
         (step_dir / FINISHED_MARKER).touch()
diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index 61a4129f05a..d2b80124d0f 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -2,9 +2,11 @@
 import threading
 import time
 from collections import defaultdict
+from concurrent.futures import Future, ThreadPoolExecutor
 from math import prod
 from pathlib import Path
 
+import numpy as np
 import torch
 import torch.distributed as dist
 import torch.nn as nn
@@ -26,7 +28,6 @@
     as_bits,
     bits_dtype,
     encode_gaps,
-    pin_in_background,
     quantize_fp8_blocks,
     read_engine_layout,
     write_anchor_shards,
@@ -46,21 +47,30 @@ class SparseDeltaEncoder:
     """Rank-local XOR deltas of the trainer's weight shards, in HF coordinates.
 
     Every rank keeps the wire-dtype bits of its own shards from the previous
-    broadcast in pinned host memory and diffs against them on device, so no
-    weights are gathered. The prime->HF conversion is replayed once on meta
-    tensors: every HF tensor must come out as a contiguous view of one trainer
-    tensor (renames, expert unstacking), which maps a changed trainer element to
-    its HF name and index.
+    broadcast (the snapshot) and diffs against them, so no weights are gathered.
+    The prime->HF conversion is replayed once on meta tensors: every HF tensor
+    must come out as a contiguous view of one trainer tensor (renames, expert
+    unstacking), which maps a changed trainer element to its HF name and index.
+
+    ``encode`` runs on a background thread and a side CUDA stream while the
+    next step trains, without collectives; it only reads the weights before
+    ``weights_read`` is set. The snapshot lives on the GPU when it fits next to
+    the first training step's peak memory, else in one host buffer. With a GPU
+    snapshot, the same host buffer holds anchor copies instead, so a rank has at
+    most one shard-sized host copy.
 
     Tensors the engines hold in fp8 (``set_fp8_layout``: HF name -> scale format,
     from the engine layout) are quantized with vLLM's own per-block quantizer,
     previous and current weights alike, and diffed as fp8 bytes. Their 128-row
-    blocks must not straddle trainer ranks, else the trainer tensor is gathered.
-    A step that changes any block scale returns no delta, so it is published as
-    a full checkpoint.
+    blocks must not straddle trainer ranks, else the trainer tensor is gathered
+    (by the caller, before ``encode``). A step that changes any block scale is
+    published as a full checkpoint.
     """
 
-    def __init__(self, model: nn.Module):
+    # Elements per transient chunk of the diff (bf16: 128 MiB per temporary)
+    CHUNK = 1 << 26
+
+    def __init__(self, model: nn.Module, snapshot: str):
         keep_in_fp32 = getattr(model, "keep_in_fp32_for_weight_transfer", None)
         self.wire_dtypes: dict[str, torch.dtype] = {}
         meta: dict[str, torch.Tensor] = {}
@@ -89,11 +99,19 @@ def __init__(self, model: nn.Module):
         self.fp8_layout: dict[str, str] | None = None
         self.gather: set[str] = set()
         self.segment_starts: dict[str, torch.Tensor] = {}
+        self.snapshot_placement = snapshot
+        # key -> (local numel, global flat offset) of this rank's shard
+        self.shards: dict[str, tuple[int, int]] = {}
+        # key -> previous wire bits: GPU tensors, or views into the host buffer
         self.snapshot: dict[str, torch.Tensor] = {}
-        self.offsets: dict[str, int] = {}
+        self.on_gpu: bool | None = None  # None until placed (auto: at the first delta)
+        self.host_buffer: torch.Tensor | None = None
+        self.host: dict[str, torch.Tensor] = {}
         self.pinning: threading.Thread | None = None
-        self.scales_changed = False
+        self.pinned = False
         self.initialized = False
+        self.device = torch.cuda.current_device()
+        self.stream = torch.cuda.Stream()
 
     def set_fp8_layout(self, model: nn.Module, fp8: dict[str, str], is_master: bool) -> None:
         """Diff the tensors the engines hold in fp8 (HF name -> scale format) as fp8 bytes from now on."""
@@ -155,129 +173,221 @@ def _find_gathered(self, model: nn.Module, is_master: bool) -> set[str]:
             dist.all_reduce(flags, op=dist.ReduceOp.MAX)
         return {key for key, flag in zip(keys, flags.tolist()) if flag}
 
-    def _fp8_diff(
-        self, key: str, fqn: str, value: torch.Tensor, bits: torch.Tensor, offset: int
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Changed flat indices, XOR values and new fp8 bits of this rank's shard, with previous and
-        current weights quantized as the engines do; advances the snapshot."""
-        ue8m0, num_cols = self.fp8[fqn]
-        previous = self.snapshot[key].to(bits.device, non_blocking=True).view(torch.bfloat16)
-        current = bits.view(torch.bfloat16)
-        if key in self.gather:
-            local_shape = (previous.numel() // prod(value.shape[1:]), *value.shape[1:])
-            previous = DTensor.from_local(
-                previous.view(local_shape),
-                value.device_mesh,
-                value.placements,
-                shape=value.shape,
-                stride=value.stride(),
-            ).full_tensor()
-            current = value.full_tensor().to(torch.bfloat16)
-        (previous_bits, previous_scales), (current_bits, current_scales) = (
-            quantize_fp8_blocks(weights.reshape(-1, num_cols), ue8m0) for weights in (previous, current)
-        )
-        if not torch.equal(previous_scales, current_scales):
-            self.scales_changed = True
-        if key in self.gather:
-            start = offset
-            previous_bits = previous_bits.view(-1)[start : start + bits.numel()]
-            current_bits = current_bits.view(-1)[start : start + bits.numel()]
-        xor = previous_bits.view(-1) ^ current_bits.view(-1)
-        self.snapshot[key].copy_(bits, non_blocking=True)
-        changed = xor.nonzero().squeeze(1)
-        return changed, xor[changed], current_bits.view(-1)
-
-    def _allocate_snapshot(self, model: nn.Module, is_master: bool) -> None:
-        """The previous broadcast's bits live in one host buffer, so trainer GPU memory is unchanged. It
-        is page-locked on a thread while the next step trains."""
-        sizes = {}
+    def local_shards(self, model: nn.Module, is_master: bool) -> list[tuple]:
+        """This rank's ``(key, fqn, flat shard, global flat offset, gathered)`` for ``encode``, taken on the
+        training thread: FSDP swaps module parameters during forward, but the sharded storage stays put.
+        ``gathered`` is ``(previous, current)`` bf16 full tensors of fp8 tensors whose blocks straddle ranks
+        (the all-gathers run here, as ``encode`` runs no collectives)."""
+        shards = []
         for key, value in model.state_dict().items():
-            shard = self._local_shard(key, value, is_master) if resolve_fqn(model, key) in self.segments else None
-            if shard is not None and shard[0].numel():
-                sizes[key] = (shard[0].numel(), shard[1], self.wire_dtypes[key])
-        aligned = {key: -(-numel * dtype.itemsize // 8) * 8 for key, (numel, _, dtype) in sizes.items()}
-        buffer, self.pinning = pin_in_background(sum(aligned.values()))
+            fqn = resolve_fqn(model, key)
+            if fqn not in self.segments:
+                continue
+            shard = self._local_shard(key, value, is_master)
+            empty = shard is None or shard[0].numel() == 0
+            gathered = None
+            if key in self.gather and self.initialized:
+                if empty:
+                    value.full_tensor()  # every rank joins the all-gathers
+                    value.full_tensor()
+                else:
+                    gathered = (self._gather_previous(key, value), value.full_tensor().to(torch.bfloat16))
+            if not empty:
+                shards.append((key, fqn, shard[0], shard[1], gathered))
+        return shards
+
+    def _gather_previous(self, key: str, value: DTensor) -> torch.Tensor:
+        previous = self.snapshot[key].to(f"cuda:{self.device}").view(torch.bfloat16)
+        local_shape = (previous.numel() // prod(value.shape[1:]), *value.shape[1:])
+        return DTensor.from_local(
+            previous.view(local_shape), value.device_mesh, value.placements, shape=value.shape, stride=value.stride()
+        ).full_tensor()
+
+    def _bits(self, key: str) -> torch.dtype:
+        return bits_dtype(str(self.wire_dtypes[key]).removeprefix("torch."))
+
+    def _allocate_host(self) -> None:
+        """One host buffer for all shards (the CPU snapshot, or anchor copies of a GPU snapshot)."""
+        sizes = {key: -(-numel * self.wire_dtypes[key].itemsize // 8) * 8 for key, (numel, _) in self.shards.items()}
+        self.host_buffer = torch.from_numpy(np.empty(sum(sizes.values()), np.uint8))
         offset = 0
-        for key, (numel, global_offset, dtype) in sizes.items():
-            bits = bits_dtype(str(dtype).removeprefix("torch."))
-            self.snapshot[key] = buffer[offset : offset + numel * dtype.itemsize].view(bits)
-            self.offsets[key] = global_offset
-            offset += aligned[key]
-
-    def anchor_parts(self, model: nn.Module) -> list[tuple[str, list[int], int, torch.Tensor]]:
-        """This rank's slices of the full HF checkpoint in the snapshot (the last broadcast's wire weights)."""
+        for key, (numel, _) in self.shards.items():
+            self.host[key] = self.host_buffer[offset : offset + numel * self.wire_dtypes[key].itemsize].view(
+                self._bits(key)
+            )
+            offset += sizes[key]
+
+    def _host_ready(self) -> None:
+        """Wait for the page-locking thread: no CUDA copy may touch the buffer while it runs."""
+        if self.pinning is not None:
+            start = time.perf_counter()
+            self.pinning.join()
+            self.pinning, self.pinned = None, True
+            get_logger().info(
+                f"Page-locked the {self.host_buffer.numel() / 2**30:.1f} GiB sparse weight host buffer in "
+                f"{time.perf_counter() - self.pin_start:.1f}s (waited {time.perf_counter() - start:.1f}s)"
+            )
+
+    def _pin_host(self) -> None:
+        if self.host_buffer is None or self.pinned or self.pinning is not None:
+            return
+        buffer = self.host_buffer
+
+        def register() -> None:
+            torch.cuda.set_device(self.device)
+            error = torch.cuda.cudart().cudaHostRegister(buffer.data_ptr(), buffer.numel(), 0)
+            assert error == 0, f"cudaHostRegister failed: {error}"
+
+        self.pin_start = time.perf_counter()
+        self.pinning = threading.Thread(target=register)
+        self.pinning.start()
+
+    def _place(self, peak_reserved: int) -> None:
+        """Keep the snapshot on the GPU if it fits next to the first training step's peak, else on the host."""
+        num_bytes = sum(numel * self.wire_dtypes[key].itemsize for key, (numel, _) in self.shards.items())
+        total = torch.cuda.get_device_properties(self.device).total_memory
+        headroom = total - peak_reserved - num_bytes
+        if self.snapshot_placement == "auto":
+            self.on_gpu = headroom >= max(4 * 2**30, total // 20)
+        else:
+            self.on_gpu = self.snapshot_placement == "gpu"
+        get_logger().info(
+            f"Sparse weight snapshot ({num_bytes / 2**30:.1f} GiB) on the {'GPU' if self.on_gpu else 'host'} "
+            f"(training peak {peak_reserved / 2**30:.1f} GiB of {total / 2**30:.1f} GiB, "
+            f"snapshot={self.snapshot_placement})"
+        )
+
+    def _fill(self, shards: list[tuple]) -> None:
+        """The first broadcast only records the weights: on the GPU if forced there, else in the host buffer
+        (placed at the first delta, against the measured training peak)."""
+        self.shards = {key: (shard.numel(), offset) for key, _, shard, offset, _ in shards}
+        if self.snapshot_placement == "gpu":
+            self.on_gpu = True
+            for key, _, shard, _, _ in shards:
+                self.snapshot[key] = as_bits(shard.to(f"cuda:{self.device}", self.wire_dtypes[key], copy=True))
+            return
+        self._allocate_host()
+        for key, _, shard, _, _ in shards:
+            for start in range(0, shard.numel(), self.CHUNK):
+                end = start + self.CHUNK
+                self.host[key][start:end].copy_(as_bits(shard[start:end].to(self.wire_dtypes[key])))
+        self.snapshot = dict(self.host)
+
+    def copy_snapshot_to_host(self) -> None:
+        """Copy a GPU snapshot (the last broadcast's weights) into the host buffer, for an anchor. Callers wait
+        for earlier readers of the host buffer first."""
+        assert self.on_gpu
+        if self.host_buffer is None:
+            self._allocate_host()
+        self._host_ready()
+        with torch.cuda.stream(self.stream):
+            for key, bits in self.snapshot.items():
+                self.host[key].copy_(bits, non_blocking=self.pinned)
+            self.stream.synchronize()
+        self._pin_host()
+
+    def anchor_parts(self) -> list[tuple[str, list[int], int, torch.Tensor]]:
+        """This rank's slices of the full HF checkpoint of the last broadcast, from host memory (the CPU
+        snapshot, or ``copy_snapshot_to_host``)."""
         parts = []
-        for key, bits in self.snapshot.items():
-            fqn = resolve_fqn(model, key)
-            values, offset = bits.view(self.wire_dtypes[key]), self.offsets[key]
-            for start, name in self.segments[fqn]:
-                begin, end = max(start, offset), min(start + self.hf_numel[name], offset + bits.numel())
+        for key, bits in self.host.items():
+            values, (numel, offset) = bits.view(self.wire_dtypes[key]), self.shards[key]
+            for start, name in self.segments[self.key_fqn[key]]:
+                begin, end = max(start, offset), min(start + self.hf_numel[name], offset + numel)
                 if begin < end:
-                    shape = self.hf_tensors[name][0]
-                    parts.append((name, shape, begin - start, values[begin - offset : end - offset]))
+                    parts.append((name, self.hf_tensors[name][0], begin - start, values[begin - offset : end - offset]))
         return parts
 
-    @staticmethod
-    def _diff(bits: torch.Tensor, previous: torch.Tensor, chunk: int = 1 << 28) -> tuple[torch.Tensor, torch.Tensor]:
-        """Changed flat indices of ``bits`` vs the pinned host ``previous`` and their XOR values,
-        streaming ``previous`` through the GPU in chunks and updating it in place."""
-        changed, changed_values = [], []
-        for start in range(0, bits.numel(), chunk):
-            current = bits[start : start + chunk]
-            xor = current ^ previous[start : start + chunk].to(bits.device, non_blocking=True)
-            previous[start : start + chunk].copy_(current, non_blocking=True)
-            index = xor.nonzero().squeeze(1)
-            changed.append(index + start)
-            changed_values.append(xor[index])
-        return torch.cat(changed), torch.cat(changed_values)
-
     @torch.no_grad()
     def encode(
-        self, model: nn.Module, is_master: bool
-    ) -> tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor, dict[str, int]] | None:
-        """Diff the current shards against the snapshot and advance it. Returns
-        ``(entries, gaps, xor value bytes, new-weight checksum per HF tensor part)``
-        on the host, or None on the first call and when an fp8 block scale changed."""
-        first = not self.initialized
-        self.initialized = True
-        self.scales_changed = False
-        if first:
-            self._allocate_snapshot(model, is_master)
-        elif self.pinning is not None:
-            wait = time.perf_counter()
-            self.pinning.join()
-            self.pinning = None
-            get_logger().info(
-                f"Page-locked the sparse weight snapshot in {time.perf_counter() - self.pin_start:.1f}s "
-                f"(waited {time.perf_counter() - wait:.1f}s)"
-            )
+        self, shards: list[tuple], peak_reserved: int, ready: torch.cuda.Event, weights_read: threading.Event
+    ) -> tuple | None:
+        """Diff the shards against the snapshot and advance it, on the side stream once ``ready`` (recorded
+        after the optimizer step) has passed. Sets ``weights_read`` once done reading the weights. Returns
+        ``(entries, gaps, xor value bytes, new-weight checksum per HF tensor part, scales changed on this
+        rank)`` on the host, or None on the first call."""
+        torch.cuda.set_device(self.device)
+        try:
+            with torch.cuda.stream(self.stream):
+                self.stream.wait_event(ready)
+                self.key_fqn = {key: fqn for key, fqn, *_ in shards}
+                if not self.initialized:
+                    self._fill(shards)
+                    self.stream.synchronize()
+                    self.initialized = True
+                    return None
+                return self._encode(shards, peak_reserved, weights_read)
+        finally:
+            weights_read.set()
+
+    def _encode(self, shards: list[tuple], peak_reserved: int, weights_read: threading.Event) -> tuple:
+        device = torch.device(f"cuda:{self.device}")
+        previous_snapshot = self.snapshot
+        if self.on_gpu is None:
+            self._place(peak_reserved)
+            if self.on_gpu:
+                self.snapshot = {key: torch.empty_like(bits, device=device) for key, bits in self.host.items()}
+        if not self.on_gpu or previous_snapshot is not self.snapshot:
+            self._host_ready()
         names, counts, gaps, values = [], [], [], []
         checksum_names, checksums = [], []
-        for key, value in model.state_dict().items():
-            fqn = resolve_fqn(model, key)
-            if fqn not in self.segments:
-                continue
-            shard = self._local_shard(key, value, is_master)
-            if shard is None or shard[0].numel() == 0:
-                if key in self.gather and not first:
-                    value.full_tensor()  # every rank joins the all-gathers of a gathered fp8 tensor
-                    value.full_tensor()
-                continue
-            bits, offset = as_bits(shard[0].to(self.wire_dtypes[key])), shard[1]
-            if first:
-                self.snapshot[key].copy_(bits)
-                continue
+        scales_changed = False
+        for key, fqn, shard, offset, gathered in shards:
+            previous, out, wire = previous_snapshot[key], self.snapshot[key], self.wire_dtypes[key]
+            segments = [
+                (start - offset, start - offset + self.hf_numel[name], name)
+                for start, name in self.segments[fqn]
+                if start < offset + shard.numel() and start + self.hf_numel[name] > offset
+            ]
+            sums = [torch.zeros((), dtype=torch.int64, device=device) for _ in segments]
+            changed, changed_values = [], []
             if fqn in self.fp8:
-                changed, changed_values, bits = self._fp8_diff(key, fqn, value, bits, offset)
+                ue8m0, num_cols = self.fp8[fqn]
+                if gathered is not None:
+                    chunks = [(0, shard.numel())]
+                else:
+                    rows = max(FP8_BLOCK[0], self.CHUNK // num_cols // FP8_BLOCK[0] * FP8_BLOCK[0])
+                    chunks = [(start, start + rows * num_cols) for start in range(0, shard.numel(), rows * num_cols)]
             else:
-                changed, changed_values = self._diff(bits, self.snapshot[key])
-            for start, name in self.segments[fqn]:
-                begin, end = max(start, offset), min(start + self.hf_numel[name], offset + bits.numel())
-                if begin < end:
-                    checksum_names.append(name)
-                    checksums.append(bits[begin - offset : end - offset].sum(dtype=torch.int64))
+                chunks = [(start, start + self.CHUNK) for start in range(0, shard.numel(), self.CHUNK)]
+            for start, end in chunks:
+                end = min(end, shard.numel())
+                current = as_bits(shard[start:end].to(device, non_blocking=True).to(wire))
+                previous_bits = previous[start:end].to(device, non_blocking=True)
+                if fqn in self.fp8:
+                    if gathered is not None:
+                        previous_weights, current_weights = gathered
+                    else:
+                        previous_weights, current_weights = (
+                            previous_bits.view(torch.bfloat16),
+                            current.view(torch.bfloat16),
+                        )
+                    (previous_fp8, previous_scales), (current_fp8, current_scales) = (
+                        quantize_fp8_blocks(weights.reshape(-1, num_cols), ue8m0)
+                        for weights in (previous_weights, current_weights)
+                    )
+                    scales_changed |= not torch.equal(previous_scales, current_scales)
+                    if gathered is not None:
+                        previous_fp8 = previous_fp8.view(-1)[offset : offset + shard.numel()]
+                        current_fp8 = current_fp8.view(-1)[offset : offset + shard.numel()]
+                    xor = previous_fp8.view(-1) ^ current_fp8.view(-1)
+                    new_bits = current_fp8.view(-1)
+                else:
+                    xor = current ^ previous_bits
+                    new_bits = current
+                out[start:end].copy_(current, non_blocking=True)
+                index = xor.nonzero().squeeze(1)
+                changed.append(index + start)
+                changed_values.append(xor[index])
+                for i, (begin, stop, _) in enumerate(segments):
+                    begin, stop = max(begin, start), min(stop, end)
+                    if begin < stop:
+                        sums[i] += new_bits[begin - start : stop - start].sum(dtype=torch.int64)
+            checksum_names.extend(name for *_, name in segments)
+            checksums.extend(sums)
+            changed, changed_values = torch.cat(changed), torch.cat(changed_values)
             if fqn not in self.segment_starts:
-                self.segment_starts[fqn] = torch.tensor([start for start, _ in self.segments[fqn]], device=bits.device)
+                self.segment_starts[fqn] = torch.tensor([start for start, _ in self.segments[fqn]], device=device)
             starts = self.segment_starts[fqn]
             position = changed + offset
             segment = torch.searchsorted(starts, position, right=True) - 1
@@ -287,31 +397,25 @@ def encode(
             values.append(changed_values.view(torch.uint8).cpu())
             counts.append(torch.bincount(segment, minlength=len(starts)))
             names.extend(name for _, name in self.segments[fqn])
-        if first:
-            self.pin_start = time.perf_counter()
-            self.pinning.start()
-        if self.fp8 and dist.is_initialized():
-            changed_anywhere = torch.tensor(int(self.scales_changed), device=next(model.parameters()).device)
-            dist.all_reduce(changed_anywhere, op=dist.ReduceOp.MAX)
-            self.scales_changed = bool(changed_anywhere.item())
-        torch.cuda.synchronize()  # the snapshot is written asynchronously
-        if first or self.scales_changed:
-            return None
+        self.stream.synchronize()
+        weights_read.set()
+        self._pin_host()
         counts = torch.cat(counts).tolist() if counts else []
         checksums = dict(zip(checksum_names, torch.stack(checksums).tolist())) if checksums else {}
-        torch.cuda.empty_cache()
         entries = [(name, count) for name, count in zip(names, counts) if count]
         empty = torch.empty(0, dtype=torch.uint8)
-        return entries, torch.cat(gaps) if gaps else empty, torch.cat(values) if values else empty, checksums
+        gaps = torch.cat(gaps) if gaps else empty
+        return entries, gaps, torch.cat(values) if values else empty, checksums, scales_changed
 
 
 class FileSystemWeightSender(WeightSender):
     """Broadcast weights by saving a HF-compatible checkpoint (or, for LoRA
     runs, the PEFT-shaped adapter) to a shared filesystem. With ``sparse``,
     steps only carry the XOR delta against the previous broadcast (see
-    ``prime_rl.utils.sparse_weights``), written on a background thread while
-    training continues; full checkpoints (anchors) serve engines that fall off
-    the delta chain."""
+    ``prime_rl.utils.sparse_weights``). The delta is encoded and written in the
+    background while the next step trains: the trainer only waits, before its
+    next optimizer step, until the encode has read the weights. Full
+    checkpoints (anchors) serve engines that fall off the delta chain."""
 
     def __init__(
         self,
@@ -328,11 +432,45 @@ def __init__(
         self.last_step: int | None = None
         self.anchor_step: int | None = None
         self.pending_anchor: int | None = None
-        # Background writes of the last step (its delta, and its anchor on anchor steps)
-        self.writers: list[threading.Thread] = []
+        # Background work, each in order: encodes, delta writes and commits, periodic anchor writes
+        self.encodes = ThreadPoolExecutor(1)
+        self.writes = ThreadPoolExecutor(1)
+        self.anchors = ThreadPoolExecutor(1)
+        self.host_readers: list[Future] = []
+        self.errors: list[BaseException] = []
+        self.weights_read: threading.Event | None = None
         self.deferred_finish = False
         self.logger.debug("Initialized filesystem weight broadcast")
 
+    def wait_for_weight_reads(self) -> None:
+        if self.weights_read is None:
+            return
+        start = time.perf_counter()
+        self.weights_read.wait()
+        self.weights_read = None
+        self.logger.debug(f"Waited {time.perf_counter() - start:.2f}s for the sparse weight encode")
+        self._raise_errors()
+
+    def close(self) -> None:
+        for pool in (self.encodes, self.writes, self.anchors):
+            pool.shutdown(wait=True)
+        self._raise_errors()
+
+    def _raise_errors(self) -> None:
+        if self.errors:
+            raise RuntimeError("Sparse weight broadcast failed in the background") from self.errors[0]
+
+    def _submit(self, pool: ThreadPoolExecutor, fn, *args) -> Future:
+        def run() -> None:
+            try:
+                fn(*args)
+            except BaseException as error:
+                self.logger.exception(f"Sparse weight broadcast failed in {fn.__name__}")
+                self.errors.append(error)
+                raise
+
+        return pool.submit(run)
+
     def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
         if self.lora_config is not None:
             # All ranks must participate in DTensor gathering, but only master saves
@@ -357,61 +495,111 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
         if not self.config.sparse:
             self._save_full(model, step_dir)
             return
-        # A step's background writes read the snapshot, which the next encode updates.
-        for writer in self.writers:
-            writer.join()
-        self.writers, self.deferred_finish = [], False
+        self.wait_for_weight_reads()
+        self._raise_errors()
+        self.deferred_finish = False
         if self.world.is_master:
             # Engines that load a step with a delta dir report their weight layout next to it.
             (step_dir / DELTA_DIR).mkdir(exist_ok=True)
         if self.encoder is None:
-            self.encoder = SparseDeltaEncoder(model)
+            self.encoder = SparseDeltaEncoder(model, self.config.snapshot)
         if self.encoder.fp8_layout is None:
             # Until every engine worker has reported which tensors it holds in fp8, steps are full checkpoints.
             layout = [read_engine_layout(step_dir.parent) if self.world.is_master else None]
             dist.broadcast_object_list(layout, src=0)
             if layout[0] is not None:
                 self.encoder.set_fp8_layout(model, layout[0], self.world.is_master)
-        first = not self.encoder.initialized
-        start = time.perf_counter()
-        delta = self.encoder.encode(model, self.world.is_master)
-        encode_time = time.perf_counter() - start
-        if self.encoder.scales_changed:
-            self.logger.info(f"fp8 block scales changed at v{step}: publishing a full checkpoint")
+        shards = self.encoder.local_shards(model, self.world.is_master)
+        first = self.last_step is None
         base_step, self.last_step = self.last_step, step
-        if first:
-            num_bytes = sum(snapshot.nbytes for snapshot in self.encoder.snapshot.values())
-            self.logger.info(
-                f"Initialized the {num_bytes / 2**30:.1f} GiB sparse weight snapshot in {encode_time:.1f}s"
-            )
-        if first and step == 0:
-            # The run's first broadcast is the checkpoint the engines booted from.
+        # The run's first broadcast is the checkpoint the engines booted from; a resumed run's is an anchor.
+        boot = first and step == 0
+        full = (first and step != 0) or (not first and self.encoder.fp8_layout is None)
+        anchor = False
+        if not first and not full:
+            # Periodic anchors only serve engines that fall off the chain: written in the background.
+            if (
+                self.world.is_master
+                and self.pending_anchor is not None
+                and anchor_complete(self.step_dir(self.pending_anchor))
+            ):
+                self.anchor_step, self.pending_anchor = max(self.anchor_step or 0, self.pending_anchor), None
+            start_anchor = [step % self.config.anchor_interval == 0 and self.pending_anchor is None]
+            dist.broadcast_object_list(start_anchor, src=0)
+            anchor = start_anchor[0]
+            if anchor:
+                self.pending_anchor = step
+        ready = torch.cuda.Event()
+        ready.record()
+        self.weights_read = threading.Event()
+        args = (step, base_step, step_dir, shards, torch.cuda.max_memory_reserved(), ready, self.weights_read)
+        self._submit(self.encodes, self._encode, *args, boot, full, anchor)
+        if boot:
             if self.world.is_master:
                 manifest = {"step": step, "base_step": None, "files": [], "tensors": self.encoder.hf_tensors}
                 (step_dir / DELTA_DIR / MANIFEST).write_text(json.dumps(manifest))
             self.anchor_step = step
             return
-        if delta is None or self.encoder.fp8_layout is None:
-            self._save_anchor(model, step_dir).join()
-            dist.barrier()
-            self.anchor_step = step
-            return
-        writer = threading.Thread(target=self._write_delta, args=(step, base_step, step_dir, delta, encode_time))
-        writer.start()
-        self.writers.append(writer)
         self.deferred_finish = True
-        # Periodic anchors only serve engines that fall off the chain: written in the background too.
-        if (
-            self.world.is_master
-            and self.pending_anchor is not None
-            and anchor_complete(self.step_dir(self.pending_anchor))
-        ):
-            self.anchor_step, self.pending_anchor = max(self.anchor_step, self.pending_anchor), None
-        start_anchor = [step % self.config.anchor_interval == 0 and self.pending_anchor is None]
-        dist.broadcast_object_list(start_anchor, src=0)
-        if start_anchor[0]:
-            self.writers.append(self._save_anchor(model, step_dir))
-            self.pending_anchor = step
+
+    def _wait_for_host_readers(self) -> None:
+        for reader in self.host_readers:
+            reader.exception()
+        self.host_readers = []
+
+    def _encode(
+        self,
+        step: int,
+        base_step: int | None,
+        step_dir: Path,
+        shards,
+        peak_reserved: int,
+        ready,
+        weights_read: threading.Event,
+        boot: bool,
+        full: bool,
+        anchor: bool,
+    ) -> None:
+        """Encode thread: diff v{step}, then hand the delta (or, on full steps, the checkpoint) to a writer."""
+        encoder = self.encoder
+        assert encoder is not None
+        if not encoder.on_gpu:
+            self._wait_for_host_readers()  # anchors written from the host snapshot, which the encode advances
+        start = time.perf_counter()
+        delta = encoder.encode(shards, peak_reserved, ready, weights_read)
+        encode_time = time.perf_counter() - start
+        if delta is None:
+            num_bytes = sum(numel * encoder.wire_dtypes[key].itemsize for key, (numel, _) in encoder.shards.items())
+            self.logger.info(f"Recorded the {num_bytes / 2**30:.1f} GiB sparse weight snapshot in {encode_time:.1f}s")
+            if boot:
+                return
+        elif encoder.fp8 and self._scales_changed_anywhere(step_dir, delta[4]):
+            self.logger.info(f"fp8 block scales changed at v{step}: publishing a full checkpoint")
+            full = True
+        if full or anchor:
+            if encoder.on_gpu:
+                self._wait_for_host_readers()
+                encoder.copy_snapshot_to_host()
+            parts = encoder.anchor_parts()
+        if full:
+            self.host_readers.append(self._submit(self.writes, self._write_full, step, step_dir, parts))
+            return
+        self._submit(self.writes, self._write_delta, step, base_step, step_dir, delta[:4], encode_time)
+        if anchor:
+            self.host_readers.append(self._submit(self.anchors, self._write_anchor, step_dir, parts))
+
+    def _scales_changed_anywhere(self, step_dir: Path, changed: bool) -> bool:
+        """Agree through the step dir (the encode thread runs no collectives) whether any rank's fp8 block
+        scales changed."""
+        delta_dir = step_dir / DELTA_DIR
+        delta_dir.mkdir(exist_ok=True)
+        tmp = delta_dir / f".rank_{self.world.rank}.scales.tmp"
+        tmp.write_text(str(int(changed)))
+        tmp.rename(delta_dir / f"rank_{self.world.rank}.scales")
+        flags = [delta_dir / f"rank_{rank}.scales" for rank in range(self.world.world_size)]
+        while not all(flag.exists() for flag in flags):
+            time.sleep(0.05)
+        return any(flag.read_text() == "1" for flag in flags)
 
     def _write_delta(self, step: int, base_step: int, step_dir: Path, delta: tuple, encode_time: float) -> None:
         """Write this rank's delta file; the master then waits for every rank's file and commits
@@ -441,15 +629,29 @@ def _write_delta(self, step: int, base_step: int, step_dir: Path, delta: tuple,
             "tensors": self.encoder.hf_tensors,
         }
         (delta_dir / MANIFEST).write_text(json.dumps(manifest))
-        super()._finish(step, step_dir)
+        WeightSender._finish(self, step, step_dir)
         num_bytes = sum((delta_dir / name).stat().st_size for name in files)
         self.logger.info(
             f"Sparse weight delta v{step} (base v{base_step}): {num_bytes / 2**20:.1f} MiB, "
             f"encode {encode_time:.2f}s, write {time.perf_counter() - start:.2f}s (in the background)"
         )
 
+    def _write_full(self, step: int, step_dir: Path, parts: list) -> None:
+        """Publish v{step} as a full checkpoint: every rank writes its slices, the master commits."""
+        self._write_anchor(step_dir, parts)
+        if self.world.is_master:
+            while not anchor_complete(step_dir):
+                time.sleep(0.1)
+            self.anchor_step = step
+            WeightSender._finish(self, step, step_dir)
+
+    def _write_anchor(self, step_dir: Path, parts: list) -> None:
+        start = time.perf_counter()
+        write_anchor_shards(parts, step_dir, self.world.rank, self.world.world_size)
+        self.logger.debug(f"Wrote the {step_dir.name} anchor in {time.perf_counter() - start:.1f}s")
+
     def _finish(self, step: int, step_dir: Path) -> None:
-        # A delta step is committed by its writer thread once every rank's file is on disk.
+        # A delta or full step is committed by the writer thread once every rank's files are on disk.
         if not self.deferred_finish:
             super()._finish(step, step_dir)
 
@@ -459,21 +661,6 @@ def _save_full(self, model: nn.Module, step_dir: Path) -> None:
         self.logger.debug(f"Saving weights to {step_dir}")
         save_state_dict_parallel(state_dict, step_dir)
 
-    def _save_anchor(self, model: nn.Module, step_dir: Path) -> threading.Thread:
-        """Write this rank's slices of the full checkpoint from the snapshot on a thread; engines treat
-        the step as an anchor once every rank's marker exists."""
-        assert self.encoder is not None
-        parts = self.encoder.anchor_parts(model)
-        start = time.perf_counter()
-
-        def write() -> None:
-            write_anchor_shards(parts, step_dir, self.world.rank, self.world.world_size)
-            self.logger.debug(f"Wrote the {step_dir.name} anchor in {time.perf_counter() - start:.1f}s")
-
-        thread = threading.Thread(target=write)
-        thread.start()
-        return thread
-
     def _clean(self, step: int) -> None:
         # Sparse chains replay from their anchor, so keep everything since it.
         super()._clean(step if self.anchor_step is None else min(step, self.anchor_step + 1))
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index 4d7775025cc..192041c4869 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -12,7 +12,6 @@
 import json
 import os
 import re
-import threading
 from collections import defaultdict
 from collections.abc import Iterable, Iterator
 from concurrent.futures import ThreadPoolExecutor
@@ -308,20 +307,6 @@ def pin(num_bytes: int) -> torch.Tensor:
     return tensor
 
 
-def pin_in_background(num_bytes: int) -> tuple[torch.Tensor, threading.Thread]:
-    """Host memory of ``num_bytes`` and a (not yet started) thread that page-locks it with one
-    ``cudaHostRegister`` (a copy may not span two registrations). Page-locking a large buffer is slow
-    on a loaded host. No CUDA copies may touch the memory while the thread runs."""
-    tensor = torch.from_numpy(np.empty(num_bytes, np.uint8))
-
-    def register() -> None:
-        if num_bytes:
-            error = torch.cuda.cudart().cudaHostRegister(tensor.data_ptr(), num_bytes, 0)
-            assert error == 0, f"cudaHostRegister failed: {error}"
-
-    return tensor, threading.Thread(target=register)
-
-
 def unpin(tensor: torch.Tensor) -> None:
     if tensor.numel():
         torch.cuda.cudart().cudaHostUnregister(tensor.data_ptr())

From 3ef620517a89b4eeed9c85f90540d5d7948a5f82 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Thu, 15 Oct 2026 14:23:38 +0000
Subject: [PATCH 22/88] fix(weights): diff fp8-held tensors on the training
 thread; pin the snapshot from its device

---
 src/prime_rl/transports/weights/filesystem.py | 78 ++++++++++++-------
 1 file changed, 50 insertions(+), 28 deletions(-)

diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index d2b80124d0f..4e36873079c 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -62,9 +62,9 @@ class SparseDeltaEncoder:
     Tensors the engines hold in fp8 (``set_fp8_layout``: HF name -> scale format,
     from the engine layout) are quantized with vLLM's own per-block quantizer,
     previous and current weights alike, and diffed as fp8 bytes. Their 128-row
-    blocks must not straddle trainer ranks, else the trainer tensor is gathered
-    (by the caller, before ``encode``). A step that changes any block scale is
-    published as a full checkpoint.
+    blocks must not straddle trainer ranks, else the trainer tensor is gathered.
+    They are diffed on the training thread (``prepare``). A step that changes
+    any block scale is published as a full checkpoint.
     """
 
     # Elements per transient chunk of the diff (bf16: 128 MiB per temporary)
@@ -104,6 +104,8 @@ def __init__(self, model: nn.Module, snapshot: str):
         self.shards: dict[str, tuple[int, int]] = {}
         # key -> previous wire bits: GPU tensors, or views into the host buffer
         self.snapshot: dict[str, torch.Tensor] = {}
+        self.previous: dict[str, torch.Tensor] = {}
+        self.fp8_part: list | None = None
         self.on_gpu: bool | None = None  # None until placed (auto: at the first delta)
         self.host_buffer: torch.Tensor | None = None
         self.host: dict[str, torch.Tensor] = {}
@@ -173,6 +175,7 @@ def _find_gathered(self, model: nn.Module, is_master: bool) -> set[str]:
             dist.all_reduce(flags, op=dist.ReduceOp.MAX)
         return {key for key, flag in zip(keys, flags.tolist()) if flag}
 
+    @torch.no_grad()
     def local_shards(self, model: nn.Module, is_master: bool) -> list[tuple]:
         """This rank's ``(key, fqn, flat shard, global flat offset, gathered)`` for ``encode``, taken on the
         training thread: FSDP swaps module parameters during forward, but the sharded storage stays put.
@@ -299,9 +302,27 @@ def anchor_parts(self) -> list[tuple[str, list[int], int, torch.Tensor]]:
         return parts
 
     @torch.no_grad()
-    def encode(
-        self, shards: list[tuple], peak_reserved: int, ready: torch.cuda.Event, weights_read: threading.Event
-    ) -> tuple | None:
+    def prepare(self, shards: list[tuple], peak_reserved: int) -> None:
+        """The training thread's part of an encode, before ``encode``: place the snapshot at the first delta,
+        and diff the tensors held in fp8. vLLM's quantizer is torch.compile'd, and running it next to the
+        training step breaks activation checkpointing there."""
+        self.previous = self.snapshot
+        self.fp8_part = None
+        if not self.initialized:
+            return
+        if self.on_gpu is None:
+            self._place(peak_reserved)
+            if self.on_gpu:
+                device = f"cuda:{self.device}"
+                self.snapshot = {key: torch.empty_like(bits, device=device) for key, bits in self.host.items()}
+        fp8 = [shard for shard in shards if shard[1] in self.fp8]
+        if fp8:
+            if not self.on_gpu or self.previous is not self.snapshot:
+                self._host_ready()
+            self.fp8_part = self._diff(fp8)
+
+    @torch.no_grad()
+    def encode(self, shards: list[tuple], ready: torch.cuda.Event, weights_read: threading.Event) -> tuple | None:
         """Diff the shards against the snapshot and advance it, on the side stream once ``ready`` (recorded
         after the optimizer step) has passed. Sets ``weights_read`` once done reading the weights. Returns
         ``(entries, gaps, xor value bytes, new-weight checksum per HF tensor part, scales changed on this
@@ -316,24 +337,33 @@ def encode(
                     self.stream.synchronize()
                     self.initialized = True
                     return None
-                return self._encode(shards, peak_reserved, weights_read)
+                if not self.on_gpu or self.previous is not self.snapshot:
+                    self._host_ready()
+                part = self._diff([shard for shard in shards if shard[1] not in self.fp8])
+                self.stream.synchronize()
         finally:
             weights_read.set()
+        self.previous = self.snapshot
+        self._pin_host()
+        if self.fp8_part is not None:
+            part = [ours + theirs for ours, theirs in zip(part[:-1], self.fp8_part[:-1])] + [self.fp8_part[-1]]
+        names, counts, gaps, values, checksum_names, checksums, scales_changed = part
+        counts = torch.cat(counts).tolist() if counts else []
+        checksums = dict(zip(checksum_names, torch.stack(checksums).tolist())) if checksums else {}
+        entries = [(name, count) for name, count in zip(names, counts) if count]
+        empty = torch.empty(0, dtype=torch.uint8)
+        gaps = torch.cat(gaps) if gaps else empty
+        return entries, gaps, torch.cat(values) if values else empty, checksums, scales_changed
 
-    def _encode(self, shards: list[tuple], peak_reserved: int, weights_read: threading.Event) -> tuple:
+    def _diff(self, shards: list[tuple]) -> list:
+        """Diff ``shards`` against ``self.previous`` on the current stream, writing the new bits to the
+        snapshot, in chunks of ``CHUNK`` elements (fp8: whole 128-row blocks)."""
         device = torch.device(f"cuda:{self.device}")
-        previous_snapshot = self.snapshot
-        if self.on_gpu is None:
-            self._place(peak_reserved)
-            if self.on_gpu:
-                self.snapshot = {key: torch.empty_like(bits, device=device) for key, bits in self.host.items()}
-        if not self.on_gpu or previous_snapshot is not self.snapshot:
-            self._host_ready()
         names, counts, gaps, values = [], [], [], []
         checksum_names, checksums = [], []
         scales_changed = False
         for key, fqn, shard, offset, gathered in shards:
-            previous, out, wire = previous_snapshot[key], self.snapshot[key], self.wire_dtypes[key]
+            previous, out, wire = self.previous[key], self.snapshot[key], self.wire_dtypes[key]
             segments = [
                 (start - offset, start - offset + self.hf_numel[name], name)
                 for start, name in self.segments[fqn]
@@ -397,15 +427,7 @@ def _encode(self, shards: list[tuple], peak_reserved: int, weights_read: threadi
             values.append(changed_values.view(torch.uint8).cpu())
             counts.append(torch.bincount(segment, minlength=len(starts)))
             names.extend(name for _, name in self.segments[fqn])
-        self.stream.synchronize()
-        weights_read.set()
-        self._pin_host()
-        counts = torch.cat(counts).tolist() if counts else []
-        checksums = dict(zip(checksum_names, torch.stack(checksums).tolist())) if checksums else {}
-        entries = [(name, count) for name, count in zip(names, counts) if count]
-        empty = torch.empty(0, dtype=torch.uint8)
-        gaps = torch.cat(gaps) if gaps else empty
-        return entries, gaps, torch.cat(values) if values else empty, checksums, scales_changed
+        return [names, counts, gaps, values, checksum_names, checksums, scales_changed]
 
 
 class FileSystemWeightSender(WeightSender):
@@ -529,10 +551,11 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
             anchor = start_anchor[0]
             if anchor:
                 self.pending_anchor = step
+        self.encoder.prepare(shards, torch.cuda.max_memory_reserved())
         ready = torch.cuda.Event()
         ready.record()
         self.weights_read = threading.Event()
-        args = (step, base_step, step_dir, shards, torch.cuda.max_memory_reserved(), ready, self.weights_read)
+        args = (step, base_step, step_dir, shards, ready, self.weights_read)
         self._submit(self.encodes, self._encode, *args, boot, full, anchor)
         if boot:
             if self.world.is_master:
@@ -553,7 +576,6 @@ def _encode(
         base_step: int | None,
         step_dir: Path,
         shards,
-        peak_reserved: int,
         ready,
         weights_read: threading.Event,
         boot: bool,
@@ -566,7 +588,7 @@ def _encode(
         if not encoder.on_gpu:
             self._wait_for_host_readers()  # anchors written from the host snapshot, which the encode advances
         start = time.perf_counter()
-        delta = encoder.encode(shards, peak_reserved, ready, weights_read)
+        delta = encoder.encode(shards, ready, weights_read)
         encode_time = time.perf_counter() - start
         if delta is None:
             num_bytes = sum(numel * encoder.wire_dtypes[key].itemsize for key, (numel, _) in encoder.shards.items())

From 3289d97ecd410bdeeea2ab9a93efb8dc84fe5937 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Thu, 15 Oct 2026 14:26:00 +0000
Subject: [PATCH 23/88] perf(weights): skip the pre-broadcast drain and cache
 release for sparse broadcasts; anchor every anchor_interval steps without a
 collective

---
 src/prime_rl/trainer/rl/train.py              | 19 +++++------
 src/prime_rl/transports/weights/base.py       |  4 +++
 src/prime_rl/transports/weights/filesystem.py | 32 +++++++------------
 3 files changed, 25 insertions(+), 30 deletions(-)

diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py
index 4d4c7891835..f97a87af978 100644
--- a/src/prime_rl/trainer/rl/train.py
+++ b/src/prime_rl/trainer/rl/train.py
@@ -591,15 +591,16 @@ def train(config: TrainerConfig):
             broadcast_weights_time = 0
         else:
             broadcast_weights_start_time = time.perf_counter()
-            # The per-layer gather + fp8 conversion peaks ~50 GiB above the
-            # resident weights; release cached blocks (incl. offload-stream
-            # pools) so the broadcast gets the full headroom. Drain all
-            # pending work first: empty_cache returns blocks to the driver,
-            # so a still-running kernel holding a cached block (e.g. the
-            # optimizer step's tail) faults with an illegal memory access
-            # once its block is freed under it.
-            torch.cuda.synchronize()
-            torch.cuda.empty_cache()
+            if weight_sender.gathers_full_weights:
+                # The per-layer gather + fp8 conversion peaks ~50 GiB above the
+                # resident weights; release cached blocks (incl. offload-stream
+                # pools) so the broadcast gets the full headroom. Drain all
+                # pending work first: empty_cache returns blocks to the driver,
+                # so a still-running kernel holding a cached block (e.g. the
+                # optimizer step's tail) faults with an illegal memory access
+                # once its block is freed under it.
+                torch.cuda.synchronize()
+                torch.cuda.empty_cache()
             weight_sender.broadcast(model, step=progress.step)
             broadcast_weights_time = time.perf_counter() - broadcast_weights_start_time
 
diff --git a/src/prime_rl/transports/weights/base.py b/src/prime_rl/transports/weights/base.py
index 13182ea0d63..f855ef567f9 100644
--- a/src/prime_rl/transports/weights/base.py
+++ b/src/prime_rl/transports/weights/base.py
@@ -44,6 +44,10 @@ class WeightSender(ABC):
     (``.finished``). The trainer therefore runs in lockstep with its consumer
     on every transport — a broadcast nobody receives blocks and then fails."""
 
+    # Gathering full weights peaks well above the resident weights, so the trainer drains the GPU and
+    # releases cached blocks before such a broadcast.
+    gathers_full_weights = True
+
     def __init__(self, output_dir: Path, timeout: int):
         self.logger = get_logger()
         self.world = get_world()
diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index 4e36873079c..982b0c4b877 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -453,7 +453,6 @@ def __init__(
         self.encoder: SparseDeltaEncoder | None = None
         self.last_step: int | None = None
         self.anchor_step: int | None = None
-        self.pending_anchor: int | None = None
         # Background work, each in order: encodes, delta writes and commits, periodic anchor writes
         self.encodes = ThreadPoolExecutor(1)
         self.writes = ThreadPoolExecutor(1)
@@ -462,6 +461,7 @@ def __init__(
         self.errors: list[BaseException] = []
         self.weights_read: threading.Event | None = None
         self.deferred_finish = False
+        self.gathers_full_weights = not config.sparse
         self.logger.debug("Initialized filesystem weight broadcast")
 
     def wait_for_weight_reads(self) -> None:
@@ -537,20 +537,9 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
         # The run's first broadcast is the checkpoint the engines booted from; a resumed run's is an anchor.
         boot = first and step == 0
         full = (first and step != 0) or (not first and self.encoder.fp8_layout is None)
-        anchor = False
-        if not first and not full:
-            # Periodic anchors only serve engines that fall off the chain: written in the background.
-            if (
-                self.world.is_master
-                and self.pending_anchor is not None
-                and anchor_complete(self.step_dir(self.pending_anchor))
-            ):
-                self.anchor_step, self.pending_anchor = max(self.anchor_step or 0, self.pending_anchor), None
-            start_anchor = [step % self.config.anchor_interval == 0 and self.pending_anchor is None]
-            dist.broadcast_object_list(start_anchor, src=0)
-            anchor = start_anchor[0]
-            if anchor:
-                self.pending_anchor = step
+        # Periodic anchors only serve engines that fall off the chain: written in the background, each after
+        # the previous one.
+        anchor = not first and not full and step % self.config.anchor_interval == 0
         self.encoder.prepare(shards, torch.cuda.max_memory_reserved())
         ready = torch.cuda.Event()
         ready.record()
@@ -608,7 +597,7 @@ def _encode(
             return
         self._submit(self.writes, self._write_delta, step, base_step, step_dir, delta[:4], encode_time)
         if anchor:
-            self.host_readers.append(self._submit(self.anchors, self._write_anchor, step_dir, parts))
+            self.host_readers.append(self._submit(self.anchors, self._write_anchor, step, step_dir, parts))
 
     def _scales_changed_anywhere(self, step_dir: Path, changed: bool) -> bool:
         """Agree through the step dir (the encode thread runs no collectives) whether any rank's fp8 block
@@ -660,17 +649,18 @@ def _write_delta(self, step: int, base_step: int, step_dir: Path, delta: tuple,
 
     def _write_full(self, step: int, step_dir: Path, parts: list) -> None:
         """Publish v{step} as a full checkpoint: every rank writes its slices, the master commits."""
-        self._write_anchor(step_dir, parts)
+        self._write_anchor(step, step_dir, parts)
         if self.world.is_master:
-            while not anchor_complete(step_dir):
-                time.sleep(0.1)
-            self.anchor_step = step
             WeightSender._finish(self, step, step_dir)
 
-    def _write_anchor(self, step_dir: Path, parts: list) -> None:
+    def _write_anchor(self, step: int, step_dir: Path, parts: list) -> None:
         start = time.perf_counter()
         write_anchor_shards(parts, step_dir, self.world.rank, self.world.world_size)
         self.logger.debug(f"Wrote the {step_dir.name} anchor in {time.perf_counter() - start:.1f}s")
+        if self.world.is_master:
+            while not anchor_complete(step_dir):
+                time.sleep(0.1)
+            self.anchor_step = max(self.anchor_step or 0, step)
 
     def _finish(self, step: int, step_dir: Path) -> None:
         # A delta or full step is committed by the writer thread once every rank's files are on disk.

From b37392dc484b65382d578c53a986ff80b6e43ebb Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Thu, 15 Oct 2026 12:25:11 +0000
Subject: [PATCH 24/88] feat(weights): store sparse deltas as zigzag raw-bit
 differences, one byte per change

---
 docs/scaling.md                               |  2 +-
 .../src/prime_rl/configs/shared.py            |  2 +-
 .../inference/vllm/worker/filesystem.py       |  6 +-
 .../inference/vllm/worker/weight_transfer.py  | 17 +++---
 src/prime_rl/transports/weights/filesystem.py | 15 ++---
 src/prime_rl/utils/sparse_weights.py          | 59 +++++++++++++++----
 tests/unit/utils/test_sparse_weights.py       |  9 +--
 7 files changed, 76 insertions(+), 34 deletions(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index f7968b09665..1b10c8cc008 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -276,7 +276,7 @@ By default, the trainer and inference worker each allocate one transfer arena. S
 
 ### Sparse filesystem weight broadcast
 
-With `[weight_broadcast] type = "filesystem"` and `sparse = true`, the trainer writes a full checkpoint only every `anchor_interval` steps (default 25) and at its first two broadcasts. In between, it writes the XOR of the bf16 weights against the previous broadcast. Each trainer rank diffs its own shards, so no weights are gathered; with AdamW at RL learning rates about 1% of bf16 elements change per step, and a delta is about 1% of the checkpoint size. Engines read and decode a delta while they keep serving, then patch their weights in place during the pause. An engine that is not at a delta's base step reloads the last anchor and replays the deltas since. On anchor steps, engines apply the delta first and check it against the full checkpoint bit for bit, so a model whose weights cannot be patched in place fails at the second broadcast. The delta goes through vLLM's own weight loader; weights that are transformed after loading (quantized or re-laid-out) are not supported. LoRA is not supported.
+With `[weight_broadcast] type = "filesystem"` and `sparse = true`, the trainer writes a full checkpoint only every `anchor_interval` steps (default 25) and at its first two broadcasts. In between, it writes the elements of the bf16 weights that changed since the previous broadcast. Each trainer rank diffs its own shards, so no weights are gathered; with AdamW at RL learning rates about 1% of bf16 elements change per step, and a delta is about 1% of the checkpoint size. Engines read and decode a delta while they keep serving, then patch their weights in place during the pause. An engine that is not at a delta's base step reloads the last anchor and replays the deltas since. On anchor steps, engines apply the delta first and check it against the full checkpoint bit for bit, so a model whose weights cannot be patched in place fails at the second broadcast. The delta goes through vLLM's own weight loader; weights that are transformed after loading (quantized or re-laid-out) are not supported. LoRA is not supported.
 
 ### Custom Templates
 
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
index 7c425a012a0..7f9ce1db473 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
@@ -49,7 +49,7 @@ class FileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig):
     type: Literal["filesystem"] = "filesystem"
 
     sparse: bool = False
-    """Between anchors, publish only the XOR delta of the wire-dtype weights against the previous
+    """Between anchors, publish only the delta of the wire-dtype weights against the previous
     broadcast. Each trainer rank diffs its own shards; engines decode the delta while serving and
     patch their weights in place during the pause. Not supported with LoRA."""
 
diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 39f5ba15c86..07592a78301 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -161,10 +161,10 @@ def _apply_chunk(
             run = list(group)
             num_changes = sum(count for _, count in run)
             itemsize = bits.itemsize
-            xor = values[num_bytes : num_bytes + num_changes * itemsize]
+            delta = values[num_bytes : num_bytes + num_changes * itemsize]
             # A run after an odd number of narrower values is misaligned for its dtype view.
-            xor = xor.view(bits) if num_bytes % itemsize == 0 else xor.clone().view(bits)
-            self.index_map.apply_xor(run, indices[offset : offset + num_changes], xor)
+            delta = delta.view(bits) if num_bytes % itemsize == 0 else delta.clone().view(bits)
+            self.index_map.apply_delta(run, indices[offset : offset + num_changes], delta)
             offset += num_changes
             num_bytes += num_changes * itemsize
 
diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py
index a7e7cb369e8..97535c5fa7f 100644
--- a/src/prime_rl/inference/vllm/worker/weight_transfer.py
+++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py
@@ -120,7 +120,7 @@ def placeholders():
         self.numel = {name: prod(shape) for name, (shape, _) in tensors.items()}
         self.checksum_plan = None
         self.unpatchable = unpatchable
-        # Checkpoint tensors this worker holds; unpatchable ones too, so a change to them is caught in apply_xor
+        # Checkpoint tensors this worker holds; unpatchable ones too, so a change to them is caught in apply_delta
         self.names = set(self.records) | unpatchable
 
     def held_checksums(self) -> tuple[list[str], torch.Tensor, torch.Tensor]:
@@ -195,10 +195,11 @@ def fp8_layout(self, read_tensor) -> dict[str, str]:
         return layout
 
     @torch.no_grad()
-    def apply_xor(
-        self, entries: list[tuple[str, int]], indices: torch.Tensor, xor: torch.Tensor, chunk: int = 1 << 25
+    def apply_delta(
+        self, entries: list[tuple[str, int]], indices: torch.Tensor, delta: torch.Tensor, chunk: int = 1 << 25
     ) -> None:
-        """XOR ``xor`` (bits) into the params at the checkpoint ``indices`` of ``entries``."""
+        """Add ``delta`` (raw-bit differences, wrapping in the bits dtype) to the params at the checkpoint
+        ``indices`` of ``entries``."""
         changed = [name for name, _ in entries if name in self.unpatchable]
         if changed:
             raise NotImplementedError(f"Sparse weight deltas cannot patch {changed[:8]}: the engine casts them on load")
@@ -221,7 +222,7 @@ def apply_xor(
         entry_ends = torch.tensor([count for _, count in entries], device=device).cumsum(0)
         for start in range(0, indices.numel(), chunk):
             index = indices[start : start + chunk]
-            source = xor[start : start + chunk]
+            source = delta[start : start + chunk]
             entry = torch.searchsorted(
                 entry_ends, torch.arange(start, start + index.numel(), device=device), right=True
             )
@@ -232,9 +233,9 @@ def apply_xor(
                 hit = (row >= 0) & (row < table[entry, 2]) & (col >= 0) & (col < table[entry, 4])
                 target = (table[entry, 6] + row * table[entry, 7] + col * table[entry, 8])[hit]
                 owner = table[entry, 5][hit]
-                self._scatter_xor(owner, target, source[hit])
+                self._scatter_add(owner, target, source[hit])
 
-    def _scatter_xor(self, owner: torch.Tensor, target: torch.Tensor, source: torch.Tensor) -> None:
+    def _scatter_add(self, owner: torch.Tensor, target: torch.Tensor, source: torch.Tensor) -> None:
         order = torch.argsort(owner, stable=True)
         owner, target, source = owner[order], target[order], source[order]
         counts = torch.bincount(owner, minlength=len(self.params)).tolist()
@@ -249,7 +250,7 @@ def _scatter_xor(self, owner: torch.Tensor, target: torch.Tensor, source: torch.
                     f"Sparse weight delta dtype {source.dtype} does not match the engine param ({flat.dtype})"
                 )
             positions = target[begin:end]
-            flat.index_put_((positions,), flat[positions] ^ source[begin:end])
+            flat.index_put_((positions,), flat[positions] + source[begin:end])
 
 
 def _restore_checkpoint_format(model: Module) -> list[tuple]:
diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index 982b0c4b877..90cda4972c4 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -32,6 +32,7 @@
     read_engine_layout,
     write_anchor_shards,
     write_rank_file,
+    zigzag,
 )
 from prime_rl.utils.weights import (
     convert_state_dict_to_hf,
@@ -44,7 +45,7 @@
 
 
 class SparseDeltaEncoder:
-    """Rank-local XOR deltas of the trainer's weight shards, in HF coordinates.
+    """Rank-local deltas of the trainer's weight shards, in HF coordinates.
 
     Every rank keeps the wire-dtype bits of its own shards from the previous
     broadcast (the snapshot) and diffs against them, so no weights are gathered.
@@ -325,7 +326,7 @@ def prepare(self, shards: list[tuple], peak_reserved: int) -> None:
     def encode(self, shards: list[tuple], ready: torch.cuda.Event, weights_read: threading.Event) -> tuple | None:
         """Diff the shards against the snapshot and advance it, on the side stream once ``ready`` (recorded
         after the optimizer step) has passed. Sets ``weights_read`` once done reading the weights. Returns
-        ``(entries, gaps, xor value bytes, new-weight checksum per HF tensor part, scales changed on this
+        ``(entries, gaps, delta value bytes, new-weight checksum per HF tensor part, scales changed on this
         rank)`` on the host, or None on the first call."""
         torch.cuda.set_device(self.device)
         try:
@@ -400,15 +401,15 @@ def _diff(self, shards: list[tuple]) -> list:
                     if gathered is not None:
                         previous_fp8 = previous_fp8.view(-1)[offset : offset + shard.numel()]
                         current_fp8 = current_fp8.view(-1)[offset : offset + shard.numel()]
-                    xor = previous_fp8.view(-1) ^ current_fp8.view(-1)
+                    delta = current_fp8.view(-1) - previous_fp8.view(-1)
                     new_bits = current_fp8.view(-1)
                 else:
-                    xor = current ^ previous_bits
+                    delta = current - previous_bits
                     new_bits = current
                 out[start:end].copy_(current, non_blocking=True)
-                index = xor.nonzero().squeeze(1)
+                index = delta.nonzero().squeeze(1)
                 changed.append(index + start)
-                changed_values.append(xor[index])
+                changed_values.append(zigzag(delta[index]))
                 for i, (begin, stop, _) in enumerate(segments):
                     begin, stop = max(begin, start), min(stop, end)
                     if begin < stop:
@@ -433,7 +434,7 @@ def _diff(self, shards: list[tuple]) -> list:
 class FileSystemWeightSender(WeightSender):
     """Broadcast weights by saving a HF-compatible checkpoint (or, for LoRA
     runs, the PEFT-shaped adapter) to a shared filesystem. With ``sparse``,
-    steps only carry the XOR delta against the previous broadcast (see
+    steps only carry the delta against the previous broadcast (see
     ``prime_rl.utils.sparse_weights``). The delta is encoded and written in the
     background while the next step trains: the trainer only waits, before its
     next optimizer step, until the encode has read the weights. Full
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index 192041c4869..34e77866453 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -4,9 +4,11 @@
 full checkpoint on anchor steps): one ``rank_{r}.safetensors`` per trainer rank
 and a ``manifest.json``, written last. A rank file stores zstd frames over its
 entries ``(hf_name, count)``: gaps between sorted flat indices (restarted per
-entry) and the XOR of the new and the previous raw bits. XOR is exact for any
-dtype, but applying it twice reverts it, so a delta must be applied on top of
-exactly its ``base_step``; per-tensor checksums of the new weights catch it if not.
+entry) and the differences of the new and the previous raw bits (wrapping, in
+the bits dtype; mostly a few ulps), zigzag-encoded. Gaps and differences are
+stored as one byte each, with values of 255 or more in an int32 escape stream.
+A delta must be applied on top of exactly its ``base_step``; per-tensor
+checksums of the new weights catch it if not.
 """
 
 import json
@@ -16,6 +18,7 @@
 from collections.abc import Iterable, Iterator
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import ExitStack
+from itertools import groupby
 from math import prod
 from pathlib import Path
 from typing import BinaryIO
@@ -57,6 +60,33 @@ def quantize_fp8_blocks(weight: torch.Tensor, ue8m0: bool) -> tuple[torch.Tensor
     return quantized.view(torch.uint8), scales
 
 
+def zigzag(delta: torch.Tensor) -> torch.Tensor:
+    """Raw-bit differences (bits dtype, wrapping) -> unsigned values of the same width, small for
+    small steps of either sign: 0, -1, 1, -2 -> 0, 1, 2, 3."""
+    signed = delta.view(torch.int8) if delta.dtype == torch.uint8 else delta
+    return ((signed << 1) ^ (signed >> (8 * signed.element_size() - 1))).view(delta.dtype)
+
+
+def _unpack_values(data: np.ndarray, entries: list[tuple[str, int]], tensors: dict) -> np.ndarray:
+    """The per-change unsigned values of ``entries`` from their raw bytes (each entry in its bits dtype)."""
+    out, offset = [], 0
+    for itemsize, group in groupby(entries, key=lambda entry: bits_dtype(tensors[entry[0]][1]).itemsize):
+        num_bytes = sum(count for _, count in group) * itemsize
+        out.append(data[offset : offset + num_bytes].view(f"<u{itemsize}").astype(np.int64))
+        offset += num_bytes
+    return np.concatenate(out) if out else np.empty(0, np.int64)
+
+
+def _pack_values(values: np.ndarray, entries: list[tuple[str, int]], tensors: dict) -> np.ndarray:
+    """Inverse of ``_unpack_values`` for signed values, wrapping to each entry's bits dtype."""
+    out, offset = [], 0
+    for itemsize, group in groupby(entries, key=lambda entry: bits_dtype(tensors[entry[0]][1]).itemsize):
+        count = sum(count for _, count in group)
+        out.append(values[offset : offset + count].astype(f"<i{itemsize}").view(np.uint8))
+        offset += count
+    return np.concatenate(out) if out else np.empty(0, np.uint8)
+
+
 def encode_gaps(indices: torch.Tensor, entry_starts: torch.Tensor) -> torch.Tensor:
     """Sorted flat indices -> ``index - previous - 1``, restarting at every
     ``entry_starts`` position (where the gap is the index itself)."""
@@ -76,8 +106,9 @@ def write_rank_file(
     frame_size: int = 1 << 20,
 ) -> None:
     """Write a rank's delta as zstd frames of about ``frame_size`` changes each (whole entries), so
-    readers can skip frames without entries they hold. Gaps are stored as uint8, with gaps of 255 or
-    more in a separate int32 stream. ``checksums`` (HF name -> ``bits_checksum`` of the new weights
+    readers can skip frames without entries they hold. ``values`` holds each entry's ``zigzag``
+    differences in its bits dtype. Gaps and values are stored as uint8; those of 255 or more go to
+    the int32 escape stream (gaps first). ``checksums`` (HF name -> ``bits_checksum`` of the new weights
     this rank holds) let engines verify every step. Written under a temporary name, then renamed."""
     options = {pyzstd.CParameter.compressionLevel: 1, pyzstd.CParameter.checksumFlag: 1}
     frames, size = [[]], 0
@@ -94,10 +125,12 @@ def compress(frame: list[int]) -> tuple[bytes, bytes, bytes]:
         first, last = frame[0], frame[-1]
         frame_gaps = gaps[(gap_ends[first - 1] if first else 0) : gap_ends[last]]
         frame_values = values[(value_ends[first - 1] if first else 0) : value_ends[last]]
+        frame_values = _unpack_values(frame_values, entries[first : last + 1], tensors)
+        escapes = np.concatenate([frame_gaps[frame_gaps >= 255], frame_values[frame_values >= 255]])
         return (
             pyzstd.compress(np.minimum(frame_gaps, 255).astype(np.uint8).tobytes(), options),
-            pyzstd.compress(frame_gaps[frame_gaps >= 255].tobytes(), options),
-            pyzstd.compress(frame_values.tobytes(), options),
+            pyzstd.compress(escapes.astype(np.int32).tobytes(), options),
+            pyzstd.compress(np.minimum(frame_values, 255).astype(np.uint8).tobytes(), options),
         )
 
     # Few threads: this runs in the background next to training (e.g. a CPU optimizer step)
@@ -146,8 +179,13 @@ def _rank_file_frames(
 def _read_frame(fd: int, frame, tensors: dict, names: set[str] | None) -> list[tuple[str, np.ndarray, np.ndarray]]:
     frame_entries, ((gaps_at, gaps_size), (escapes_at, escapes_size), (values_at, values_size)) = frame
     gaps = np.frombuffer(pyzstd.decompress(os.pread(fd, gaps_size, gaps_at)), np.uint8).astype(np.int32)
-    gaps[gaps == 255] = np.frombuffer(pyzstd.decompress(os.pread(fd, escapes_size, escapes_at)), np.int32)
-    values = np.frombuffer(pyzstd.decompress(os.pread(fd, values_size, values_at)), np.uint8)
+    escapes = np.frombuffer(pyzstd.decompress(os.pread(fd, escapes_size, escapes_at)), np.int32)
+    gap_escapes = gaps == 255
+    num_gap_escapes = int(gap_escapes.sum())
+    gaps[gap_escapes] = escapes[:num_gap_escapes]
+    values = np.frombuffer(pyzstd.decompress(os.pread(fd, values_size, values_at)), np.uint8).astype(np.int64)
+    values[values == 255] = escapes[num_gap_escapes:].astype(np.int64) & 0xFFFFFFFF
+    values = _pack_values((values >> 1) ^ -(values & 1), frame_entries, tensors)
     out, gap_offset, value_offset = [], 0, 0
     for name, count in frame_entries:
         num_bytes = count * bits_dtype(tensors[name][1]).itemsize
@@ -274,7 +312,8 @@ def read_delta(
     step_dir: Path, names: set[str] | None = None, num_threads: int = 16
 ) -> tuple[list[tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]], dict[str, int]]:
     """Read and decompress every rank file of a delta step dir into ``(entries,
-    gaps, xor value bytes)`` in host memory; ``decode_gaps`` turns the
+    gaps, value bytes)`` in host memory (values: new minus previous raw bits, in each entry's bits
+    dtype); ``decode_gaps`` turns the
     gaps into flat indices on the device that applies them. With ``names``, keep
     only those tensors (frames without any are not read).
     Also returns the new weights' ``bits_checksum`` per HF tensor, summed over ranks."""
diff --git a/tests/unit/utils/test_sparse_weights.py b/tests/unit/utils/test_sparse_weights.py
index fad2c4fc915..824186d6d46 100644
--- a/tests/unit/utils/test_sparse_weights.py
+++ b/tests/unit/utils/test_sparse_weights.py
@@ -11,6 +11,7 @@
     encode_gaps,
     read_delta,
     write_rank_file,
+    zigzag,
 )
 
 
@@ -30,13 +31,13 @@ def test_sparse_delta_roundtrip_is_exact(tmp_path):
         for name, rows in names:
             checksums[name] = bits_checksum(as_bits(current[name][rows]))
             offset = 0 if rows.start is None else rows.start * 32
-            xor = as_bits(current[name][rows]).reshape(-1) ^ as_bits(previous[name][rows]).reshape(-1)
-            changed = xor.nonzero().squeeze(1)
+            delta = as_bits(current[name][rows]).reshape(-1) - as_bits(previous[name][rows]).reshape(-1)
+            changed = delta.nonzero().squeeze(1)
             starts = torch.zeros_like(changed, dtype=torch.bool)
             starts[:1] = True
             entries.append((name, len(changed)))
             gaps.append(encode_gaps(changed + offset, starts))
-            values.append(xor[changed].view(torch.uint8))
+            values.append(zigzag(delta[changed]).view(torch.uint8))
         gaps, values = torch.cat(gaps).numpy(), torch.cat(values).numpy()
         # Small frames, so a frame boundary falls inside rank 1's entries
         write_rank_file(delta_dir / f"rank_{rank}.safetensors", entries, gaps, values, tensors, checksums, frame_size=8)
@@ -52,7 +53,7 @@ def apply(names):
             i = b = 0
             for name, count in entries:
                 itemsize = patched[name].element_size()
-                patched[name][indices[i : i + count]] ^= (
+                patched[name][indices[i : i + count]] += (
                     values[b : b + count * itemsize].clone().view(patched[name].dtype)
                 )
                 i += count

From fa14059e97ccd423776a27e974e33e049707c41f Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Thu, 15 Oct 2026 14:27:46 +0000
Subject: [PATCH 25/88] docs: sparse weight broadcast snapshot placement and
 background encode

---
 docs/scaling.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index 1b10c8cc008..9fc1508cd61 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -276,7 +276,7 @@ By default, the trainer and inference worker each allocate one transfer arena. S
 
 ### Sparse filesystem weight broadcast
 
-With `[weight_broadcast] type = "filesystem"` and `sparse = true`, the trainer writes a full checkpoint only every `anchor_interval` steps (default 25) and at its first two broadcasts. In between, it writes the elements of the bf16 weights that changed since the previous broadcast. Each trainer rank diffs its own shards, so no weights are gathered; with AdamW at RL learning rates about 1% of bf16 elements change per step, and a delta is about 1% of the checkpoint size. Engines read and decode a delta while they keep serving, then patch their weights in place during the pause. An engine that is not at a delta's base step reloads the last anchor and replays the deltas since. On anchor steps, engines apply the delta first and check it against the full checkpoint bit for bit, so a model whose weights cannot be patched in place fails at the second broadcast. The delta goes through vLLM's own weight loader; weights that are transformed after loading (quantized or re-laid-out) are not supported. LoRA is not supported.
+With `[weight_broadcast] type = "filesystem"` and `sparse = true`, the trainer publishes only the elements of the weights that changed since the previous broadcast, in the engine's weight format (bf16, or fp8 for engines that quantize online to `fp8_per_block`). Each trainer rank diffs its own shards against a copy of the previous broadcast (`snapshot = "auto"` keeps it on the GPU when it fits next to the first training step's peak memory, else in pinned host memory), so no weights are gathered. The diff runs on a side CUDA stream while the next step trains; the trainer only waits for it before the next optimizer step. With AdamW at RL learning rates about 1% of bf16 elements change per step. A full checkpoint (anchor) is written in the background every `anchor_interval` steps (default 25); the first broadcast is the checkpoint the engines booted from. Engines read and decode a delta while they keep serving, patch their weights in place during the pause and verify every step against trainer checksums. An engine that is not at a delta's base step reloads the last anchor and replays the deltas since. Weights the engine transforms after loading other than online `fp8_per_block`, and LoRA, are not supported.
 
 ### Custom Templates
 

From 9afadaa4985a34629c7fd214959d7795c833d4f3 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Thu, 15 Oct 2026 18:08:29 +0000
Subject: [PATCH 26/88] feat(weights): stream sparse deltas in parts; record
 the v0 snapshot during the handshake; never page-lock during training

- Trainer ranks hand the writer a part file every 64M changes while encoding; the manifest comes first and a per-rank parts marker last. Receivers start the engines' prefetch at the manifest, and engines read part files as they appear, on reader/copy threads bound to the GPU's NUMA node.
- The first broadcast records the snapshot (and fills the host buffer, in every mode) while the master waits for the receivers, joined before the first training step.
- The host buffer is page-locked only while no step runs (snapshot=cpu at v0, auto->host at v1): registering it in the background stalled the trainer's CUDA calls (hang at the first anchor with snapshot=gpu).
- Fewer kernel launches and host syncs per step in the encode: one reduction per run of equal-size segments for checksums, one D2H per part; fp8 scale changes agreed with an all_reduce on the training thread.
---
 .../inference/vllm/worker/filesystem.py       |  62 ++-
 src/prime_rl/transports/weights/base.py       |   4 +
 src/prime_rl/transports/weights/filesystem.py | 422 ++++++++++--------
 src/prime_rl/utils/sparse_weights.py          |  99 ++--
 tests/unit/utils/test_sparse_weights.py       |   9 +-
 5 files changed, 380 insertions(+), 216 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 07592a78301..415fba31270 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -6,6 +6,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING
 
+import numpy as np
 import torch
 from safetensors import safe_open
 from torch.nn import Module
@@ -28,6 +29,7 @@
     anchor_tensors,
     bits_dtype,
     decode_gaps,
+    numa_cpus,
     pin,
     read_delta,
     read_manifest,
@@ -45,6 +47,7 @@
 logger = init_logger("vllm.inference.vllm.worker_filesystem")
 
 _prefetch_pool = ThreadPoolExecutor(1)
+_copy_pool: ThreadPoolExecutor | None = None
 
 _SAFETENSORS_DTYPES = {"BF16": "bfloat16", "F16": "float16", "F32": "float32"}
 
@@ -83,24 +86,59 @@ def _read_delta(
         self, step_dir: Path
     ) -> tuple[list[tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]], dict[str, int]]:
         """Read this worker's part of a delta into a pinned host buffer that is reused across steps
-        (allocating pinned memory every step is slow on a loaded host)."""
+        (allocating pinned memory every step is slow on a loaded host), in chunks of up to 2^26 changes
+        of consecutive frames. Reads and copies run on the CPUs of the GPU's NUMA node."""
+        global _copy_pool
+        cpus = numa_cpus(torch.cuda.current_device())
+        if cpus is not None:
+            os.sched_setaffinity(0, cpus)
+        if _copy_pool is None:
+            _copy_pool = ThreadPoolExecutor(
+                16, initializer=None if cpus is None else lambda: os.sched_setaffinity(0, cpus)
+            )
         start = time.perf_counter()
-        delta, checksums = read_delta(step_dir, None if self.index_map is None else self.index_map.names)
+        frames, checksums = read_delta(step_dir, None if self.index_map is None else self.index_map.names, cpus=cpus)
         read_time = time.perf_counter() - start
-        num_bytes = sum(-(-gaps.nbytes // 4) * 4 + -(-values.nbytes // 4) * 4 for _, gaps, values in delta)
+        chunks, size = [[]], 0
+        for frame in frames:
+            if size >= 1 << 26:
+                chunks.append([])
+                size = 0
+            chunks[-1].append(frame)
+            size += len(frame[1])
+        layout, num_bytes = [], 0
+        for chunk in chunks:
+            gaps_bytes = -(-sum(gaps.nbytes for _, gaps, _ in chunk) // 4) * 4
+            values_bytes = -(-sum(values.nbytes for _, _, values in chunk) // 4) * 4
+            layout.append((num_bytes, num_bytes + gaps_bytes))
+            num_bytes += gaps_bytes + values_bytes
         if self.host_buffer is None or self.host_buffer.numel() < num_bytes:
             if self.host_buffer is not None:
                 unpin(self.host_buffer)
             self.host_buffer = pin(num_bytes * 5 // 4)
-        pinned, offset = [], 0
-        for entries, gaps, values in delta:
-            parts = []
-            for array in (gaps, values):
-                part = self.host_buffer[offset : offset + array.nbytes].view(array.dtype)
-                part.copy_(array)
-                parts.append(part)
-                offset += -(-array.nbytes // 4) * 4
-            pinned.append((entries, *parts))
+        buffer = self.host_buffer.numpy()
+        copies, pinned = [], []
+        for chunk, (gaps_at, values_at) in zip(chunks, layout):
+            gaps_end, values_end = gaps_at, values_at
+            for _, gaps, values in chunk:
+                copies.append(
+                    _copy_pool.submit(np.copyto, buffer[gaps_end : gaps_end + gaps.nbytes].view(gaps.dtype), gaps)
+                )
+                copies.append(
+                    _copy_pool.submit(np.copyto, buffer[values_end : values_end + values.nbytes], values.view(np.uint8))
+                )
+                gaps_end += gaps.nbytes
+                values_end += values.nbytes
+            if chunk:
+                pinned.append(
+                    (
+                        [entry for entries, _, _ in chunk for entry in entries],
+                        self.host_buffer[gaps_at:gaps_end].view(torch.int32),
+                        self.host_buffer[values_at:values_end],
+                    )
+                )
+        for copy in copies:
+            copy.result()
         logger.info(
             f"Read sparse weight delta {step_dir.name}: {num_bytes / 2**30:.2f} GiB, "
             f"read+decompress {read_time:.2f}s, pin {time.perf_counter() - start - read_time:.2f}s"
diff --git a/src/prime_rl/transports/weights/base.py b/src/prime_rl/transports/weights/base.py
index f855ef567f9..baeeb8ef6f9 100644
--- a/src/prime_rl/transports/weights/base.py
+++ b/src/prime_rl/transports/weights/base.py
@@ -59,6 +59,7 @@ def broadcast(self, model: nn.Module, step: int) -> None:
         """Broadcast policy v{step} to the inference pool."""
         start_time = time.perf_counter()
         step_dir = self.step_dir(step)
+        self._prepare(model, step)
         if self.world.is_master:
             # Reset per attempt so a re-broadcast (e.g. on resume) never trips
             # the consumer or the trainer on stale markers of a previous run.
@@ -96,6 +97,9 @@ def _wait_for_receiver_ready(self, step_dir: Path) -> None:
             time.sleep(0.1)
         self.logger.debug("Receiver ready, starting the transfer")
 
+    def _prepare(self, model: nn.Module, step: int) -> None:
+        """Work on v{step} that needs no consumer, started before the handshake."""
+
     @abstractmethod
     def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
         """Move v{step}'s weights to the consumer. All ranks enter only after
diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index 90cda4972c4..f9716bde229 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -1,7 +1,10 @@
+import asyncio
 import json
 import threading
 import time
+from bisect import bisect_left, bisect_right
 from collections import defaultdict
+from collections.abc import Callable
 from concurrent.futures import Future, ThreadPoolExecutor
 from math import prod
 from pathlib import Path
@@ -29,6 +32,8 @@
     bits_dtype,
     encode_gaps,
     quantize_fp8_blocks,
+    rank_part,
+    rank_parts_marker,
     read_engine_layout,
     write_anchor_shards,
     write_rank_file,
@@ -55,10 +60,12 @@ class SparseDeltaEncoder:
 
     ``encode`` runs on a background thread and a side CUDA stream while the
     next step trains, without collectives; it only reads the weights before
-    ``weights_read`` is set. The snapshot lives on the GPU when it fits next to
-    the first training step's peak memory, else in one host buffer. With a GPU
-    snapshot, the same host buffer holds anchor copies instead, so a rank has at
-    most one shard-sized host copy.
+    ``weights_read`` is set, and hands the delta to the writer in parts as it
+    goes. The snapshot lives on the GPU when it fits next to the first training
+    step's peak memory, else in one page-locked host buffer (locked only between
+    training steps: registering it can stall the trainer's CUDA calls). With a
+    GPU snapshot, the same host buffer (pageable) holds anchor copies instead, so
+    a rank has at most one shard-sized host copy.
 
     Tensors the engines hold in fp8 (``set_fp8_layout``: HF name -> scale format,
     from the engine layout) are quantized with vLLM's own per-block quantizer,
@@ -70,6 +77,8 @@ class SparseDeltaEncoder:
 
     # Elements per transient chunk of the diff (bf16: 128 MiB per temporary)
     CHUNK = 1 << 26
+    # Changes per part handed to the writer
+    PART = 1 << 26
 
     def __init__(self, model: nn.Module, snapshot: str):
         keep_in_fp32 = getattr(model, "keep_in_fp32_for_weight_transfer", None)
@@ -100,17 +109,17 @@ def __init__(self, model: nn.Module, snapshot: str):
         self.fp8_layout: dict[str, str] | None = None
         self.gather: set[str] = set()
         self.segment_starts: dict[str, torch.Tensor] = {}
+        self.shard_segments: dict[str, tuple[list[str], list[int], list[int]]] = {}
         self.snapshot_placement = snapshot
         # key -> (local numel, global flat offset) of this rank's shard
         self.shards: dict[str, tuple[int, int]] = {}
         # key -> previous wire bits: GPU tensors, or views into the host buffer
         self.snapshot: dict[str, torch.Tensor] = {}
         self.previous: dict[str, torch.Tensor] = {}
-        self.fp8_part: list | None = None
+        self.fp8_parts: list[tuple] = []
         self.on_gpu: bool | None = None  # None until placed (auto: at the first delta)
         self.host_buffer: torch.Tensor | None = None
         self.host: dict[str, torch.Tensor] = {}
-        self.pinning: threading.Thread | None = None
         self.pinned = False
         self.initialized = False
         self.device = torch.cuda.current_device()
@@ -221,30 +230,17 @@ def _allocate_host(self) -> None:
             )
             offset += sizes[key]
 
-    def _host_ready(self) -> None:
-        """Wait for the page-locking thread: no CUDA copy may touch the buffer while it runs."""
-        if self.pinning is not None:
-            start = time.perf_counter()
-            self.pinning.join()
-            self.pinning, self.pinned = None, True
-            get_logger().info(
-                f"Page-locked the {self.host_buffer.numel() / 2**30:.1f} GiB sparse weight host buffer in "
-                f"{time.perf_counter() - self.pin_start:.1f}s (waited {time.perf_counter() - start:.1f}s)"
-            )
-
-    def _pin_host(self) -> None:
-        if self.host_buffer is None or self.pinned or self.pinning is not None:
-            return
-        buffer = self.host_buffer
-
-        def register() -> None:
-            torch.cuda.set_device(self.device)
-            error = torch.cuda.cudart().cudaHostRegister(buffer.data_ptr(), buffer.numel(), 0)
-            assert error == 0, f"cudaHostRegister failed: {error}"
-
-        self.pin_start = time.perf_counter()
-        self.pinning = threading.Thread(target=register)
-        self.pinning.start()
+    def _page_lock(self) -> None:
+        """Page-lock the host buffer with one ``cudaHostRegister`` (a copy may not span two registrations).
+        Only while no training step runs: registering a large buffer can stall the trainer's CUDA calls."""
+        start = time.perf_counter()
+        error = torch.cuda.cudart().cudaHostRegister(self.host_buffer.data_ptr(), self.host_buffer.numel(), 0)
+        assert error == 0, f"cudaHostRegister failed: {error}"
+        self.pinned = True
+        get_logger().info(
+            f"Page-locked the {self.host_buffer.numel() / 2**30:.1f} GiB sparse weight host buffer in "
+            f"{time.perf_counter() - start:.1f}s"
+        )
 
     def _place(self, peak_reserved: int) -> None:
         """Keep the snapshot on the GPU if it fits next to the first training step's peak, else on the host."""
@@ -262,33 +258,32 @@ def _place(self, peak_reserved: int) -> None:
         )
 
     def _fill(self, shards: list[tuple]) -> None:
-        """The first broadcast only records the weights: on the GPU if forced there, else in the host buffer
-        (placed at the first delta, against the measured training peak)."""
+        """The first broadcast only records the weights, in the host buffer, and on the GPU if forced there
+        (else placed at the first delta, against the measured training peak). The host buffer is filled in
+        every mode, so its pages are resident before training: faulting them in later stalls the trainer."""
         self.shards = {key: (shard.numel(), offset) for key, _, shard, offset, _ in shards}
-        if self.snapshot_placement == "gpu":
-            self.on_gpu = True
-            for key, _, shard, _, _ in shards:
-                self.snapshot[key] = as_bits(shard.to(f"cuda:{self.device}", self.wire_dtypes[key], copy=True))
-            return
         self._allocate_host()
+        if self.snapshot_placement == "cpu":
+            self._page_lock()
         for key, _, shard, _, _ in shards:
             for start in range(0, shard.numel(), self.CHUNK):
                 end = start + self.CHUNK
-                self.host[key][start:end].copy_(as_bits(shard[start:end].to(self.wire_dtypes[key])))
-        self.snapshot = dict(self.host)
+                self.host[key][start:end].copy_(as_bits(shard[start:end].to(self.wire_dtypes[key])), non_blocking=True)
+        if self.snapshot_placement == "gpu":
+            self.on_gpu = True
+            for key, _, shard, _, _ in shards:
+                self.snapshot[key] = as_bits(shard.to(f"cuda:{self.device}", self.wire_dtypes[key], copy=True))
+        else:
+            self.snapshot = dict(self.host)
 
     def copy_snapshot_to_host(self) -> None:
         """Copy a GPU snapshot (the last broadcast's weights) into the host buffer, for an anchor. Callers wait
         for earlier readers of the host buffer first."""
         assert self.on_gpu
-        if self.host_buffer is None:
-            self._allocate_host()
-        self._host_ready()
         with torch.cuda.stream(self.stream):
             for key, bits in self.snapshot.items():
                 self.host[key].copy_(bits, non_blocking=self.pinned)
             self.stream.synchronize()
-        self._pin_host()
 
     def anchor_parts(self) -> list[tuple[str, list[int], int, torch.Tensor]]:
         """This rank's slices of the full HF checkpoint of the last broadcast, from host memory (the CPU
@@ -303,31 +298,33 @@ def anchor_parts(self) -> list[tuple[str, list[int], int, torch.Tensor]]:
         return parts
 
     @torch.no_grad()
-    def prepare(self, shards: list[tuple], peak_reserved: int) -> None:
+    def prepare(self, shards: list[tuple], peak_reserved: int) -> bool:
         """The training thread's part of an encode, before ``encode``: place the snapshot at the first delta,
-        and diff the tensors held in fp8. vLLM's quantizer is torch.compile'd, and running it next to the
-        training step breaks activation checkpointing there."""
+        and diff the tensors held in fp8 (vLLM's quantizer is torch.compile'd, and running it next to the
+        training step breaks activation checkpointing there). Returns whether an fp8 block scale changed
+        on this rank."""
         self.previous = self.snapshot
-        self.fp8_part = None
+        self.fp8_parts = []
         if not self.initialized:
-            return
+            return False
         if self.on_gpu is None:
             self._place(peak_reserved)
             if self.on_gpu:
                 device = f"cuda:{self.device}"
                 self.snapshot = {key: torch.empty_like(bits, device=device) for key, bits in self.host.items()}
+            elif not self.pinned:
+                self._page_lock()
         fp8 = [shard for shard in shards if shard[1] in self.fp8]
-        if fp8:
-            if not self.on_gpu or self.previous is not self.snapshot:
-                self._host_ready()
-            self.fp8_part = self._diff(fp8)
+        return self._diff(fp8, self.fp8_parts.append) if fp8 else False
 
     @torch.no_grad()
-    def encode(self, shards: list[tuple], ready: torch.cuda.Event, weights_read: threading.Event) -> tuple | None:
+    def encode(
+        self, shards: list[tuple], ready: torch.cuda.Event, weights_read: threading.Event, emit: Callable
+    ) -> None:
         """Diff the shards against the snapshot and advance it, on the side stream once ``ready`` (recorded
-        after the optimizer step) has passed. Sets ``weights_read`` once done reading the weights. Returns
-        ``(entries, gaps, delta value bytes, new-weight checksum per HF tensor part, scales changed on this
-        rank)`` on the host, or None on the first call."""
+        after the optimizer step) has passed; the first call only records the snapshot. Hands the delta to
+        ``emit`` in parts ``(entries, gaps, value bytes, new-weight checksum per HF tensor part)`` on the
+        host as they are encoded. Sets ``weights_read`` once done reading the weights."""
         torch.cuda.set_device(self.device)
         try:
             with torch.cuda.stream(self.stream):
@@ -337,40 +334,63 @@ def encode(self, shards: list[tuple], ready: torch.cuda.Event, weights_read: thr
                     self._fill(shards)
                     self.stream.synchronize()
                     self.initialized = True
-                    return None
-                if not self.on_gpu or self.previous is not self.snapshot:
-                    self._host_ready()
-                part = self._diff([shard for shard in shards if shard[1] not in self.fp8])
+                    return
+                for part in self.fp8_parts:
+                    emit(part)
+                self._diff([shard for shard in shards if shard[1] not in self.fp8], emit)
                 self.stream.synchronize()
         finally:
             weights_read.set()
         self.previous = self.snapshot
-        self._pin_host()
-        if self.fp8_part is not None:
-            part = [ours + theirs for ours, theirs in zip(part[:-1], self.fp8_part[:-1])] + [self.fp8_part[-1]]
-        names, counts, gaps, values, checksum_names, checksums, scales_changed = part
-        counts = torch.cat(counts).tolist() if counts else []
-        checksums = dict(zip(checksum_names, torch.stack(checksums).tolist())) if checksums else {}
-        entries = [(name, count) for name, count in zip(names, counts) if count]
-        empty = torch.empty(0, dtype=torch.uint8)
-        gaps = torch.cat(gaps) if gaps else empty
-        return entries, gaps, torch.cat(values) if values else empty, checksums, scales_changed
-
-    def _diff(self, shards: list[tuple]) -> list:
+
+    def _segments(self, key: str, fqn: str, offset: int, numel: int) -> tuple[list[str], list[int], list[int]]:
+        """The HF tensor parts in this rank's shard of ``key``: names, and begin/end in the shard."""
+        if key not in self.shard_segments:
+            names, begins, ends = [], [], []
+            for start, name in self.segments[fqn]:
+                if start < offset + numel and start + self.hf_numel[name] > offset:
+                    names.append(name)
+                    begins.append(max(start - offset, 0))
+                    ends.append(min(start - offset + self.hf_numel[name], numel))
+            self.shard_segments[key] = (names, begins, ends)
+        return self.shard_segments[key]
+
+    @staticmethod
+    def _segment_sums(bits: torch.Tensor, begins: list[int], ends: list[int]) -> torch.Tensor:
+        """Wrapping int64 sums of ``bits[begin:end]``, one reduction per run of adjacent equal-length
+        segments (e.g. unstacked experts)."""
+        sums, i = [], 0
+        while i < len(begins):
+            length, j = ends[i] - begins[i], i + 1
+            while j < len(begins) and begins[j] == ends[j - 1] and ends[j] - begins[j] == length:
+                j += 1
+            sums.append(bits[begins[i] : begins[i] + (j - i) * length].view(j - i, length).sum(1, dtype=torch.int64))
+            i = j
+        return torch.cat(sums)
+
+    def _diff(self, shards: list[tuple], emit: Callable) -> bool:
         """Diff ``shards`` against ``self.previous`` on the current stream, writing the new bits to the
-        snapshot, in chunks of ``CHUNK`` elements (fp8: whole 128-row blocks)."""
+        snapshot, in chunks of ``CHUNK`` elements (fp8: whole 128-row blocks). Copies the delta to the
+        host and hands it to ``emit`` every ``PART`` changes. Returns whether an fp8 block scale changed."""
         device = torch.device(f"cuda:{self.device}")
-        names, counts, gaps, values = [], [], [], []
-        checksum_names, checksums = [], []
+        pending: list[tuple] = []
+        num_pending = 0
         scales_changed = False
+
+        def flush() -> None:
+            names = [name for part in pending for name in part[0]]
+            counts = torch.cat([part[1] for part in pending]).tolist()
+            gaps = torch.cat([part[2] for part in pending]).cpu()
+            values = torch.cat([part[3] for part in pending]).cpu()
+            checksum_names = [name for part in pending for name in part[4]]
+            checksums = dict(zip(checksum_names, torch.cat([part[5] for part in pending]).tolist()))
+            emit(([(name, count) for name, count in zip(names, counts) if count], gaps, values, checksums))
+            pending.clear()
+
         for key, fqn, shard, offset, gathered in shards:
             previous, out, wire = self.previous[key], self.snapshot[key], self.wire_dtypes[key]
-            segments = [
-                (start - offset, start - offset + self.hf_numel[name], name)
-                for start, name in self.segments[fqn]
-                if start < offset + shard.numel() and start + self.hf_numel[name] > offset
-            ]
-            sums = [torch.zeros((), dtype=torch.int64, device=device) for _ in segments]
+            segment_names, begins, ends = self._segments(key, fqn, offset, shard.numel())
+            sums = torch.zeros(len(segment_names), dtype=torch.int64, device=device)
             changed, changed_values = [], []
             if fqn in self.fp8:
                 ue8m0, num_cols = self.fp8[fqn]
@@ -410,12 +430,13 @@ def _diff(self, shards: list[tuple]) -> list:
                 index = delta.nonzero().squeeze(1)
                 changed.append(index + start)
                 changed_values.append(zigzag(delta[index]))
-                for i, (begin, stop, _) in enumerate(segments):
-                    begin, stop = max(begin, start), min(stop, end)
-                    if begin < stop:
-                        sums[i] += new_bits[begin - start : stop - start].sum(dtype=torch.int64)
-            checksum_names.extend(name for *_, name in segments)
-            checksums.extend(sums)
+                first, last = bisect_right(ends, start), bisect_left(begins, end)
+                if first < last:
+                    sums[first:last] += self._segment_sums(
+                        new_bits,
+                        [max(begin, start) - start for begin in begins[first:last]],
+                        [min(stop, end) - start for stop in ends[first:last]],
+                    )
             changed, changed_values = torch.cat(changed), torch.cat(changed_values)
             if fqn not in self.segment_starts:
                 self.segment_starts[fqn] = torch.tensor([start for start, _ in self.segments[fqn]], device=device)
@@ -424,11 +445,31 @@ def _diff(self, shards: list[tuple]) -> list:
             segment = torch.searchsorted(starts, position, right=True) - 1
             entry_starts = torch.ones_like(segment, dtype=torch.bool)
             entry_starts[1:] = segment[1:] != segment[:-1]
-            gaps.append(encode_gaps(position - starts[segment], entry_starts).cpu())
-            values.append(changed_values.view(torch.uint8).cpu())
-            counts.append(torch.bincount(segment, minlength=len(starts)))
-            names.extend(name for _, name in self.segments[fqn])
-        return [names, counts, gaps, values, checksum_names, checksums, scales_changed]
+            gaps = encode_gaps(position - starts[segment], entry_starts)
+            counts = torch.bincount(segment, minlength=len(starts))
+            pending.append(
+                (
+                    [name for _, name in self.segments[fqn]],
+                    counts,
+                    gaps,
+                    changed_values.view(torch.uint8),
+                    segment_names,
+                    sums,
+                )
+            )
+            num_pending += len(gaps)
+            if num_pending >= self.PART:
+                flush()
+                num_pending = 0
+        if pending:
+            flush()
+        return bool(scales_changed)
+
+
+def _write_json(path: Path, value) -> None:
+    tmp = path.with_name(f".{path.name}.tmp")
+    tmp.write_text(json.dumps(value))
+    tmp.rename(path)
 
 
 class FileSystemWeightSender(WeightSender):
@@ -518,14 +559,27 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
         if not self.config.sparse:
             self._save_full(model, step_dir)
             return
-        self.wait_for_weight_reads()
-        self._raise_errors()
         self.deferred_finish = False
         if self.world.is_master:
             # Engines that load a step with a delta dir report their weight layout next to it.
             (step_dir / DELTA_DIR).mkdir(exist_ok=True)
-        if self.encoder is None:
-            self.encoder = SparseDeltaEncoder(model, self.config.snapshot)
+        assert self.encoder is not None
+        first = self.last_step is None
+        base_step, self.last_step = self.last_step, step
+        if first:
+            # The snapshot is recorded (``_prepare``) before the first training step.
+            self.wait_for_weight_reads()
+            if step == 0:
+                # The run's first broadcast is the checkpoint the engines booted from.
+                if self.world.is_master:
+                    manifest = {"step": step, "base_step": None, "ranks": 0, "tensors": self.encoder.hf_tensors}
+                    _write_json(step_dir / DELTA_DIR / MANIFEST, manifest)
+                self.anchor_step = step
+                return
+            # A resumed run's first broadcast is a full checkpoint.
+            self._submit(self.encodes, self._publish_full, step, step_dir)
+            self.deferred_finish = True
+            return
         if self.encoder.fp8_layout is None:
             # Until every engine worker has reported which tensors it holds in fp8, steps are full checkpoints.
             layout = [read_engine_layout(step_dir.parent) if self.world.is_master else None]
@@ -533,119 +587,135 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
             if layout[0] is not None:
                 self.encoder.set_fp8_layout(model, layout[0], self.world.is_master)
         shards = self.encoder.local_shards(model, self.world.is_master)
-        first = self.last_step is None
-        base_step, self.last_step = self.last_step, step
-        # The run's first broadcast is the checkpoint the engines booted from; a resumed run's is an anchor.
-        boot = first and step == 0
-        full = (first and step != 0) or (not first and self.encoder.fp8_layout is None)
+        scales_changed = self.encoder.prepare(shards, torch.cuda.max_memory_reserved())
+        if self.encoder.fp8:
+            changed = torch.tensor(int(scales_changed), device=f"cuda:{self.encoder.device}")
+            dist.all_reduce(changed, op=dist.ReduceOp.MAX)
+            scales_changed = bool(changed.item())
+            if scales_changed:
+                self.logger.info(f"fp8 block scales changed at v{step}: publishing a full checkpoint")
+        full = self.encoder.fp8_layout is None or scales_changed
         # Periodic anchors only serve engines that fall off the chain: written in the background, each after
         # the previous one.
-        anchor = not first and not full and step % self.config.anchor_interval == 0
-        self.encoder.prepare(shards, torch.cuda.max_memory_reserved())
+        anchor = not full and step % self.config.anchor_interval == 0
         ready = torch.cuda.Event()
         ready.record()
         self.weights_read = threading.Event()
-        args = (step, base_step, step_dir, shards, ready, self.weights_read)
-        self._submit(self.encodes, self._encode, *args, boot, full, anchor)
-        if boot:
-            if self.world.is_master:
-                manifest = {"step": step, "base_step": None, "files": [], "tensors": self.encoder.hf_tensors}
-                (step_dir / DELTA_DIR / MANIFEST).write_text(json.dumps(manifest))
-            self.anchor_step = step
-            return
+        args = (step, base_step, step_dir, shards, ready, self.weights_read, full, anchor)
+        self._submit(self.encodes, self._encode, *args)
         self.deferred_finish = True
 
+    def _prepare(self, model: nn.Module, step: int) -> None:
+        if self.lora_config is not None or not self.config.sparse:
+            return
+        self.wait_for_weight_reads()
+        self._raise_errors()
+        if self.encoder is None:
+            # The first broadcast only records the snapshot, while the master waits for the receivers.
+            self.encoder = SparseDeltaEncoder(model, self.config.snapshot)
+            shards = self.encoder.local_shards(model, self.world.is_master)
+            ready = torch.cuda.Event()
+            ready.record()
+            self.weights_read = threading.Event()
+            self._submit(self.encodes, self._record, shards, ready, self.weights_read)
+
+    def _record(self, shards, ready, weights_read: threading.Event) -> None:
+        encoder = self.encoder
+        assert encoder is not None
+        start = time.perf_counter()
+        encoder.encode(shards, ready, weights_read, emit=None)
+        num_bytes = sum(numel * encoder.wire_dtypes[key].itemsize for key, (numel, _) in encoder.shards.items())
+        self.logger.info(
+            f"Recorded the {num_bytes / 2**30:.1f} GiB sparse weight snapshot in {time.perf_counter() - start:.1f}s"
+        )
+
     def _wait_for_host_readers(self) -> None:
         for reader in self.host_readers:
             reader.exception()
         self.host_readers = []
 
+    def _anchor_parts(self) -> list:
+        """This rank's slices of the last broadcast's full checkpoint, in host memory."""
+        assert self.encoder is not None
+        if self.encoder.on_gpu:
+            self._wait_for_host_readers()
+            self.encoder.copy_snapshot_to_host()
+        return self.encoder.anchor_parts()
+
+    def _publish_full(self, step: int, step_dir: Path) -> None:
+        self.host_readers.append(self._submit(self.writes, self._write_full, step, step_dir, self._anchor_parts()))
+
     def _encode(
         self,
         step: int,
-        base_step: int | None,
+        base_step: int,
         step_dir: Path,
         shards,
         ready,
         weights_read: threading.Event,
-        boot: bool,
         full: bool,
         anchor: bool,
     ) -> None:
-        """Encode thread: diff v{step}, then hand the delta (or, on full steps, the checkpoint) to a writer."""
+        """Encode thread: diff v{step}, handing the delta to the writer in parts as it is encoded (on full
+        steps, only advance the snapshot and hand the checkpoint to the writer)."""
         encoder = self.encoder
         assert encoder is not None
         if not encoder.on_gpu:
             self._wait_for_host_readers()  # anchors written from the host snapshot, which the encode advances
         start = time.perf_counter()
-        delta = encoder.encode(shards, ready, weights_read)
-        encode_time = time.perf_counter() - start
-        if delta is None:
-            num_bytes = sum(numel * encoder.wire_dtypes[key].itemsize for key, (numel, _) in encoder.shards.items())
-            self.logger.info(f"Recorded the {num_bytes / 2**30:.1f} GiB sparse weight snapshot in {encode_time:.1f}s")
-            if boot:
-                return
-        elif encoder.fp8 and self._scales_changed_anywhere(step_dir, delta[4]):
-            self.logger.info(f"fp8 block scales changed at v{step}: publishing a full checkpoint")
-            full = True
-        if full or anchor:
-            if encoder.on_gpu:
-                self._wait_for_host_readers()
-                encoder.copy_snapshot_to_host()
-            parts = encoder.anchor_parts()
         if full:
-            self.host_readers.append(self._submit(self.writes, self._write_full, step, step_dir, parts))
+            encoder.encode(shards, ready, weights_read, emit=lambda part: None)
+            self._publish_full(step, step_dir)
             return
-        self._submit(self.writes, self._write_delta, step, base_step, step_dir, delta[:4], encode_time)
-        if anchor:
-            self.host_readers.append(self._submit(self.anchors, self._write_anchor, step, step_dir, parts))
+        if self.world.is_master:
+            self._submit(self.writes, self._start_delta, step, base_step, step_dir)
+        num_parts = 0
 
-    def _scales_changed_anywhere(self, step_dir: Path, changed: bool) -> bool:
-        """Agree through the step dir (the encode thread runs no collectives) whether any rank's fp8 block
-        scales changed."""
-        delta_dir = step_dir / DELTA_DIR
-        delta_dir.mkdir(exist_ok=True)
-        tmp = delta_dir / f".rank_{self.world.rank}.scales.tmp"
-        tmp.write_text(str(int(changed)))
-        tmp.rename(delta_dir / f"rank_{self.world.rank}.scales")
-        flags = [delta_dir / f"rank_{rank}.scales" for rank in range(self.world.world_size)]
-        while not all(flag.exists() for flag in flags):
-            time.sleep(0.05)
-        return any(flag.read_text() == "1" for flag in flags)
+        def emit(part: tuple) -> None:
+            nonlocal num_parts
+            self._submit(self.writes, self._write_part, step_dir, num_parts, part)
+            num_parts += 1
 
-    def _write_delta(self, step: int, base_step: int, step_dir: Path, delta: tuple, encode_time: float) -> None:
-        """Write this rank's delta file; the master then waits for every rank's file and commits
-        the step (manifest, then the finished marker). Runs on a thread, without collectives."""
+        encoder.encode(shards, ready, weights_read, emit)
+        self._submit(self.writes, self._finish_delta, step, base_step, step_dir, num_parts, time.perf_counter() - start)
+        if anchor:
+            self.host_readers.append(
+                self._submit(self.anchors, self._write_anchor, step, step_dir, self._anchor_parts())
+            )
+
+    def _start_delta(self, step: int, base_step: int, step_dir: Path) -> None:
+        """The manifest comes first, so engines can read part files as they appear."""
         assert self.encoder is not None
-        start = time.perf_counter()
-        entries, gaps, values, checksums = delta
-        delta_dir = step_dir / DELTA_DIR
-        files = [f"rank_{rank}.safetensors" for rank in range(self.world.world_size)]
-        write_rank_file(
-            delta_dir / files[self.world.rank],
-            entries,
-            gaps.numpy(),
-            values.numpy(),
-            self.encoder.hf_tensors,
-            checksums,
-        )
-        if not self.world.is_master:
-            return
-        while not all((delta_dir / name).exists() for name in files):
-            time.sleep(0.05)
         manifest = {
             "step": step,
             "base_step": base_step,
             "anchor_step": self.anchor_step,
-            "files": files,
+            "ranks": self.world.world_size,
             "tensors": self.encoder.hf_tensors,
         }
-        (delta_dir / MANIFEST).write_text(json.dumps(manifest))
+        _write_json(step_dir / DELTA_DIR / MANIFEST, manifest)
+
+    def _write_part(self, step_dir: Path, index: int, part: tuple) -> None:
+        assert self.encoder is not None
+        entries, gaps, values, checksums = part
+        path = step_dir / DELTA_DIR / rank_part(self.world.rank, index)
+        write_rank_file(path, entries, gaps.numpy(), values.numpy(), self.encoder.hf_tensors, checksums)
+
+    def _finish_delta(self, step: int, base_step: int, step_dir: Path, num_parts: int, encode_time: float) -> None:
+        """After this rank's last part: its parts marker. The master then waits for every rank's marker and
+        commits the step. Runs on the writer thread, without collectives."""
+        delta_dir = step_dir / DELTA_DIR
+        _write_json(delta_dir / rank_parts_marker(self.world.rank), num_parts)
+        if not self.world.is_master:
+            return
+        markers = [delta_dir / rank_parts_marker(rank) for rank in range(self.world.world_size)]
+        while not all(marker.exists() for marker in markers):
+            time.sleep(0.05)
         WeightSender._finish(self, step, step_dir)
-        num_bytes = sum((delta_dir / name).stat().st_size for name in files)
+        num_bytes = sum(path.stat().st_size for path in delta_dir.glob("rank_*.safetensors"))
         self.logger.info(
             f"Sparse weight delta v{step} (base v{base_step}): {num_bytes / 2**20:.1f} MiB, "
-            f"encode {encode_time:.2f}s, write {time.perf_counter() - start:.2f}s (in the background)"
+            f"encode {encode_time:.2f}s (parts written as encoded)"
         )
 
     def _write_full(self, step: int, step_dir: Path, parts: list) -> None:
@@ -688,14 +758,18 @@ class FileSystemWeightReceiver(WeightReceiver):
 
     async def receive(self, step: int) -> None:
         weights_dir = self.step_dir(step)
+        manifest, finished = weights_dir / DELTA_DIR / MANIFEST, weights_dir / FINISHED_MARKER
         self._ack(step)
-        await wait_for_path(weights_dir / FINISHED_MARKER)
-        if (weights_dir / "adapter_config.json").exists():
-            await load_lora_adapter(self.admin_plane, self.model_name, weights_dir)
-            return
-        if (weights_dir / DELTA_DIR / MANIFEST).exists():
-            # Read and decode the delta while the engines keep serving; the pause only applies it.
+        # A sparse delta's manifest comes first and its part files as the trainer encodes: the engines
+        # read and decode them while they keep serving, and the pause only applies the delta.
+        while not (manifest.exists() or finished.exists()):
+            await asyncio.sleep(0.05)
+        if manifest.exists():
             start = time.perf_counter()
             await self.admin_plane.prefetch_weights(weights_dir)
             self.logger.debug(f"Prefetched sparse weight delta v{step} in {time.perf_counter() - start:.2f}s")
+        await wait_for_path(finished)
+        if (weights_dir / "adapter_config.json").exists():
+            await load_lora_adapter(self.admin_plane, self.model_name, weights_dir)
+            return
         await self.admin_plane.update_weights(weights_dir, transport="filesystem", step=step)
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index 34e77866453..9692defe1e6 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -1,8 +1,9 @@
 """Sparse weight deltas for the filesystem weight broadcast.
 
 A delta lives in the ``delta/`` subdir of a broadcast step dir (next to the
-full checkpoint on anchor steps): one ``rank_{r}.safetensors`` per trainer rank
-and a ``manifest.json``, written last. A rank file stores zstd frames over its
+full checkpoint on anchor steps): a ``manifest.json``, written first, and per
+trainer rank part files ``rank_{r}_{p}.safetensors``, written as the encode
+produces them, then a ``rank_{r}.parts`` marker. A part file stores zstd frames over its
 entries ``(hf_name, count)``: gaps between sorted flat indices (restarted per
 entry) and the differences of the new and the previous raw bits (wrapping, in
 the bits dtype; mostly a few ulps), zigzag-encoded. Gaps and differences are
@@ -14,6 +15,7 @@
 import json
 import os
 import re
+import time
 from collections import defaultdict
 from collections.abc import Iterable, Iterator
 from concurrent.futures import ThreadPoolExecutor
@@ -176,7 +178,10 @@ def _rank_file_frames(
     return frames, json.loads(metadata["checksums"])
 
 
-def _read_frame(fd: int, frame, tensors: dict, names: set[str] | None) -> list[tuple[str, np.ndarray, np.ndarray]]:
+def _read_frame(
+    fd: int, frame, tensors: dict, names: set[str] | None
+) -> tuple[list[tuple[str, int]], np.ndarray, np.ndarray]:
+    """A frame's kept ``(entries, gaps, value bytes)``."""
     frame_entries, ((gaps_at, gaps_size), (escapes_at, escapes_size), (values_at, values_size)) = frame
     gaps = np.frombuffer(pyzstd.decompress(os.pread(fd, gaps_size, gaps_at)), np.uint8).astype(np.int32)
     escapes = np.frombuffer(pyzstd.decompress(os.pread(fd, escapes_size, escapes_at)), np.int32)
@@ -186,14 +191,18 @@ def _read_frame(fd: int, frame, tensors: dict, names: set[str] | None) -> list[t
     values = np.frombuffer(pyzstd.decompress(os.pread(fd, values_size, values_at)), np.uint8).astype(np.int64)
     values[values == 255] = escapes[num_gap_escapes:].astype(np.int64) & 0xFFFFFFFF
     values = _pack_values((values >> 1) ^ -(values & 1), frame_entries, tensors)
-    out, gap_offset, value_offset = [], 0, 0
+    if names is None or all(name in names for name, _ in frame_entries):
+        return frame_entries, gaps, values
+    kept, kept_gaps, kept_values, gap_offset, value_offset = [], [], [], 0, 0
     for name, count in frame_entries:
         num_bytes = count * bits_dtype(tensors[name][1]).itemsize
-        if names is None or name in names:
-            out.append((name, gaps[gap_offset : gap_offset + count], values[value_offset : value_offset + num_bytes]))
+        if name in names:
+            kept.append((name, count))
+            kept_gaps.append(gaps[gap_offset : gap_offset + count])
+            kept_values.append(values[value_offset : value_offset + num_bytes])
         gap_offset += count
         value_offset += num_bytes
-    return out
+    return kept, np.concatenate(kept_gaps), np.concatenate(kept_values)
 
 
 def _write_atomically(path: Path, data: bytes) -> None:
@@ -308,32 +317,68 @@ def read_manifest(step_dir: Path) -> dict:
     return json.loads((step_dir / DELTA_DIR / MANIFEST).read_text())
 
 
+def rank_part(rank: int, part: int) -> str:
+    return f"rank_{rank}_{part}.safetensors"
+
+
+def rank_parts_marker(rank: int) -> str:
+    """Written by a trainer rank after its last part file: the number of part files."""
+    return f"rank_{rank}.parts"
+
+
 def read_delta(
-    step_dir: Path, names: set[str] | None = None, num_threads: int = 16
-) -> tuple[list[tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]], dict[str, int]]:
-    """Read and decompress every rank file of a delta step dir into ``(entries,
-    gaps, value bytes)`` in host memory (values: new minus previous raw bits, in each entry's bits
-    dtype); ``decode_gaps`` turns the
-    gaps into flat indices on the device that applies them. With ``names``, keep
-    only those tensors (frames without any are not read).
-    Also returns the new weights' ``bits_checksum`` per HF tensor, summed over ranks."""
+    step_dir: Path, names: set[str] | None = None, num_threads: int = 16, cpus: set[int] | None = None
+) -> tuple[list[tuple[list[tuple[str, int]], np.ndarray, np.ndarray]], dict[str, int]]:
+    """Read and decompress a delta step dir into ``(entries, gaps, value bytes)`` per frame in host
+    memory (values: new minus previous raw bits, in each entry's bits dtype); ``decode_gaps`` turns
+    the gaps into flat indices on the device that applies them. With ``names``, keep only those
+    tensors (frames without any are not read). Part files are read as the trainer writes them, until
+    every rank has written its parts marker. ``cpus``: CPUs for the reader threads (the device's NUMA
+    node). Also returns the new weights' ``bits_checksum`` per HF tensor, summed over ranks."""
     manifest = read_manifest(step_dir)
     tensors = manifest["tensors"]
+    delta_dir = step_dir / DELTA_DIR
     checksums: dict[str, int] = {}
-    with ExitStack() as stack, ThreadPoolExecutor(num_threads) as pool:
-        files = [stack.enter_context(open(step_dir / DELTA_DIR / name, "rb")) for name in manifest["files"]]
-        jobs = []
-        for file in files:
-            frames, partial = _rank_file_frames(file, tensors, names)
-            jobs.extend((file.fileno(), frame) for frame in frames)
-            for tensor_name, value in partial.items():
-                checksums[tensor_name] = add_checksums(checksums.get(tensor_name, 0), value)
-        kept = [entry for frame in pool.map(lambda job: _read_frame(*job, tensors, names), jobs) for entry in frame]
+    num_parts: list[int | None] = [None] * manifest["ranks"]
+    next_part = [0] * manifest["ranks"]
+    initializer = None if cpus is None else (lambda: os.sched_setaffinity(0, cpus))
+    with ExitStack() as stack, ThreadPoolExecutor(num_threads, initializer=initializer) as pool:
+        files, jobs = [], []
+        while any(count is None or next_part[rank] < count for rank, count in enumerate(num_parts)):
+            found = False
+            for rank in range(len(num_parts)):
+                if num_parts[rank] is None and (delta_dir / rank_parts_marker(rank)).exists():
+                    num_parts[rank] = int((delta_dir / rank_parts_marker(rank)).read_text())
+                path = delta_dir / rank_part(rank, next_part[rank])
+                if (num_parts[rank] is None or next_part[rank] < num_parts[rank]) and path.exists():
+                    file = stack.enter_context(open(path, "rb"))
+                    files.append(file)
+                    frames, partial = _rank_file_frames(file, tensors, names)
+                    jobs.extend(pool.submit(_read_frame, file.fileno(), frame, tensors, names) for frame in frames)
+                    for tensor_name, value in partial.items():
+                        checksums[tensor_name] = add_checksums(checksums.get(tensor_name, 0), value)
+                    next_part[rank] += 1
+                    found = True
+            if not found:
+                time.sleep(0.02)
+        frames = [frame for frame in (job.result() for job in jobs) if frame[0]]
         for file in files:
             drop_page_cache(file.fileno())
-    gaps = torch.from_numpy(np.concatenate([g for _, g, _ in kept]) if kept else np.empty(0, np.int32))
-    values = torch.from_numpy(np.concatenate([v for _, _, v in kept]) if kept else np.empty(0, np.uint8))
-    return [([(name, len(g)) for name, g, _ in kept], gaps, values)], checksums
+    return frames, checksums
+
+
+def numa_cpus(device: int) -> set[int] | None:
+    """The CPUs of the NUMA node the CUDA ``device`` hangs off, or None if unknown."""
+    props = torch.cuda.get_device_properties(device)
+    sysfs = Path(f"/sys/bus/pci/devices/{props.pci_domain_id:04x}:{props.pci_bus_id:02x}:{props.pci_device_id:02x}.0")
+    node = int((sysfs / "numa_node").read_text()) if (sysfs / "numa_node").exists() else -1
+    if node < 0:
+        return None
+    cpus = set()
+    for part in Path(f"/sys/devices/system/node/node{node}/cpulist").read_text().strip().split(","):
+        first, _, last = part.partition("-")
+        cpus.update(range(int(first), int(last or first) + 1))
+    return cpus
 
 
 def pin(num_bytes: int) -> torch.Tensor:
diff --git a/tests/unit/utils/test_sparse_weights.py b/tests/unit/utils/test_sparse_weights.py
index 824186d6d46..8ac79ce677d 100644
--- a/tests/unit/utils/test_sparse_weights.py
+++ b/tests/unit/utils/test_sparse_weights.py
@@ -9,6 +9,8 @@
     bits_checksum,
     decode_gaps,
     encode_gaps,
+    rank_part,
+    rank_parts_marker,
     read_delta,
     write_rank_file,
     zigzag,
@@ -40,15 +42,16 @@ def test_sparse_delta_roundtrip_is_exact(tmp_path):
             values.append(zigzag(delta[changed]).view(torch.uint8))
         gaps, values = torch.cat(gaps).numpy(), torch.cat(values).numpy()
         # Small frames, so a frame boundary falls inside rank 1's entries
-        write_rank_file(delta_dir / f"rank_{rank}.safetensors", entries, gaps, values, tensors, checksums, frame_size=8)
-    files = ["rank_0.safetensors", "rank_1.safetensors"]
-    (delta_dir / MANIFEST).write_text(json.dumps({"files": files, "tensors": tensors}))
+        write_rank_file(delta_dir / rank_part(rank, 0), entries, gaps, values, tensors, checksums, frame_size=8)
+        (delta_dir / rank_parts_marker(rank)).write_text("1")
+    (delta_dir / MANIFEST).write_text(json.dumps({"ranks": 2, "tensors": tensors}))
 
     def apply(names):
         patched = {name: as_bits(tensor.clone()).view(-1) for name, tensor in previous.items()}
         delta, checksums = read_delta(tmp_path, names)
         assert checksums == {name: bits_checksum(as_bits(tensor)) for name, tensor in current.items()}
         for entries, gaps, values in delta:
+            gaps, values = torch.from_numpy(gaps), torch.from_numpy(values)
             indices = decode_gaps(gaps, torch.tensor([count for _, count in entries]))
             i = b = 0
             for name, count in entries:

From d3a9e95c93847e9af3281c2e67d8b65748e6e946 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Thu, 15 Oct 2026 19:41:13 +0000
Subject: [PATCH 27/88] perf(weights): copy sparse delta parts to the host
 through a pinned staging buffer

---
 src/prime_rl/transports/weights/filesystem.py | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index f9716bde229..5e5e40783ee 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -262,6 +262,7 @@ def _fill(self, shards: list[tuple]) -> None:
         (else placed at the first delta, against the measured training peak). The host buffer is filled in
         every mode, so its pages are resident before training: faulting them in later stalls the trainer."""
         self.shards = {key: (shard.numel(), offset) for key, _, shard, offset, _ in shards}
+        self.staging = torch.empty(1 << 28, dtype=torch.uint8, pin_memory=True)
         self._allocate_host()
         if self.snapshot_placement == "cpu":
             self._page_lock()
@@ -343,6 +344,18 @@ def encode(
             weights_read.set()
         self.previous = self.snapshot
 
+    def _to_host(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Copy a delta part to pageable host memory through the pinned staging buffer: a copy straight
+        into pageable memory holds the CUDA driver and stalls the training step's own CUDA calls."""
+        data, host = tensor.view(torch.uint8), torch.empty(tensor.shape, dtype=tensor.dtype)
+        host_bytes = host.view(torch.uint8)
+        for begin in range(0, data.numel(), self.staging.numel()):
+            chunk = self.staging[: min(self.staging.numel(), data.numel() - begin)]
+            chunk.copy_(data[begin : begin + chunk.numel()], non_blocking=True)
+            torch.cuda.current_stream().synchronize()
+            host_bytes[begin : begin + chunk.numel()].copy_(chunk)
+        return host
+
     def _segments(self, key: str, fqn: str, offset: int, numel: int) -> tuple[list[str], list[int], list[int]]:
         """The HF tensor parts in this rank's shard of ``key``: names, and begin/end in the shard."""
         if key not in self.shard_segments:
@@ -380,8 +393,8 @@ def _diff(self, shards: list[tuple], emit: Callable) -> bool:
         def flush() -> None:
             names = [name for part in pending for name in part[0]]
             counts = torch.cat([part[1] for part in pending]).tolist()
-            gaps = torch.cat([part[2] for part in pending]).cpu()
-            values = torch.cat([part[3] for part in pending]).cpu()
+            gaps = self._to_host(torch.cat([part[2] for part in pending]))
+            values = self._to_host(torch.cat([part[3] for part in pending]))
             checksum_names = [name for part in pending for name in part[4]]
             checksums = dict(zip(checksum_names, torch.cat([part[5] for part in pending]).tolist()))
             emit(([(name, count) for name, count in zip(names, counts) if count], gaps, values, checksums))

From 253f1a5aba64970f899b64f093859d3439d284ef Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Thu, 15 Oct 2026 22:38:45 +0000
Subject: [PATCH 28/88] perf(weights): diff fp8-held tensors on the encode
 thread with an ahead-of-time compiled quantizer

vLLM's per_block_cast_to_fp8 is compiled once per run with AOTInductor (bit-identical to its torch.compile'd version for columns a multiple of 128), which, unlike the torch.compile'd function, can run next to the compiled, checkpointed training step. Tensors with other column counts stay on the training thread. Ranks agree on fp8 block-scale changes through per-rank flag files from the encode thread; a step's parts are held until then.
---
 src/prime_rl/transports/weights/filesystem.py | 76 ++++++++++++++-----
 src/prime_rl/utils/sparse_weights.py          | 50 +++++++++++-
 2 files changed, 101 insertions(+), 25 deletions(-)

diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index 5e5e40783ee..a475d7728b8 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -30,6 +30,7 @@
     anchor_complete,
     as_bits,
     bits_dtype,
+    compile_fp8_quantizer,
     encode_gaps,
     quantize_fp8_blocks,
     rank_part,
@@ -107,6 +108,9 @@ def __init__(self, model: nn.Module, snapshot: str):
         # fqn -> (ue8m0 scales, columns) for trainer tensors held in fp8; quantized as one [rows, columns] matrix
         self.fp8: dict[str, tuple[bool, int]] = {}
         self.fp8_layout: dict[str, str] | None = None
+        self.fp8_background: set[str] = set()
+        self.compiled_fp8 = None
+        self.scales_changed = False
         self.gather: set[str] = set()
         self.segment_starts: dict[str, torch.Tensor] = {}
         self.shard_segments: dict[str, tuple[list[str], list[int], list[int]]] = {}
@@ -148,6 +152,12 @@ def set_fp8_layout(self, model: nn.Module, fp8: dict[str, str], is_master: bool)
             for _, name in segs:
                 self.hf_tensors[name] = (self.hf_tensors[name][0], "float8_e4m3fn")
         self.gather = self._find_gathered(model, is_master)
+        # Diffed on the encode thread with the ahead-of-time compiled quantizer; others on the training thread
+        self.fp8_background = {fqn for fqn, (_, num_cols) in self.fp8.items() if num_cols % FP8_BLOCK[1] == 0}
+        if self.fp8_background:
+            start = time.perf_counter()
+            self.compiled_fp8 = compile_fp8_quantizer()
+            get_logger().info(f"Compiled the fp8 quantizer ahead of time in {time.perf_counter() - start:.1f}s")
 
     def _local_shard(self, key: str, value: torch.Tensor, is_master: bool) -> tuple[torch.Tensor, int] | None:
         """This rank's contiguous flat slice of ``value`` and its global flat offset,
@@ -299,15 +309,17 @@ def anchor_parts(self) -> list[tuple[str, list[int], int, torch.Tensor]]:
         return parts
 
     @torch.no_grad()
-    def prepare(self, shards: list[tuple], peak_reserved: int) -> bool:
+    def prepare(self, shards: list[tuple], peak_reserved: int) -> None:
         """The training thread's part of an encode, before ``encode``: place the snapshot at the first delta,
-        and diff the tensors held in fp8 (vLLM's quantizer is torch.compile'd, and running it next to the
-        training step breaks activation checkpointing there). Returns whether an fp8 block scale changed
-        on this rank."""
+        and diff the tensors held in fp8 that the ahead-of-time compiled quantizer does not cover (vLLM's
+        quantizer is torch.compile'd, and running it next to the training step breaks activation
+        checkpointing there). ``scales_changed``: whether an fp8 block scale changed on this rank, after
+        ``encode``."""
         self.previous = self.snapshot
         self.fp8_parts = []
+        self.scales_changed = False
         if not self.initialized:
-            return False
+            return
         if self.on_gpu is None:
             self._place(peak_reserved)
             if self.on_gpu:
@@ -315,8 +327,9 @@ def prepare(self, shards: list[tuple], peak_reserved: int) -> bool:
                 self.snapshot = {key: torch.empty_like(bits, device=device) for key, bits in self.host.items()}
             elif not self.pinned:
                 self._page_lock()
-        fp8 = [shard for shard in shards if shard[1] in self.fp8]
-        return self._diff(fp8, self.fp8_parts.append) if fp8 else False
+        here = [shard for shard in shards if shard[1] in self.fp8 and shard[1] not in self.fp8_background]
+        if here:
+            self.scales_changed = self._diff(here, self.fp8_parts.append)
 
     @torch.no_grad()
     def encode(
@@ -338,7 +351,8 @@ def encode(
                     return
                 for part in self.fp8_parts:
                     emit(part)
-                self._diff([shard for shard in shards if shard[1] not in self.fp8], emit)
+                rest = [shard for shard in shards if shard[1] not in self.fp8 or shard[1] in self.fp8_background]
+                self.scales_changed |= self._diff(rest, emit)
                 self.stream.synchronize()
         finally:
             weights_read.set()
@@ -427,7 +441,11 @@ def flush() -> None:
                             current.view(torch.bfloat16),
                         )
                     (previous_fp8, previous_scales), (current_fp8, current_scales) = (
-                        quantize_fp8_blocks(weights.reshape(-1, num_cols), ue8m0)
+                        quantize_fp8_blocks(
+                            weights.reshape(-1, num_cols),
+                            ue8m0,
+                            self.compiled_fp8 if fqn in self.fp8_background else None,
+                        )
                         for weights in (previous_weights, current_weights)
                     )
                     scales_changed |= not torch.equal(previous_scales, current_scales)
@@ -600,14 +618,8 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
             if layout[0] is not None:
                 self.encoder.set_fp8_layout(model, layout[0], self.world.is_master)
         shards = self.encoder.local_shards(model, self.world.is_master)
-        scales_changed = self.encoder.prepare(shards, torch.cuda.max_memory_reserved())
-        if self.encoder.fp8:
-            changed = torch.tensor(int(scales_changed), device=f"cuda:{self.encoder.device}")
-            dist.all_reduce(changed, op=dist.ReduceOp.MAX)
-            scales_changed = bool(changed.item())
-            if scales_changed:
-                self.logger.info(f"fp8 block scales changed at v{step}: publishing a full checkpoint")
-        full = self.encoder.fp8_layout is None or scales_changed
+        self.encoder.prepare(shards, torch.cuda.max_memory_reserved())
+        full = self.encoder.fp8_layout is None
         # Periodic anchors only serve engines that fall off the chain: written in the background, each after
         # the previous one.
         anchor = not full and step % self.config.anchor_interval == 0
@@ -670,7 +682,8 @@ def _encode(
         anchor: bool,
     ) -> None:
         """Encode thread: diff v{step}, handing the delta to the writer in parts as it is encoded (on full
-        steps, only advance the snapshot and hand the checkpoint to the writer)."""
+        steps, only advance the snapshot and hand the checkpoint to the writer). With fp8 deltas, the parts
+        are held until the ranks agree that no block scale changed, else the step is a full checkpoint."""
         encoder = self.encoder
         assert encoder is not None
         if not encoder.on_gpu:
@@ -680,22 +693,43 @@ def _encode(
             encoder.encode(shards, ready, weights_read, emit=lambda part: None)
             self._publish_full(step, step_dir)
             return
-        if self.world.is_master:
-            self._submit(self.writes, self._start_delta, step, base_step, step_dir)
+        held: list[tuple] = []
         num_parts = 0
 
         def emit(part: tuple) -> None:
             nonlocal num_parts
+            if num_parts == 0 and self.world.is_master:
+                self._submit(self.writes, self._start_delta, step, base_step, step_dir)
             self._submit(self.writes, self._write_part, step_dir, num_parts, part)
             num_parts += 1
 
-        encoder.encode(shards, ready, weights_read, emit)
+        encoder.encode(shards, ready, weights_read, held.append if encoder.fp8 else emit)
+        if encoder.fp8:
+            if self._scales_changed_anywhere(step_dir, encoder.scales_changed):
+                self.logger.info(f"fp8 block scales changed at v{step}: publishing a full checkpoint")
+                self._publish_full(step, step_dir)
+                return
+            for part in held:
+                emit(part)
+        if num_parts == 0 and self.world.is_master:
+            self._submit(self.writes, self._start_delta, step, base_step, step_dir)
         self._submit(self.writes, self._finish_delta, step, base_step, step_dir, num_parts, time.perf_counter() - start)
         if anchor:
             self.host_readers.append(
                 self._submit(self.anchors, self._write_anchor, step, step_dir, self._anchor_parts())
             )
 
+    def _scales_changed_anywhere(self, step_dir: Path, changed: bool) -> bool:
+        """Agree through the step dir whether any rank's fp8 block scales changed: the encode thread runs no
+        collectives, and the training thread may be waiting for a batch that needs this step."""
+        delta_dir = step_dir / DELTA_DIR
+        delta_dir.mkdir(exist_ok=True)
+        _write_json(delta_dir / f"rank_{self.world.rank}.scales", int(changed))
+        flags = [delta_dir / f"rank_{rank}.scales" for rank in range(self.world.world_size)]
+        while not all(flag.exists() for flag in flags):
+            time.sleep(0.05)
+        return any(flag.read_text() == "1" for flag in flags)
+
     def _start_delta(self, step: int, base_step: int, step_dir: Path) -> None:
         """The manifest comes first, so engines can read part files as they appear."""
         assert self.encoder is not None
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index 9692defe1e6..f80bcdcfeca 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -49,19 +49,61 @@ def bits_dtype(dtype: str) -> torch.dtype:
     return _BITS_DTYPE[getattr(torch, dtype).itemsize]
 
 
-def quantize_fp8_blocks(weight: torch.Tensor, ue8m0: bool) -> tuple[torch.Tensor, torch.Tensor]:
+def quantize_fp8_blocks(weight: torch.Tensor, ue8m0: bool, compiled=None) -> tuple[torch.Tensor, torch.Tensor]:
     """vLLM's online ``fp8_per_block`` quantization of a 2-D weight: fp8 bits (uint8) and fp32 block
     scales. With ``ue8m0``, also DeepGEMM's requantization to power-of-two scales. Uses vLLM's own
-    compiled kernels: a reimplementation rounds differently on a few elements."""
+    compiled kernels: a reimplementation rounds differently on a few elements. ``compiled``: the
+    ahead-of-time compiled kernels from ``compile_fp8_quantizer`` (columns a multiple of 128), which,
+    unlike vLLM's torch.compile'd function, may run on a thread next to a compiled training step."""
     from vllm.model_executor.layers.quantization.utils.fp8_utils import requant_weight_ue8m0_inplace
     from vllm.utils.deep_gemm import per_block_cast_to_fp8
 
-    quantized, scales = per_block_cast_to_fp8(weight, FP8_BLOCK, use_ue8m0=False)
+    if compiled is None:
+        quantized, scales = per_block_cast_to_fp8(weight, FP8_BLOCK, use_ue8m0=False)
+        if ue8m0:
+            requant_weight_ue8m0_inplace(quantized, scales, FP8_BLOCK)
+        return quantized.view(torch.uint8), scales
+    cast, cast_ue8m0 = compiled
+    quantized, scales = cast(weight)
     if ue8m0:
-        requant_weight_ue8m0_inplace(quantized, scales, FP8_BLOCK)
+        # requant_weight_ue8m0_inplace, step by step
+        rows, cols = quantized.shape
+        expanded = scales.repeat_interleave(FP8_BLOCK[0], 0).repeat_interleave(FP8_BLOCK[1], 1)[:rows, :cols]
+        quantized, scales = cast_ue8m0(quantized.to(torch.float32) * expanded)
     return quantized.view(torch.uint8), scales
 
 
+def compile_fp8_quantizer() -> tuple:
+    """vLLM's ``per_block_cast_to_fp8`` compiled ahead of time with AOTInductor for any shape (bf16 input;
+    and fp32 input with ue8m0 scales, for the requantization). Matches vLLM's compiled function bit for
+    bit for columns a multiple of 128. Compile on the thread that trains: compiling next to a compiled
+    training step breaks activation checkpointing there."""
+    import tempfile
+
+    from torch.export import Dim
+    from vllm.utils.deep_gemm import per_block_cast_to_fp8
+
+    cast = per_block_cast_to_fp8._torchdynamo_orig_callable
+
+    class Cast(torch.nn.Module):
+        def __init__(self, ue8m0: bool):
+            super().__init__()
+            self.ue8m0 = ue8m0
+
+        def forward(self, weight: torch.Tensor):
+            return cast(weight, FP8_BLOCK, use_ue8m0=self.ue8m0)
+
+    compiled = []
+    for ue8m0, dtype in ((False, torch.bfloat16), (True, torch.float32)):
+        example = torch.zeros(384, 640, dtype=dtype, device="cuda")
+        program = torch.export.export(Cast(ue8m0), (example,), dynamic_shapes={"weight": {0: Dim.AUTO, 1: Dim.AUTO}})
+        path = str(Path(tempfile.mkdtemp()) / "quantize.pt2")
+        compiled.append(
+            torch._inductor.aoti_load_package(torch._inductor.aoti_compile_and_package(program, package_path=path))
+        )
+    return tuple(compiled)
+
+
 def zigzag(delta: torch.Tensor) -> torch.Tensor:
     """Raw-bit differences (bits dtype, wrapping) -> unsigned values of the same width, small for
     small steps of either sign: 0, -1, 1, -2 -> 0, 1, 2, 3."""

From fd11f1582712a19cdebcfda8ce1e1c4d1ed2f0eb Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Fri, 16 Oct 2026 11:41:16 +0000
Subject: [PATCH 29/88] fix(weights): engine fp8 layout only from checkpoint
 weights; position-weighted checksums

- An engine reports its fp8 layout while its weights equal a checkpoint (the boot checkpoint at v0, a loaded full checkpoint, or the anchor a replay starts from). A fresh fp8 engine joining at a non-anchor step compared the boot checkpoint against replayed weights and failed.
- Per-step checksums weight each element by its checkpoint flat index + 1, so a change written at a wrong position inside a tensor is caught every step. The full anchor reload that checked the chain during the pause is gone.
- One page-lock helper.
---
 .../inference/vllm/worker/filesystem.py       | 51 +++++++------
 .../inference/vllm/worker/weight_transfer.py  | 67 +++++++++--------
 src/prime_rl/transports/weights/filesystem.py | 72 +++++++++++--------
 src/prime_rl/utils/sparse_weights.py          | 21 ++++--
 tests/unit/utils/test_sparse_weights.py       |  2 +-
 5 files changed, 118 insertions(+), 95 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 415fba31270..833671b1aad 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -18,7 +18,6 @@
     CheckpointIndexMap,
     load_weights_checkpoint_layerwise,
     update_mla_absorbed_weights,
-    weight_checksums,
 )
 from prime_rl.utils.sparse_weights import (
     DELTA_DIR,
@@ -154,29 +153,21 @@ def update_weights_from_path(self, weight_path: str) -> None:
         step = int(step_dir.name.removeprefix("step_"))
         if not (step_dir / DELTA_DIR / MANIFEST).exists():
             self._load_checkpoint(step_dir)
+            if (step_dir / DELTA_DIR).is_dir():
+                # A sparse run's full checkpoint
+                self._write_engine_layout(step_dir, step_dir)
         else:
             base_step = read_manifest(step_dir)["base_step"]
             if base_step is None:
                 # The run's first broadcast: the weights the engine booted with.
                 if self.weight_step is not None:
                     self._load_checkpoint(self._boot_dir())
+                self._write_engine_layout(step_dir, self._boot_dir())
             elif base_step == self.weight_step:
                 self._apply_delta(step_dir)
-                if anchor_complete(step_dir):
-                    # Anchor step: check that the delta chain reproduced the full checkpoint bit for bit.
-                    chained = weight_checksums(self._model())
-                    self._load_checkpoint(step_dir)
-                    current = weight_checksums(self._model())
-                    mismatched = [name for name, value in current.items() if chained[name] != value]
-                    if mismatched:
-                        raise RuntimeError(f"Sparse weight deltas diverged from the v{step} anchor: {mismatched[:8]}")
-                    logger.info(f"Sparse weight delta chain matches the v{step} anchor")
             else:
                 self._replay(step_dir)
         self.weight_step = step
-        if (step_dir / DELTA_DIR).is_dir() and not self.wrote_layout:
-            self._write_engine_layout(step_dir)
-            self.wrote_layout = True
 
     def _boot_dir(self) -> Path:
         """The local dir of the checkpoint the engine booted from."""
@@ -206,21 +197,25 @@ def _apply_chunk(
             offset += num_changes
             num_bytes += num_changes * itemsize
 
-    def _write_engine_layout(self, step_dir: Path) -> None:
+    def _write_engine_layout(self, step_dir: Path, checkpoint_dir: Path) -> None:
         """Tell the trainer which checkpoint tensors this engine holds in fp8 (and with which
-        scales), so it can send their deltas in the fp8 domain. Reads one checkpoint tensor per fp8
-        param from the step's full checkpoint, or from the boot checkpoint."""
-        if anchor_complete(step_dir):
-            index = anchor_index(step_dir)
+        scales), so it can send their deltas in the fp8 domain. Once per engine, while its weights
+        are those of ``checkpoint_dir`` (an anchor or the boot checkpoint): it reads one checkpoint
+        tensor per fp8 param from there."""
+        if self.wrote_layout:
+            return
+        self.wrote_layout = True
+        if anchor_complete(checkpoint_dir):
+            index = anchor_index(checkpoint_dir)
             tensors = index[1]
 
             def read_tensor(name: str) -> torch.Tensor:
-                return next(anchor_tensors(step_dir, [name], index))[1]
+                return next(anchor_tensors(checkpoint_dir, [name], index))[1]
 
         else:
             files = {}
             tensors = {}
-            for path in self._boot_dir().glob("*.safetensors"):
+            for path in checkpoint_dir.glob("*.safetensors"):
                 with safe_open(path, framework="pt") as f:
                     for name in f.keys():
                         tensor_slice = f.get_slice(name)
@@ -252,15 +247,17 @@ def read_tensor(name: str) -> torch.Tensor:
     def _replay(self, step_dir: Path) -> None:
         """Load the latest complete anchor at or before ``step_dir`` (or the boot checkpoint, for a
         chain from the run's first broadcast) and replay the deltas since."""
-        chain = []
-        while not anchor_complete(step_dir):
-            base_step = read_manifest(step_dir)["base_step"]
+        chain, source = [], step_dir
+        while not anchor_complete(source):
+            base_step = read_manifest(source)["base_step"]
             if base_step is None:
                 break
-            chain.append(step_dir)
-            step_dir = step_dir.parent / f"step_{base_step}"
-        logger.info(f"Replaying {len(chain)} sparse weight deltas on top of {step_dir.name}")
-        self._load_checkpoint(step_dir if anchor_complete(step_dir) else self._boot_dir())
+            chain.append(source)
+            source = source.parent / f"step_{base_step}"
+        logger.info(f"Replaying {len(chain)} sparse weight deltas on top of {source.name}")
+        source = source if anchor_complete(source) else self._boot_dir()
+        self._load_checkpoint(source)
+        self._write_engine_layout(step_dir, source)
         for delta_dir in reversed(chain):
             self._apply_delta(delta_dir)
 
diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py
index 97535c5fa7f..dc109bfabbe 100644
--- a/src/prime_rl/inference/vllm/worker/weight_transfer.py
+++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py
@@ -124,42 +124,61 @@ def placeholders():
         self.names = set(self.records) | unpatchable
 
     def held_checksums(self) -> tuple[list[str], torch.Tensor, torch.Tensor]:
-        """``bits_checksum`` of the live param elements each checkpoint tensor was written to: the
-        names, their sums, and a mask of the ones this worker holds only part of (TP slices)."""
+        """``bits_checksum`` of the live param elements each checkpoint tensor was written to, weighted by
+        their checkpoint positions: the names, their sums, and a mask of the ones this worker holds only
+        part of (TP slices)."""
         if self.checksum_plan is None:
             self.checksum_plan = self._plan_checksums()
         names, partial, row_ranges, other_writes = self.checksum_plan
         totals = torch.zeros(len(names), dtype=torch.int64, device=partial.device)
-        for index, (starts, ends, ids) in row_ranges.items():
-            param = self.params[index]
-            row_sums = as_bits(param.data).reshape(-1, param.shape[-1]).sum(dim=1, dtype=torch.int64)
-            cumulative = torch.cat([row_sums.new_zeros(1), row_sums.cumsum(0)])
-            totals.index_add_(0, ids, cumulative[ends] - cumulative[starts])
-        for name_id, rows, cols, index, offset, *strides in other_writes:
+        for index, (starts, ends, ids, num_cols, firsts) in row_ranges.items():
+            bits = as_bits(self.params[index].data)
+            bits = bits.reshape(-1, bits.shape[-1])
+            row_sums = bits.sum(dim=1, dtype=torch.int64)
+            # Σ_c c·bits[row, c], in blocks of rows that bound the int64 temporary
+            columns = torch.arange(bits.shape[1], device=bits.device)
+            moments = torch.cat([(block * columns).sum(1) for block in bits.split(max(1, (1 << 24) // bits.shape[1]))])
+            rows = torch.arange(bits.shape[0], device=bits.device)
+            plain, by_row, by_col = (
+                cumulative[ends] - cumulative[starts]
+                for cumulative in (
+                    torch.cat([values.new_zeros(1), values.cumsum(0)])
+                    for values in (row_sums, rows * row_sums, moments)
+                )
+            )
+            totals.index_add_(0, ids, num_cols * by_row + firsts * plain + by_col)
+        for name_id, num_cols, first, rows, cols, index, offset, *strides in other_writes:
             view = as_bits(self.params[index].data).view(-1).as_strided((rows, cols), strides, offset)
-            totals[name_id] += view.sum(dtype=torch.int64)
+            row_weights = first + num_cols * torch.arange(rows, device=view.device)
+            totals[name_id] += (row_weights * view.sum(1, dtype=torch.int64)).sum()
+            totals[name_id] += (torch.arange(cols, device=view.device) * view.sum(0, dtype=torch.int64)).sum()
         return names, totals, partial
 
     def _plan_checksums(self):
-        """Writes that cover whole rows of a param become row ranges, summed with one cumsum per param."""
+        """Writes that cover whole rows of a param become row ranges, summed with cumsums over the rows
+        of each param. Element ``(row, col)`` of a write is weighted by its checkpoint flat index + 1:
+        ``row * num_cols + first + col`` for a row range, with ``first`` per write."""
         names = [
             name
             for name, writes in self.records.items()
             if name not in self.unpatchable and sum(rows * cols for _, rows, _, cols, *_ in writes) <= self.numel[name]
         ]
         device = self.params[0].device
-        ranges: dict[int, list[list[int]]] = defaultdict(lambda: [[], [], []])
+        ranges: dict[int, list[list[int]]] = defaultdict(lambda: [[], [], [], [], []])
         other_writes = []
         for name_id, name in enumerate(names):
-            for _, rows, _, cols, index, offset, row_stride, col_stride in self.records[name]:
+            num_cols = self.num_cols[name]
+            for row0, rows, col0, cols, index, offset, row_stride, col_stride in self.records[name]:
                 width = self.params[index].shape[-1]
                 if col_stride == 1 and row_stride == width and cols == width and offset % width == 0:
-                    starts, ends, ids = ranges[index]
-                    starts.append(offset // width)
-                    ends.append(offset // width + rows)
-                    ids.append(name_id)
+                    start = offset // width
+                    for values, value in zip(
+                        ranges[index], (start, start + rows, name_id, num_cols, (row0 - start) * num_cols + col0 + 1)
+                    ):
+                        values.append(value)
                 else:
-                    other_writes.append((name_id, rows, cols, index, offset, row_stride, col_stride))
+                    first = row0 * num_cols + col0 + 1
+                    other_writes.append((name_id, num_cols, first, rows, cols, index, offset, row_stride, col_stride))
         row_ranges = {
             index: tuple(torch.tensor(values, device=device) for values in lists) for index, lists in ranges.items()
         }
@@ -291,20 +310,6 @@ def _restore_kernel_format(swapped: list[tuple]) -> None:
             layer.register_buffer(name, buffer, persistent=name not in non_persistent)
 
 
-@torch.no_grad()
-def weight_checksums(model: Module, chunk: int = 1 << 24) -> dict[str, int]:
-    """Position-weighted sum of every parameter's raw bits (exact, any dtype)."""
-    sums = {}
-    for name, param in model.named_parameters():
-        bits = as_bits(param.data).reshape(-1)
-        total = torch.zeros((), dtype=torch.int64, device=bits.device)
-        for start in range(0, bits.numel(), chunk):
-            part = bits[start : start + chunk].long()
-            total += (part * torch.arange(start + 1, start + 1 + part.numel(), device=bits.device)).sum()
-        sums[name] = total
-    return {name: int(total) for name, total in sums.items()}
-
-
 @torch.no_grad()
 def update_mla_absorbed_weights(model: Module) -> None:
     """Recompute MLA absorbed KV weights after in-place kv_b_proj updates."""
diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index a475d7728b8..00a6d855cce 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -32,6 +32,7 @@
     bits_dtype,
     compile_fp8_quantizer,
     encode_gaps,
+    page_lock,
     quantize_fp8_blocks,
     rank_part,
     rank_parts_marker,
@@ -49,6 +50,9 @@
     save_state_dict_parallel,
 )
 
+# Elements per block of the position-weighted checksums (int64 temporaries of 8 B each)
+SUM_BLOCK = 1 << 22
+
 
 class SparseDeltaEncoder:
     """Rank-local deltas of the trainer's weight shards, in HF coordinates.
@@ -241,11 +245,10 @@ def _allocate_host(self) -> None:
             offset += sizes[key]
 
     def _page_lock(self) -> None:
-        """Page-lock the host buffer with one ``cudaHostRegister`` (a copy may not span two registrations).
-        Only while no training step runs: registering a large buffer can stall the trainer's CUDA calls."""
+        """Page-lock the host buffer, only while no training step runs: registering a large buffer can
+        stall the trainer's CUDA calls."""
         start = time.perf_counter()
-        error = torch.cuda.cudart().cudaHostRegister(self.host_buffer.data_ptr(), self.host_buffer.numel(), 0)
-        assert error == 0, f"cudaHostRegister failed: {error}"
+        page_lock(self.host_buffer)
         self.pinned = True
         get_logger().info(
             f"Page-locked the {self.host_buffer.numel() / 2**30:.1f} GiB sparse weight host buffer in "
@@ -293,8 +296,7 @@ def copy_snapshot_to_host(self) -> None:
         assert self.on_gpu
         with torch.cuda.stream(self.stream):
             for key, bits in self.snapshot.items():
-                self.host[key].copy_(bits, non_blocking=self.pinned)
-            self.stream.synchronize()
+                self._to_host(bits, self.host[key])
 
     def anchor_parts(self) -> list[tuple[str, list[int], int, torch.Tensor]]:
         """This rank's slices of the full HF checkpoint of the last broadcast, from host memory (the CPU
@@ -358,11 +360,12 @@ def encode(
             weights_read.set()
         self.previous = self.snapshot
 
-    def _to_host(self, tensor: torch.Tensor) -> torch.Tensor:
-        """Copy a delta part to pageable host memory through the pinned staging buffer: a copy straight
-        into pageable memory holds the CUDA driver and stalls the training step's own CUDA calls."""
-        data, host = tensor.view(torch.uint8), torch.empty(tensor.shape, dtype=tensor.dtype)
-        host_bytes = host.view(torch.uint8)
+    def _to_host(self, tensor: torch.Tensor, host: torch.Tensor | None = None) -> torch.Tensor:
+        """Copy a GPU tensor to pageable host memory (``host``, or a new tensor) through the pinned staging
+        buffer: a copy straight into pageable memory holds the CUDA driver and stalls the training step's
+        own CUDA calls."""
+        host = torch.empty(tensor.shape, dtype=tensor.dtype) if host is None else host
+        data, host_bytes = tensor.view(torch.uint8), host.view(torch.uint8)
         for begin in range(0, data.numel(), self.staging.numel()):
             chunk = self.staging[: min(self.staging.numel(), data.numel() - begin)]
             chunk.copy_(data[begin : begin + chunk.numel()], non_blocking=True)
@@ -370,30 +373,40 @@ def _to_host(self, tensor: torch.Tensor) -> torch.Tensor:
             host_bytes[begin : begin + chunk.numel()].copy_(chunk)
         return host
 
-    def _segments(self, key: str, fqn: str, offset: int, numel: int) -> tuple[list[str], list[int], list[int]]:
-        """The HF tensor parts in this rank's shard of ``key``: names, and begin/end in the shard."""
+    def _segments(
+        self, key: str, fqn: str, offset: int, numel: int
+    ) -> tuple[list[str], list[int], list[int], list[int]]:
+        """The HF tensor parts in this rank's shard of ``key``: names, begin/end in the shard, and the shift
+        from a shard position to the HF flat index."""
         if key not in self.shard_segments:
-            names, begins, ends = [], [], []
+            names, begins, ends, shifts = [], [], [], []
             for start, name in self.segments[fqn]:
-                if start < offset + numel and start + self.hf_numel[name] > offset:
+                hf_numel = prod(self.hf_tensors[name][0])
+                if start < offset + numel and start + hf_numel > offset:
                     names.append(name)
                     begins.append(max(start - offset, 0))
-                    ends.append(min(start - offset + self.hf_numel[name], numel))
-            self.shard_segments[key] = (names, begins, ends)
+                    ends.append(min(start - offset + hf_numel, numel))
+                    shifts.append(offset - start)
+            self.shard_segments[key] = (names, begins, ends, shifts)
         return self.shard_segments[key]
 
     @staticmethod
-    def _segment_sums(bits: torch.Tensor, begins: list[int], ends: list[int]) -> torch.Tensor:
-        """Wrapping int64 sums of ``bits[begin:end]``, one reduction per run of adjacent equal-length
-        segments (e.g. unstacked experts)."""
-        sums, i = [], 0
-        while i < len(begins):
-            length, j = ends[i] - begins[i], i + 1
-            while j < len(begins) and begins[j] == ends[j - 1] and ends[j] - begins[j] == length:
-                j += 1
-            sums.append(bits[begins[i] : begins[i] + (j - i) * length].view(j - i, length).sum(1, dtype=torch.int64))
-            i = j
-        return torch.cat(sums)
+    def _segment_sums(bits: torch.Tensor, begins: list[int], ends: list[int], firsts: list[int]) -> torch.Tensor:
+        """``bits_checksum`` of each ``bits[begin:end]`` whose first element is at HF flat index ``first``:
+        Σ bits[k]·(first + k - begin + 1), wrapping int64, from cumulative sums over blocks of
+        ``SUM_BLOCK`` elements."""
+        device = bits.device
+        begins, ends, firsts = (torch.tensor(values, device=device) for values in (begins, ends, firsts))
+        plain = torch.zeros(len(begins), dtype=torch.int64, device=device)
+        weighted = torch.zeros_like(plain)
+        for start in range(0, bits.numel(), SUM_BLOCK):
+            values = bits[start : start + SUM_BLOCK].long()
+            first, last = (bound.clamp(start, start + values.numel()) - start for bound in (begins, ends))
+            positions = torch.arange(start + 1, start + 1 + values.numel(), device=device)
+            for total, summed in ((plain, values), (weighted, values * positions)):
+                cumulative = torch.cat([summed.new_zeros(1), summed.cumsum(0)])
+                total += cumulative[last] - cumulative[first]
+        return weighted + (firsts - begins) * plain
 
     def _diff(self, shards: list[tuple], emit: Callable) -> bool:
         """Diff ``shards`` against ``self.previous`` on the current stream, writing the new bits to the
@@ -416,7 +429,7 @@ def flush() -> None:
 
         for key, fqn, shard, offset, gathered in shards:
             previous, out, wire = self.previous[key], self.snapshot[key], self.wire_dtypes[key]
-            segment_names, begins, ends = self._segments(key, fqn, offset, shard.numel())
+            segment_names, begins, ends, shifts = self._segments(key, fqn, offset, shard.numel())
             sums = torch.zeros(len(segment_names), dtype=torch.int64, device=device)
             changed, changed_values = [], []
             if fqn in self.fp8:
@@ -467,6 +480,7 @@ def flush() -> None:
                         new_bits,
                         [max(begin, start) - start for begin in begins[first:last]],
                         [min(stop, end) - start for stop in ends[first:last]],
+                        [max(begin, start) + shift for begin, shift in zip(begins[first:last], shifts[first:last])],
                     )
             changed, changed_values = torch.cat(changed), torch.cat(changed_values)
             if fqn not in self.segment_starts:
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index f80bcdcfeca..7c74c57d412 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -426,9 +426,13 @@ def numa_cpus(device: int) -> set[int] | None:
 def pin(num_bytes: int) -> torch.Tensor:
     """Page-locked host memory of exactly ``num_bytes``, outside torch's caching host allocator
     (which rounds sizes up and keeps freed blocks). Release it with ``unpin``."""
-    tensor = torch.from_numpy(np.empty(num_bytes, np.uint8))
-    if num_bytes:
-        error = torch.cuda.cudart().cudaHostRegister(tensor.data_ptr(), num_bytes, 0)
+    return page_lock(torch.from_numpy(np.empty(num_bytes, np.uint8)))
+
+
+def page_lock(tensor: torch.Tensor) -> torch.Tensor:
+    """Page-lock a host tensor in place with one ``cudaHostRegister`` (a copy may not span two)."""
+    if tensor.numel():
+        error = torch.cuda.cudart().cudaHostRegister(tensor.data_ptr(), tensor.nbytes, 0)
         assert error == 0, f"cudaHostRegister failed: {error}"
     return tensor
 
@@ -443,10 +447,13 @@ def drop_page_cache(fd: int) -> None:
     os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)
 
 
-def bits_checksum(bits: torch.Tensor) -> int:
-    """Wrapping int64 sum of raw bits: additive over any split of a tensor, so trainer
-    ranks and engine workers can each sum the part they hold."""
-    return int(bits.sum(dtype=torch.int64))
+def bits_checksum(bits: torch.Tensor, start: int = 0) -> int:
+    """Wrapping int64 sum of the raw bits of a flat tensor slice, each weighted by its flat index + 1
+    in the full tensor (the slice starts at ``start``): additive over any split of a tensor, so
+    trainer ranks and engine workers can each sum the part they hold, and a change written at a
+    wrong position is caught."""
+    positions = torch.arange(start + 1, start + 1 + bits.numel(), device=bits.device)
+    return int((bits.reshape(-1).long() * positions).sum())
 
 
 def add_checksums(a: int, b: int) -> int:
diff --git a/tests/unit/utils/test_sparse_weights.py b/tests/unit/utils/test_sparse_weights.py
index 8ac79ce677d..cad430c980d 100644
--- a/tests/unit/utils/test_sparse_weights.py
+++ b/tests/unit/utils/test_sparse_weights.py
@@ -31,8 +31,8 @@ def test_sparse_delta_roundtrip_is_exact(tmp_path):
     for rank, names in enumerate([[("a", slice(0, 32))], [("a", slice(32, 64)), ("b", slice(None))]]):
         entries, gaps, values, checksums = [], [], [], {}
         for name, rows in names:
-            checksums[name] = bits_checksum(as_bits(current[name][rows]))
             offset = 0 if rows.start is None else rows.start * 32
+            checksums[name] = bits_checksum(as_bits(current[name][rows]), offset)
             delta = as_bits(current[name][rows]).reshape(-1) - as_bits(previous[name][rows]).reshape(-1)
             changed = delta.nonzero().squeeze(1)
             starts = torch.zeros_like(changed, dtype=torch.bool)

From c87e75e56e421b8d87a0409a8f845785366ce4fb Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Fri, 16 Oct 2026 11:44:42 +0000
Subject: [PATCH 30/88] refactor(weights): simplify the sparse weight broadcast

- The trainer waits for the engine weight layout at the second broadcast instead of publishing full checkpoints until every engine worker reported it.
- fp8 tensors with columns that are not a multiple of 128 raise; the training-thread fp8 diff is gone.
- One write_file helper for every atomic write. Parts markers carry the part count and bytes; the manifest drops anchor_step.
- Pollers of files written on another host revalidate the NFS dir first (cached "not found" lookups last ~30 s).
- Engines apply each read chunk at once, share one reader pool, and one safetensors dtype map.
- snapshot drops the forced "gpu" value.
---
 .../src/prime_rl/configs/shared.py            |   4 +-
 .../inference/vllm/worker/filesystem.py       |  48 ++----
 .../inference/vllm/worker/weight_transfer.py  |   4 -
 src/prime_rl/transports/weights/filesystem.py | 138 ++++++++----------
 src/prime_rl/utils/pathing.py                 |  16 +-
 src/prime_rl/utils/sparse_weights.py          |  75 ++++++----
 tests/unit/utils/test_sparse_weights.py       |   6 +-
 7 files changed, 133 insertions(+), 158 deletions(-)

diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
index 7f9ce1db473..4ac8a6f2675 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
@@ -57,10 +57,10 @@ class FileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig):
     """With ``sparse``, publish a full checkpoint every this many steps (and at the first broadcast).
     Engines that fall off the delta chain reload the last anchor and replay the deltas since."""
 
-    snapshot: Literal["auto", "gpu", "cpu"] = "auto"
+    snapshot: Literal["auto", "cpu"] = "auto"
     """With ``sparse``, where each trainer rank keeps the previous broadcast's weights to diff against.
     ``auto`` keeps them on the GPU if they fit next to the first training step's peak memory, else in
-    page-locked host memory."""
+    page-locked host memory; ``cpu`` always in host memory."""
 
 
 class InMemoryWeightBroadcastConfig(BaseWeightBroadcastConfig):
diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 833671b1aad..7b7a60ae2a5 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -23,6 +23,7 @@
     DELTA_DIR,
     ENGINE_LAYOUT,
     MANIFEST,
+    SAFETENSORS_DTYPES,
     anchor_complete,
     anchor_index,
     anchor_tensors,
@@ -33,6 +34,7 @@
     read_delta,
     read_manifest,
     unpin,
+    write_file,
 )
 
 # This is to get type hints for the Worker class but not actually extend it at runtime as this is required by vLLM worker extension
@@ -46,9 +48,7 @@
 logger = init_logger("vllm.inference.vllm.worker_filesystem")
 
 _prefetch_pool = ThreadPoolExecutor(1)
-_copy_pool: ThreadPoolExecutor | None = None
-
-_SAFETENSORS_DTYPES = {"BF16": "bfloat16", "F16": "float16", "F32": "float32"}
+_reader_pool: ThreadPoolExecutor | None = None
 
 
 class FileSystemWeightUpdateWorker(Worker):
@@ -87,16 +87,16 @@ def _read_delta(
         """Read this worker's part of a delta into a pinned host buffer that is reused across steps
         (allocating pinned memory every step is slow on a loaded host), in chunks of up to 2^26 changes
         of consecutive frames. Reads and copies run on the CPUs of the GPU's NUMA node."""
-        global _copy_pool
+        global _reader_pool
         cpus = numa_cpus(torch.cuda.current_device())
         if cpus is not None:
             os.sched_setaffinity(0, cpus)
-        if _copy_pool is None:
-            _copy_pool = ThreadPoolExecutor(
+        if _reader_pool is None:
+            _reader_pool = ThreadPoolExecutor(
                 16, initializer=None if cpus is None else lambda: os.sched_setaffinity(0, cpus)
             )
         start = time.perf_counter()
-        frames, checksums = read_delta(step_dir, None if self.index_map is None else self.index_map.names, cpus=cpus)
+        frames, checksums = read_delta(step_dir, _reader_pool, None if self.index_map is None else self.index_map.names)
         read_time = time.perf_counter() - start
         chunks, size = [[]], 0
         for frame in frames:
@@ -121,10 +121,12 @@ def _read_delta(
             gaps_end, values_end = gaps_at, values_at
             for _, gaps, values in chunk:
                 copies.append(
-                    _copy_pool.submit(np.copyto, buffer[gaps_end : gaps_end + gaps.nbytes].view(gaps.dtype), gaps)
+                    _reader_pool.submit(np.copyto, buffer[gaps_end : gaps_end + gaps.nbytes].view(gaps.dtype), gaps)
                 )
                 copies.append(
-                    _copy_pool.submit(np.copyto, buffer[values_end : values_end + values.nbytes], values.view(np.uint8))
+                    _reader_pool.submit(
+                        np.copyto, buffer[values_end : values_end + values.nbytes], values.view(np.uint8)
+                    )
                 )
                 gaps_end += gaps.nbytes
                 values_end += values.nbytes
@@ -220,7 +222,7 @@ def read_tensor(name: str) -> torch.Tensor:
                     for name in f.keys():
                         tensor_slice = f.get_slice(name)
                         files[name] = path
-                        tensors[name] = (tensor_slice.get_shape(), _SAFETENSORS_DTYPES[tensor_slice.get_dtype()])
+                        tensors[name] = (tensor_slice.get_shape(), SAFETENSORS_DTYPES[tensor_slice.get_dtype()])
 
             def read_tensor(name: str) -> torch.Tensor:
                 with safe_open(files[name], framework="pt") as f:
@@ -239,9 +241,7 @@ def read_tensor(name: str) -> torch.Tensor:
         layout = {"world_size": group.world_size, "fp8": self.index_map.fp8_layout(read_tensor)}
         path = step_dir.parent / ENGINE_LAYOUT / f"rank_{group.rank_in_group}.json"
         path.parent.mkdir(exist_ok=True)
-        tmp = path.with_suffix(f".tmp{os.getpid()}")
-        tmp.write_text(json.dumps(layout))
-        os.replace(tmp, path)
+        write_file(path, json.dumps(layout).encode())
         logger.info(f"Wrote the engine layout: {len(layout['fp8'])} checkpoint tensors held in fp8")
 
     def _replay(self, step_dir: Path) -> None:
@@ -277,23 +277,7 @@ def _apply_delta(self, step_dir: Path) -> None:
             self.index_map = CheckpointIndexMap(model, tensors)
             logger.info(f"Built the checkpoint index map in {time.perf_counter() - start - read_time:.2f}s")
         for entries, gaps, values in delta:
-            # Decode in chunks of whole entries: the int64 indices take 8 B per change on the GPU.
-            first = change_offset = byte_offset = num_changes = 0
-            for last in range(len(entries)):
-                num_changes += entries[last][1]
-                if num_changes < 1 << 26 and last + 1 < len(entries):
-                    continue
-                chunk = entries[first : last + 1]
-                num_bytes = sum(count * bits_dtype(tensors[name][1]).itemsize for name, count in chunk)
-                self._apply_chunk(
-                    chunk,
-                    gaps[change_offset : change_offset + num_changes],
-                    values[byte_offset : byte_offset + num_bytes],
-                    tensors,
-                    device,
-                )
-                first, change_offset, byte_offset = last + 1, change_offset + num_changes, byte_offset + num_bytes
-                num_changes = 0
+            self._apply_chunk(entries, gaps, values, tensors, device)
         update_mla_absorbed_weights(model)
         torch.cuda.synchronize()
         apply_time = time.perf_counter() - start - read_time
@@ -306,7 +290,8 @@ def _apply_delta(self, step_dir: Path) -> None:
 
     def _verify(self, step_dir: Path, checksums: dict[str, int]) -> None:
         """Check every tensor this worker holds against the trainer's checksum of the new weights.
-        Tensors split across TP workers are summed over the TP group first."""
+        Tensors split across TP workers are summed over the TP group first; without TP, tensors a worker
+        holds only part of are not checked."""
         assert self.index_map is not None
         names, totals, partial = self.index_map.held_checksums()
         tp_group = get_tp_group()
@@ -315,7 +300,6 @@ def _verify(self, step_dir: Path, checksums: dict[str, int]) -> None:
             torch.distributed.all_reduce(reduced, group=tp_group.device_group)
             totals[partial] = reduced
         elif partial.any():
-            totals[partial] = 0
             names = [name if not is_partial else None for name, is_partial in zip(names, partial.tolist())]
         mismatched = [
             name
diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py
index dc109bfabbe..36b10bfbbea 100644
--- a/src/prime_rl/inference/vllm/worker/weight_transfer.py
+++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py
@@ -264,10 +264,6 @@ def _scatter_add(self, owner: torch.Tensor, target: torch.Tensor, source: torch.
             if not count:
                 continue
             flat = as_bits(self.params[index].data).view(-1)
-            if flat.dtype != source.dtype:
-                raise RuntimeError(
-                    f"Sparse weight delta dtype {source.dtype} does not match the engine param ({flat.dtype})"
-                )
             positions = target[begin:end]
             flat.index_put_((positions,), flat[positions] + source[begin:end])
 
diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index 00a6d855cce..d58db24e95c 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -22,7 +22,6 @@
 from prime_rl.trainer.lora import get_lora_state, save_lora_config
 from prime_rl.transports.weights.base import FINISHED_MARKER, WeightReceiver, WeightSender
 from prime_rl.utils.logger import get_logger
-from prime_rl.utils.pathing import wait_for_path
 from prime_rl.utils.sparse_weights import (
     DELTA_DIR,
     FP8_BLOCK,
@@ -37,7 +36,9 @@
     rank_part,
     rank_parts_marker,
     read_engine_layout,
+    revalidate,
     write_anchor_shards,
+    write_file,
     write_rank_file,
     zigzag,
 )
@@ -76,8 +77,9 @@ class SparseDeltaEncoder:
     from the engine layout) are quantized with vLLM's own per-block quantizer,
     previous and current weights alike, and diffed as fp8 bytes. Their 128-row
     blocks must not straddle trainer ranks, else the trainer tensor is gathered.
-    They are diffed on the training thread (``prepare``). A step that changes
-    any block scale is published as a full checkpoint.
+    The quantizer is compiled ahead of time, so it can run on the encode thread
+    (columns a multiple of 128). A step that changes any block scale is
+    published as a full checkpoint.
     """
 
     # Elements per transient chunk of the diff (bf16: 128 MiB per temporary)
@@ -97,7 +99,6 @@ def __init__(self, model: nn.Module, snapshot: str):
 
         segments = defaultdict(list)
         self.hf_tensors: dict[str, tuple[list[int], str]] = {}
-        self.hf_numel: dict[str, int] = {}
         for name, tensor in convert_state_dict_to_hf(model, dict(meta)).items():
             base = tensor if tensor._base is None else tensor._base
             if id(base) not in owner or not tensor.is_contiguous():
@@ -106,25 +107,22 @@ def __init__(self, model: nn.Module, snapshot: str):
                 raise NotImplementedError(f"Sparse weight broadcast indexes tensors below 2^32 elements ({name})")
             segments[owner[id(base)]].append((tensor.storage_offset(), name))
             self.hf_tensors[name] = (list(tensor.shape), str(tensor.dtype).removeprefix("torch."))
-            self.hf_numel[name] = tensor.numel()
         self.segments = {fqn: sorted(segs) for fqn, segs in segments.items()}
         self.numel = {fqn: tensor.numel() for fqn, tensor in meta.items()}
         # fqn -> (ue8m0 scales, columns) for trainer tensors held in fp8; quantized as one [rows, columns] matrix
         self.fp8: dict[str, tuple[bool, int]] = {}
         self.fp8_layout: dict[str, str] | None = None
-        self.fp8_background: set[str] = set()
         self.compiled_fp8 = None
         self.scales_changed = False
         self.gather: set[str] = set()
         self.segment_starts: dict[str, torch.Tensor] = {}
-        self.shard_segments: dict[str, tuple[list[str], list[int], list[int]]] = {}
+        self.shard_segments: dict[str, tuple[list[str], list[int], list[int], list[int]]] = {}
         self.snapshot_placement = snapshot
         # key -> (local numel, global flat offset) of this rank's shard
         self.shards: dict[str, tuple[int, int]] = {}
         # key -> previous wire bits: GPU tensors, or views into the host buffer
         self.snapshot: dict[str, torch.Tensor] = {}
         self.previous: dict[str, torch.Tensor] = {}
-        self.fp8_parts: list[tuple] = []
         self.on_gpu: bool | None = None  # None until placed (auto: at the first delta)
         self.host_buffer: torch.Tensor | None = None
         self.host: dict[str, torch.Tensor] = {}
@@ -148,6 +146,7 @@ def set_fp8_layout(self, model: nn.Module, fp8: dict[str, str], is_master: bool)
             ]
             if (
                 len(modes) > 1
+                or num_cols % FP8_BLOCK[1]
                 or any(len(shape) != 2 or shape[1] != num_cols for shape in shapes)
                 or any(row % FP8_BLOCK[0] and row != num_rows for row in boundaries)
             ):
@@ -156,9 +155,7 @@ def set_fp8_layout(self, model: nn.Module, fp8: dict[str, str], is_master: bool)
             for _, name in segs:
                 self.hf_tensors[name] = (self.hf_tensors[name][0], "float8_e4m3fn")
         self.gather = self._find_gathered(model, is_master)
-        # Diffed on the encode thread with the ahead-of-time compiled quantizer; others on the training thread
-        self.fp8_background = {fqn for fqn, (_, num_cols) in self.fp8.items() if num_cols % FP8_BLOCK[1] == 0}
-        if self.fp8_background:
+        if self.fp8:
             start = time.perf_counter()
             self.compiled_fp8 = compile_fp8_quantizer()
             get_logger().info(f"Compiled the fp8 quantizer ahead of time in {time.perf_counter() - start:.1f}s")
@@ -260,10 +257,7 @@ def _place(self, peak_reserved: int) -> None:
         num_bytes = sum(numel * self.wire_dtypes[key].itemsize for key, (numel, _) in self.shards.items())
         total = torch.cuda.get_device_properties(self.device).total_memory
         headroom = total - peak_reserved - num_bytes
-        if self.snapshot_placement == "auto":
-            self.on_gpu = headroom >= max(4 * 2**30, total // 20)
-        else:
-            self.on_gpu = self.snapshot_placement == "gpu"
+        self.on_gpu = self.snapshot_placement == "auto" and headroom >= max(4 * 2**30, total // 20)
         get_logger().info(
             f"Sparse weight snapshot ({num_bytes / 2**30:.1f} GiB) on the {'GPU' if self.on_gpu else 'host'} "
             f"(training peak {peak_reserved / 2**30:.1f} GiB of {total / 2**30:.1f} GiB, "
@@ -271,9 +265,9 @@ def _place(self, peak_reserved: int) -> None:
         )
 
     def _fill(self, shards: list[tuple]) -> None:
-        """The first broadcast only records the weights, in the host buffer, and on the GPU if forced there
-        (else placed at the first delta, against the measured training peak). The host buffer is filled in
-        every mode, so its pages are resident before training: faulting them in later stalls the trainer."""
+        """The first broadcast only records the weights, in the host buffer (placed at the first delta,
+        against the measured training peak). The host buffer is filled in every mode, so its pages are
+        resident before training: faulting them in later stalls the trainer."""
         self.shards = {key: (shard.numel(), offset) for key, _, shard, offset, _ in shards}
         self.staging = torch.empty(1 << 28, dtype=torch.uint8, pin_memory=True)
         self._allocate_host()
@@ -283,12 +277,7 @@ def _fill(self, shards: list[tuple]) -> None:
             for start in range(0, shard.numel(), self.CHUNK):
                 end = start + self.CHUNK
                 self.host[key][start:end].copy_(as_bits(shard[start:end].to(self.wire_dtypes[key])), non_blocking=True)
-        if self.snapshot_placement == "gpu":
-            self.on_gpu = True
-            for key, _, shard, _, _ in shards:
-                self.snapshot[key] = as_bits(shard.to(f"cuda:{self.device}", self.wire_dtypes[key], copy=True))
-        else:
-            self.snapshot = dict(self.host)
+        self.snapshot = dict(self.host)
 
     def copy_snapshot_to_host(self) -> None:
         """Copy a GPU snapshot (the last broadcast's weights) into the host buffer, for an anchor. Callers wait
@@ -305,33 +294,23 @@ def anchor_parts(self) -> list[tuple[str, list[int], int, torch.Tensor]]:
         for key, bits in self.host.items():
             values, (numel, offset) = bits.view(self.wire_dtypes[key]), self.shards[key]
             for start, name in self.segments[self.key_fqn[key]]:
-                begin, end = max(start, offset), min(start + self.hf_numel[name], offset + numel)
+                begin, end = max(start, offset), min(start + prod(self.hf_tensors[name][0]), offset + numel)
                 if begin < end:
                     parts.append((name, self.hf_tensors[name][0], begin - start, values[begin - offset : end - offset]))
         return parts
 
     @torch.no_grad()
-    def prepare(self, shards: list[tuple], peak_reserved: int) -> None:
-        """The training thread's part of an encode, before ``encode``: place the snapshot at the first delta,
-        and diff the tensors held in fp8 that the ahead-of-time compiled quantizer does not cover (vLLM's
-        quantizer is torch.compile'd, and running it next to the training step breaks activation
-        checkpointing there). ``scales_changed``: whether an fp8 block scale changed on this rank, after
-        ``encode``."""
+    def prepare(self, peak_reserved: int) -> None:
+        """The training thread's part of an encode, before ``encode``: place the snapshot at the first delta."""
         self.previous = self.snapshot
-        self.fp8_parts = []
-        self.scales_changed = False
-        if not self.initialized:
+        if not self.initialized or self.on_gpu is not None:
             return
-        if self.on_gpu is None:
-            self._place(peak_reserved)
-            if self.on_gpu:
-                device = f"cuda:{self.device}"
-                self.snapshot = {key: torch.empty_like(bits, device=device) for key, bits in self.host.items()}
-            elif not self.pinned:
-                self._page_lock()
-        here = [shard for shard in shards if shard[1] in self.fp8 and shard[1] not in self.fp8_background]
-        if here:
-            self.scales_changed = self._diff(here, self.fp8_parts.append)
+        self._place(peak_reserved)
+        if self.on_gpu:
+            device = f"cuda:{self.device}"
+            self.snapshot = {key: torch.empty_like(bits, device=device) for key, bits in self.host.items()}
+        elif not self.pinned:
+            self._page_lock()
 
     @torch.no_grad()
     def encode(
@@ -340,7 +319,8 @@ def encode(
         """Diff the shards against the snapshot and advance it, on the side stream once ``ready`` (recorded
         after the optimizer step) has passed; the first call only records the snapshot. Hands the delta to
         ``emit`` in parts ``(entries, gaps, value bytes, new-weight checksum per HF tensor part)`` on the
-        host as they are encoded. Sets ``weights_read`` once done reading the weights."""
+        host as they are encoded. Sets ``weights_read`` once done reading the weights. ``scales_changed``:
+        whether an fp8 block scale changed on this rank."""
         torch.cuda.set_device(self.device)
         try:
             with torch.cuda.stream(self.stream):
@@ -351,14 +331,10 @@ def encode(
                     self.stream.synchronize()
                     self.initialized = True
                     return
-                for part in self.fp8_parts:
-                    emit(part)
-                rest = [shard for shard in shards if shard[1] not in self.fp8 or shard[1] in self.fp8_background]
-                self.scales_changed |= self._diff(rest, emit)
+                self.scales_changed = self._diff(shards, emit)
                 self.stream.synchronize()
         finally:
             weights_read.set()
-        self.previous = self.snapshot
 
     def _to_host(self, tensor: torch.Tensor, host: torch.Tensor | None = None) -> torch.Tensor:
         """Copy a GPU tensor to pageable host memory (``host``, or a new tensor) through the pinned staging
@@ -454,11 +430,7 @@ def flush() -> None:
                             current.view(torch.bfloat16),
                         )
                     (previous_fp8, previous_scales), (current_fp8, current_scales) = (
-                        quantize_fp8_blocks(
-                            weights.reshape(-1, num_cols),
-                            ue8m0,
-                            self.compiled_fp8 if fqn in self.fp8_background else None,
-                        )
+                        quantize_fp8_blocks(weights.reshape(-1, num_cols), ue8m0, self.compiled_fp8)
                         for weights in (previous_weights, current_weights)
                     )
                     scales_changed |= not torch.equal(previous_scales, current_scales)
@@ -512,9 +484,7 @@ def flush() -> None:
 
 
 def _write_json(path: Path, value) -> None:
-    tmp = path.with_name(f".{path.name}.tmp")
-    tmp.write_text(json.dumps(value))
-    tmp.rename(path)
+    write_file(path, json.dumps(value).encode())
 
 
 class FileSystemWeightSender(WeightSender):
@@ -548,6 +518,7 @@ def __init__(
         self.errors: list[BaseException] = []
         self.weights_read: threading.Event | None = None
         self.deferred_finish = False
+        self.part_bytes = 0
         self.gathers_full_weights = not config.sparse
         self.logger.debug("Initialized filesystem weight broadcast")
 
@@ -626,24 +597,31 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
             self.deferred_finish = True
             return
         if self.encoder.fp8_layout is None:
-            # Until every engine worker has reported which tensors it holds in fp8, steps are full checkpoints.
-            layout = [read_engine_layout(step_dir.parent) if self.world.is_master else None]
+            # Every engine worker reports which tensors it holds in fp8 when it loads the first broadcast.
+            layout = [self._wait_for_engine_layout(step_dir.parent) if self.world.is_master else None]
             dist.broadcast_object_list(layout, src=0)
-            if layout[0] is not None:
-                self.encoder.set_fp8_layout(model, layout[0], self.world.is_master)
+            self.encoder.set_fp8_layout(model, layout[0], self.world.is_master)
         shards = self.encoder.local_shards(model, self.world.is_master)
-        self.encoder.prepare(shards, torch.cuda.max_memory_reserved())
-        full = self.encoder.fp8_layout is None
+        self.encoder.prepare(torch.cuda.max_memory_reserved())
         # Periodic anchors only serve engines that fall off the chain: written in the background, each after
         # the previous one.
-        anchor = not full and step % self.config.anchor_interval == 0
+        anchor = step % self.config.anchor_interval == 0
         ready = torch.cuda.Event()
         ready.record()
         self.weights_read = threading.Event()
-        args = (step, base_step, step_dir, shards, ready, self.weights_read, full, anchor)
+        args = (step, base_step, step_dir, shards, ready, self.weights_read, anchor)
         self._submit(self.encodes, self._encode, *args)
         self.deferred_finish = True
 
+    def _wait_for_engine_layout(self, broadcast_dir: Path) -> dict[str, str]:
+        start = time.monotonic()
+        while (layout := read_engine_layout(broadcast_dir)) is None:
+            if time.monotonic() - start > self.timeout:
+                raise TimeoutError(f"Not every engine worker reported its weight layout within {self.timeout}s")
+            time.sleep(0.1)
+        self.logger.debug(f"Waited {time.monotonic() - start:.2f}s for the engine weight layout")
+        return layout
+
     def _prepare(self, model: nn.Module, step: int) -> None:
         if self.lora_config is not None or not self.config.sparse:
             return
@@ -692,21 +670,16 @@ def _encode(
         shards,
         ready,
         weights_read: threading.Event,
-        full: bool,
         anchor: bool,
     ) -> None:
-        """Encode thread: diff v{step}, handing the delta to the writer in parts as it is encoded (on full
-        steps, only advance the snapshot and hand the checkpoint to the writer). With fp8 deltas, the parts
-        are held until the ranks agree that no block scale changed, else the step is a full checkpoint."""
+        """Encode thread: diff v{step}, handing the delta to the writer in parts as it is encoded. With fp8
+        deltas, the parts are held until the ranks agree that no block scale changed, else the step is a
+        full checkpoint."""
         encoder = self.encoder
         assert encoder is not None
         if not encoder.on_gpu:
             self._wait_for_host_readers()  # anchors written from the host snapshot, which the encode advances
         start = time.perf_counter()
-        if full:
-            encoder.encode(shards, ready, weights_read, emit=lambda part: None)
-            self._publish_full(step, step_dir)
-            return
         held: list[tuple] = []
         num_parts = 0
 
@@ -742,6 +715,7 @@ def _scales_changed_anywhere(self, step_dir: Path, changed: bool) -> bool:
         flags = [delta_dir / f"rank_{rank}.scales" for rank in range(self.world.world_size)]
         while not all(flag.exists() for flag in flags):
             time.sleep(0.05)
+            revalidate(delta_dir)
         return any(flag.read_text() == "1" for flag in flags)
 
     def _start_delta(self, step: int, base_step: int, step_dir: Path) -> None:
@@ -750,7 +724,6 @@ def _start_delta(self, step: int, base_step: int, step_dir: Path) -> None:
         manifest = {
             "step": step,
             "base_step": base_step,
-            "anchor_step": self.anchor_step,
             "ranks": self.world.world_size,
             "tensors": self.encoder.hf_tensors,
         }
@@ -760,20 +733,24 @@ def _write_part(self, step_dir: Path, index: int, part: tuple) -> None:
         assert self.encoder is not None
         entries, gaps, values, checksums = part
         path = step_dir / DELTA_DIR / rank_part(self.world.rank, index)
-        write_rank_file(path, entries, gaps.numpy(), values.numpy(), self.encoder.hf_tensors, checksums)
+        self.part_bytes += write_rank_file(
+            path, entries, gaps.numpy(), values.numpy(), self.encoder.hf_tensors, checksums
+        )
 
     def _finish_delta(self, step: int, base_step: int, step_dir: Path, num_parts: int, encode_time: float) -> None:
         """After this rank's last part: its parts marker. The master then waits for every rank's marker and
         commits the step. Runs on the writer thread, without collectives."""
         delta_dir = step_dir / DELTA_DIR
-        _write_json(delta_dir / rank_parts_marker(self.world.rank), num_parts)
+        _write_json(delta_dir / rank_parts_marker(self.world.rank), {"parts": num_parts, "bytes": self.part_bytes})
+        self.part_bytes = 0
         if not self.world.is_master:
             return
         markers = [delta_dir / rank_parts_marker(rank) for rank in range(self.world.world_size)]
         while not all(marker.exists() for marker in markers):
             time.sleep(0.05)
+            revalidate(delta_dir)
         WeightSender._finish(self, step, step_dir)
-        num_bytes = sum(path.stat().st_size for path in delta_dir.glob("rank_*.safetensors"))
+        num_bytes = sum(json.loads(marker.read_text())["bytes"] for marker in markers)
         self.logger.info(
             f"Sparse weight delta v{step} (base v{base_step}): {num_bytes / 2**20:.1f} MiB, "
             f"encode {encode_time:.2f}s (parts written as encoded)"
@@ -825,11 +802,14 @@ async def receive(self, step: int) -> None:
         # read and decode them while they keep serving, and the pause only applies the delta.
         while not (manifest.exists() or finished.exists()):
             await asyncio.sleep(0.05)
+            revalidate(manifest.parent)
         if manifest.exists():
             start = time.perf_counter()
             await self.admin_plane.prefetch_weights(weights_dir)
             self.logger.debug(f"Prefetched sparse weight delta v{step} in {time.perf_counter() - start:.2f}s")
-        await wait_for_path(finished)
+        while not finished.exists():
+            await asyncio.sleep(0.05)
+            revalidate(weights_dir)
         if (weights_dir / "adapter_config.json").exists():
             await load_lora_adapter(self.admin_plane, self.model_name, weights_dir)
             return
diff --git a/src/prime_rl/utils/pathing.py b/src/prime_rl/utils/pathing.py
index fd169d7d6b9..a02c19f6006 100644
--- a/src/prime_rl/utils/pathing.py
+++ b/src/prime_rl/utils/pathing.py
@@ -435,13 +435,15 @@ def sync_wait_for_path(path: Path, interval: int = 1, log_interval: int = 10) ->
         wait_time += interval
 
 
-async def wait_for_path(path: Path, interval: float = 0.1, log_interval: float = 10) -> None:
+async def wait_for_path(path: Path, interval: int = 1, log_interval: int = 10) -> None:
     logger = get_logger()
+    wait_time = 0
     logger.debug(f"Waiting for path `{path}`")
-    start = last_log = time.monotonic()
-    while not path.exists():
-        if time.monotonic() - last_log >= log_interval:
-            last_log = time.monotonic()
-            logger.debug(f"Waiting for path `{path}` for {last_log - start:.0f} seconds")
+    while True:
+        if path.exists():
+            logger.debug(f"Found path `{path}`")
+            break
+        if wait_time % log_interval == 0 and wait_time > 0:  # Every log_interval seconds
+            logger.debug(f"Waiting for path `{path}` for {wait_time} seconds")
         await asyncio.sleep(interval)
-    logger.debug(f"Found path `{path}`")
+        wait_time += interval
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index 7c74c57d412..e90672e7cd2 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -15,6 +15,7 @@
 import json
 import os
 import re
+import threading
 import time
 from collections import defaultdict
 from collections.abc import Iterable, Iterator
@@ -36,6 +37,8 @@
 # Dir next to the step dirs where every engine worker writes which checkpoint tensors it holds in fp8.
 ENGINE_LAYOUT = "engine_layout"
 FP8_BLOCK = [128, 128]
+# safetensors dtype -> torch dtype name
+SAFETENSORS_DTYPES = {"BF16": "bfloat16", "F16": "float16", "F32": "float32"}
 
 _BITS_DTYPE = {1: torch.uint8, 2: torch.int16, 4: torch.int32, 8: torch.int64}
 
@@ -148,12 +151,12 @@ def write_rank_file(
     tensors: dict,
     checksums: dict[str, int],
     frame_size: int = 1 << 20,
-) -> None:
+) -> int:
     """Write a rank's delta as zstd frames of about ``frame_size`` changes each (whole entries), so
     readers can skip frames without entries they hold. ``values`` holds each entry's ``zigzag``
     differences in its bits dtype. Gaps and values are stored as uint8; those of 255 or more go to
     the int32 escape stream (gaps first). ``checksums`` (HF name -> ``bits_checksum`` of the new weights
-    this rank holds) let engines verify every step. Written under a temporary name, then renamed."""
+    this rank holds) let engines verify every step. Returns the file size."""
     options = {pyzstd.CParameter.compressionLevel: 1, pyzstd.CParameter.checksumFlag: 1}
     frames, size = [[]], 0
     for index, (_, count) in enumerate(entries):
@@ -189,7 +192,9 @@ def compress(frame: list[int]) -> tuple[bytes, bytes, bytes]:
         "frames": json.dumps([[len(frame), *map(len, parts)] for frame, parts in zip(frames, compressed)]),
         "checksums": json.dumps(checksums),
     }
-    _write_atomically(path, save(arrays, metadata=metadata))
+    data = save(arrays, metadata=metadata)
+    write_file(path, data)
+    return len(data)
 
 
 def decode_gaps(gaps: torch.Tensor, counts: torch.Tensor) -> torch.Tensor:
@@ -247,12 +252,14 @@ def _read_frame(
     return kept, np.concatenate(kept_gaps), np.concatenate(kept_values)
 
 
-def _write_atomically(path: Path, data: bytes) -> None:
-    """Write under a temporary name, then rename. Serializing in memory and writing with plain file
-    I/O releases the GIL during the (possibly slow) write, so background writes don't stall training."""
-    tmp = path.with_name(f".{path.name}.tmp")
+def write_file(path: Path, data: bytes | list) -> None:
+    """Write ``data`` (bytes, or buffers written back to back) under a temporary name, then rename, so
+    readers never see a partial file. Plain file I/O releases the GIL during the (possibly slow) write,
+    so background writes don't stall training."""
+    tmp = path.with_name(f".{path.name}.{os.getpid()}.{threading.get_ident()}.tmp")
     with open(tmp, "wb") as file:
-        file.write(data)
+        for buffer in [data] if isinstance(data, bytes) else data:
+            file.write(buffer)
         drop_page_cache(file.fileno())
     tmp.rename(path)
 
@@ -273,7 +280,7 @@ def write_anchor_shards(
     for index, file_parts in enumerate(files):
         header, offset = {"__metadata__": {"shapes": json.dumps({name: shape for name, shape, _, _ in file_parts})}}, 0
         for name, _, start, values in file_parts:
-            dtype = _SAFETENSORS_DTYPES[values.dtype]
+            dtype = next(name for name, dtype in SAFETENSORS_DTYPES.items() if getattr(torch, dtype) == values.dtype)
             header[f"{name}@{start}"] = {
                 "dtype": dtype,
                 "shape": [values.numel()],
@@ -282,18 +289,10 @@ def write_anchor_shards(
             offset += values.nbytes
         encoded = json.dumps(header).encode()
         encoded += b" " * (-len(encoded) % 8)
-        path = step_dir / f"model-rank{rank:05d}-{index:05d}.safetensors"
-        tmp = path.with_name(f".{path.name}.tmp")
-        with open(tmp, "wb") as file:
-            file.write(len(encoded).to_bytes(8, "little") + encoded)
-            for *_, values in file_parts:
-                file.write(memoryview(values.view(torch.uint8).numpy()))
-            drop_page_cache(file.fileno())
-        tmp.rename(path)
-    (step_dir / f"anchor_rank{rank}.done").write_text(str(world_size))
-
-
-_SAFETENSORS_DTYPES = {torch.bfloat16: "BF16", torch.float16: "F16", torch.float32: "F32"}
+        buffers = [len(encoded).to_bytes(8, "little") + encoded]
+        buffers.extend(values.view(torch.uint8).numpy() for *_, values in file_parts)
+        write_file(step_dir / f"model-rank{rank:05d}-{index:05d}.safetensors", buffers)
+    write_file(step_dir / f"anchor_rank{rank}.done", str(world_size).encode())
 
 
 def anchor_index(step_dir: Path) -> tuple[dict[str, list[tuple[Path, str, int]]], dict[str, tuple[list[int], str]]]:
@@ -307,7 +306,7 @@ def anchor_index(step_dir: Path) -> tuple[dict[str, list[tuple[Path, str, int]]]
                 name, start = key.rsplit("@", 1)
                 slices[name].append((path, key, int(start)))
                 dtype = str(file.get_slice(key).get_dtype())
-                tensors[name] = (shapes[name], {"BF16": "bfloat16", "F16": "float16", "F32": "float32"}[dtype])
+                tensors[name] = (shapes[name], SAFETENSORS_DTYPES[dtype])
     return dict(slices), tensors
 
 
@@ -364,33 +363,34 @@ def rank_part(rank: int, part: int) -> str:
 
 
 def rank_parts_marker(rank: int) -> str:
-    """Written by a trainer rank after its last part file: the number of part files."""
+    """Written by a trainer rank after its last part file: the number of part files and their bytes."""
     return f"rank_{rank}.parts"
 
 
 def read_delta(
-    step_dir: Path, names: set[str] | None = None, num_threads: int = 16, cpus: set[int] | None = None
+    step_dir: Path, pool: ThreadPoolExecutor, names: set[str] | None = None
 ) -> tuple[list[tuple[list[tuple[str, int]], np.ndarray, np.ndarray]], dict[str, int]]:
     """Read and decompress a delta step dir into ``(entries, gaps, value bytes)`` per frame in host
-    memory (values: new minus previous raw bits, in each entry's bits dtype); ``decode_gaps`` turns
-    the gaps into flat indices on the device that applies them. With ``names``, keep only those
+    memory (values: new minus previous raw bits, in each entry's bits dtype) on ``pool``; ``decode_gaps``
+    turns the gaps into flat indices on the device that applies them. With ``names``, keep only those
     tensors (frames without any are not read). Part files are read as the trainer writes them, until
-    every rank has written its parts marker. ``cpus``: CPUs for the reader threads (the device's NUMA
-    node). Also returns the new weights' ``bits_checksum`` per HF tensor, summed over ranks."""
+    every rank has written its parts marker. Also returns the new weights' ``bits_checksum`` per HF
+    tensor, summed over ranks."""
     manifest = read_manifest(step_dir)
     tensors = manifest["tensors"]
     delta_dir = step_dir / DELTA_DIR
     checksums: dict[str, int] = {}
     num_parts: list[int | None] = [None] * manifest["ranks"]
     next_part = [0] * manifest["ranks"]
-    initializer = None if cpus is None else (lambda: os.sched_setaffinity(0, cpus))
-    with ExitStack() as stack, ThreadPoolExecutor(num_threads, initializer=initializer) as pool:
+    with ExitStack() as stack:
         files, jobs = [], []
         while any(count is None or next_part[rank] < count for rank, count in enumerate(num_parts)):
             found = False
+            revalidate(delta_dir)
             for rank in range(len(num_parts)):
-                if num_parts[rank] is None and (delta_dir / rank_parts_marker(rank)).exists():
-                    num_parts[rank] = int((delta_dir / rank_parts_marker(rank)).read_text())
+                marker = delta_dir / rank_parts_marker(rank)
+                if num_parts[rank] is None and marker.exists():
+                    num_parts[rank] = json.loads(marker.read_text())["parts"]
                 path = delta_dir / rank_part(rank, next_part[rank])
                 if (num_parts[rank] is None or next_part[rank] < num_parts[rank]) and path.exists():
                     file = stack.enter_context(open(path, "rb"))
@@ -442,6 +442,17 @@ def unpin(tensor: torch.Tensor) -> None:
         torch.cuda.cudart().cudaHostUnregister(tensor.data_ptr())
 
 
+def revalidate(directory: Path) -> None:
+    """Make an NFS client see files another host just created in ``directory``: opening a dir (and
+    the one above, which may hold ``directory`` itself) revalidates it, instead of answering from
+    cached "not found" lookups for up to ~30 s."""
+    for path in (directory.parent, directory):
+        try:
+            os.close(os.open(path, os.O_RDONLY | os.O_DIRECTORY))
+        except FileNotFoundError:
+            pass
+
+
 def drop_page_cache(fd: int) -> None:
     """Deltas and anchors are read once; don't let them fill the inference host's page cache."""
     os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)
diff --git a/tests/unit/utils/test_sparse_weights.py b/tests/unit/utils/test_sparse_weights.py
index cad430c980d..0d4af944e3d 100644
--- a/tests/unit/utils/test_sparse_weights.py
+++ b/tests/unit/utils/test_sparse_weights.py
@@ -1,4 +1,5 @@
 import json
+from concurrent.futures import ThreadPoolExecutor
 
 import torch
 
@@ -43,12 +44,13 @@ def test_sparse_delta_roundtrip_is_exact(tmp_path):
         gaps, values = torch.cat(gaps).numpy(), torch.cat(values).numpy()
         # Small frames, so a frame boundary falls inside rank 1's entries
         write_rank_file(delta_dir / rank_part(rank, 0), entries, gaps, values, tensors, checksums, frame_size=8)
-        (delta_dir / rank_parts_marker(rank)).write_text("1")
+        (delta_dir / rank_parts_marker(rank)).write_text(json.dumps({"parts": 1}))
     (delta_dir / MANIFEST).write_text(json.dumps({"ranks": 2, "tensors": tensors}))
 
     def apply(names):
         patched = {name: as_bits(tensor.clone()).view(-1) for name, tensor in previous.items()}
-        delta, checksums = read_delta(tmp_path, names)
+        with ThreadPoolExecutor(4) as pool:
+            delta, checksums = read_delta(tmp_path, pool, names)
         assert checksums == {name: bits_checksum(as_bits(tensor)) for name, tensor in current.items()}
         for entries, gaps, values in delta:
             gaps, values = torch.from_numpy(gaps), torch.from_numpy(values)

From dbe267ba22ea445fe738945a1ec35461845ea46a Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Fri, 16 Oct 2026 11:47:47 +0000
Subject: [PATCH 31/88] refactor(weights): sparse sender as a
 FileSystemWeightSender subclass in its own module

setup_weight_sender picks SparseFileSystemWeightSender for sparse filesystem broadcasts; the dense filesystem sender is unchanged from main. LoRA with a sparse broadcast is rejected by the trainer config.
---
 src/prime_rl/transports/weights/__init__.py   |   3 +
 src/prime_rl/transports/weights/filesystem.py | 746 +-----------------
 src/prime_rl/transports/weights/sparse.py     | 734 +++++++++++++++++
 src/prime_rl/utils/sparse_weights.py          |   2 +-
 4 files changed, 747 insertions(+), 738 deletions(-)
 create mode 100644 src/prime_rl/transports/weights/sparse.py

diff --git a/src/prime_rl/transports/weights/__init__.py b/src/prime_rl/transports/weights/__init__.py
index 128fd54be92..c5150bf4682 100644
--- a/src/prime_rl/transports/weights/__init__.py
+++ b/src/prime_rl/transports/weights/__init__.py
@@ -10,6 +10,7 @@
 from prime_rl.transports.weights.filesystem import FileSystemWeightReceiver, FileSystemWeightSender
 from prime_rl.transports.weights.nccl import NCCLWeightReceiver, NCCLWeightSender
 from prime_rl.transports.weights.nixl import NIXLWeightReceiver, NIXLWeightSender
+from prime_rl.transports.weights.sparse import SparseFileSystemWeightSender
 
 __all__ = [
     "WeightReceiver",
@@ -28,6 +29,8 @@ def setup_weight_sender(
 ) -> WeightSender:
     if config.type == "nccl":
         return NCCLWeightSender(output_dir, config, torch.cuda.current_device())
+    elif config.type == "filesystem" and config.sparse:
+        return SparseFileSystemWeightSender(output_dir, config)
     elif config.type == "filesystem":
         return FileSystemWeightSender(output_dir, config, lora_config)
     elif config.type == "nixl":
diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index d58db24e95c..0b45219cdf5 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -1,500 +1,27 @@
 import asyncio
-import json
-import threading
 import time
-from bisect import bisect_left, bisect_right
-from collections import defaultdict
-from collections.abc import Callable
-from concurrent.futures import Future, ThreadPoolExecutor
-from math import prod
 from pathlib import Path
 
-import numpy as np
-import torch
-import torch.distributed as dist
 import torch.nn as nn
 from torch.distributed.tensor import DTensor
-from torch.distributed.tensor._utils import compute_local_shape_and_global_offset
 
 from prime_rl.configs.shared import FileSystemWeightBroadcastConfig
 from prime_rl.configs.trainer import LoRAConfig
 from prime_rl.orchestrator.clients import load_lora_adapter
 from prime_rl.trainer.lora import get_lora_state, save_lora_config
 from prime_rl.transports.weights.base import FINISHED_MARKER, WeightReceiver, WeightSender
-from prime_rl.utils.logger import get_logger
-from prime_rl.utils.sparse_weights import (
-    DELTA_DIR,
-    FP8_BLOCK,
-    MANIFEST,
-    anchor_complete,
-    as_bits,
-    bits_dtype,
-    compile_fp8_quantizer,
-    encode_gaps,
-    page_lock,
-    quantize_fp8_blocks,
-    rank_part,
-    rank_parts_marker,
-    read_engine_layout,
-    revalidate,
-    write_anchor_shards,
-    write_file,
-    write_rank_file,
-    zigzag,
-)
+from prime_rl.utils.sparse_weights import DELTA_DIR, MANIFEST, revalidate
 from prime_rl.utils.weights import (
     convert_state_dict_to_hf,
     gather_weights_parallel,
-    resolve_fqn,
-    resolve_wire_dtype,
     save_state_dict,
     save_state_dict_parallel,
 )
 
-# Elements per block of the position-weighted checksums (int64 temporaries of 8 B each)
-SUM_BLOCK = 1 << 22
-
-
-class SparseDeltaEncoder:
-    """Rank-local deltas of the trainer's weight shards, in HF coordinates.
-
-    Every rank keeps the wire-dtype bits of its own shards from the previous
-    broadcast (the snapshot) and diffs against them, so no weights are gathered.
-    The prime->HF conversion is replayed once on meta tensors: every HF tensor
-    must come out as a contiguous view of one trainer tensor (renames, expert
-    unstacking), which maps a changed trainer element to its HF name and index.
-
-    ``encode`` runs on a background thread and a side CUDA stream while the
-    next step trains, without collectives; it only reads the weights before
-    ``weights_read`` is set, and hands the delta to the writer in parts as it
-    goes. The snapshot lives on the GPU when it fits next to the first training
-    step's peak memory, else in one page-locked host buffer (locked only between
-    training steps: registering it can stall the trainer's CUDA calls). With a
-    GPU snapshot, the same host buffer (pageable) holds anchor copies instead, so
-    a rank has at most one shard-sized host copy.
-
-    Tensors the engines hold in fp8 (``set_fp8_layout``: HF name -> scale format,
-    from the engine layout) are quantized with vLLM's own per-block quantizer,
-    previous and current weights alike, and diffed as fp8 bytes. Their 128-row
-    blocks must not straddle trainer ranks, else the trainer tensor is gathered.
-    The quantizer is compiled ahead of time, so it can run on the encode thread
-    (columns a multiple of 128). A step that changes any block scale is
-    published as a full checkpoint.
-    """
-
-    # Elements per transient chunk of the diff (bf16: 128 MiB per temporary)
-    CHUNK = 1 << 26
-    # Changes per part handed to the writer
-    PART = 1 << 26
-
-    def __init__(self, model: nn.Module, snapshot: str):
-        keep_in_fp32 = getattr(model, "keep_in_fp32_for_weight_transfer", None)
-        self.wire_dtypes: dict[str, torch.dtype] = {}
-        meta: dict[str, torch.Tensor] = {}
-        for key, value in model.state_dict().items():
-            dtype = resolve_wire_dtype(keep_in_fp32, key, torch.bfloat16) if isinstance(value, DTensor) else value.dtype
-            self.wire_dtypes[key] = dtype
-            meta[resolve_fqn(model, key)] = torch.empty(value.shape, dtype=dtype, device="meta")
-        owner = {id(tensor): fqn for fqn, tensor in meta.items()}
-
-        segments = defaultdict(list)
-        self.hf_tensors: dict[str, tuple[list[int], str]] = {}
-        for name, tensor in convert_state_dict_to_hf(model, dict(meta)).items():
-            base = tensor if tensor._base is None else tensor._base
-            if id(base) not in owner or not tensor.is_contiguous():
-                raise NotImplementedError(f"Sparse weight broadcast cannot map {name} to a contiguous trainer slice")
-            if tensor.numel() >= 2**32:
-                raise NotImplementedError(f"Sparse weight broadcast indexes tensors below 2^32 elements ({name})")
-            segments[owner[id(base)]].append((tensor.storage_offset(), name))
-            self.hf_tensors[name] = (list(tensor.shape), str(tensor.dtype).removeprefix("torch."))
-        self.segments = {fqn: sorted(segs) for fqn, segs in segments.items()}
-        self.numel = {fqn: tensor.numel() for fqn, tensor in meta.items()}
-        # fqn -> (ue8m0 scales, columns) for trainer tensors held in fp8; quantized as one [rows, columns] matrix
-        self.fp8: dict[str, tuple[bool, int]] = {}
-        self.fp8_layout: dict[str, str] | None = None
-        self.compiled_fp8 = None
-        self.scales_changed = False
-        self.gather: set[str] = set()
-        self.segment_starts: dict[str, torch.Tensor] = {}
-        self.shard_segments: dict[str, tuple[list[str], list[int], list[int], list[int]]] = {}
-        self.snapshot_placement = snapshot
-        # key -> (local numel, global flat offset) of this rank's shard
-        self.shards: dict[str, tuple[int, int]] = {}
-        # key -> previous wire bits: GPU tensors, or views into the host buffer
-        self.snapshot: dict[str, torch.Tensor] = {}
-        self.previous: dict[str, torch.Tensor] = {}
-        self.on_gpu: bool | None = None  # None until placed (auto: at the first delta)
-        self.host_buffer: torch.Tensor | None = None
-        self.host: dict[str, torch.Tensor] = {}
-        self.pinned = False
-        self.initialized = False
-        self.device = torch.cuda.current_device()
-        self.stream = torch.cuda.Stream()
-
-    def set_fp8_layout(self, model: nn.Module, fp8: dict[str, str], is_master: bool) -> None:
-        """Diff the tensors the engines hold in fp8 (HF name -> scale format) as fp8 bytes from now on."""
-        self.fp8_layout = fp8
-        for fqn, segs in self.segments.items():
-            modes = {fp8.get(name) for _, name in segs}
-            if modes == {None}:
-                continue
-            shapes = [self.hf_tensors[name][0] for _, name in segs]
-            num_cols = shapes[0][-1]
-            num_rows = self.numel[fqn] // num_cols
-            boundaries = [start // num_cols for start, _ in segs] + [
-                start // num_cols + shape[0] for (start, _), shape in zip(segs, shapes)
-            ]
-            if (
-                len(modes) > 1
-                or num_cols % FP8_BLOCK[1]
-                or any(len(shape) != 2 or shape[1] != num_cols for shape in shapes)
-                or any(row % FP8_BLOCK[0] and row != num_rows for row in boundaries)
-            ):
-                raise NotImplementedError(f"Sparse weight broadcast cannot quantize {fqn} to fp8 as one matrix")
-            self.fp8[fqn] = (modes.pop() == "ue8m0", num_cols)
-            for _, name in segs:
-                self.hf_tensors[name] = (self.hf_tensors[name][0], "float8_e4m3fn")
-        self.gather = self._find_gathered(model, is_master)
-        if self.fp8:
-            start = time.perf_counter()
-            self.compiled_fp8 = compile_fp8_quantizer()
-            get_logger().info(f"Compiled the fp8 quantizer ahead of time in {time.perf_counter() - start:.1f}s")
-
-    def _local_shard(self, key: str, value: torch.Tensor, is_master: bool) -> tuple[torch.Tensor, int] | None:
-        """This rank's contiguous flat slice of ``value`` and its global flat offset,
-        or None when another rank owns the elements."""
-        if not isinstance(value, DTensor):
-            return (value.detach().reshape(-1), 0) if is_master else None
-        if all(placement.is_replicate() for placement in value.placements):
-            return (value.to_local().detach().reshape(-1), 0) if is_master else None
-        local_shape, global_offset = compute_local_shape_and_global_offset(
-            value.shape, value.device_mesh, value.placements
-        )
-        if tuple(local_shape[1:]) != tuple(value.shape[1:]) or any(global_offset[1:]):
-            raise NotImplementedError(f"Sparse weight broadcast needs dim-0 shards ({key}: {value.placements})")
-        local = value.to_local().detach()[: local_shape[0]]
-        return local.reshape(-1), global_offset[0] * prod(value.shape[1:])
-
-    def _find_gathered(self, model: nn.Module, is_master: bool) -> set[str]:
-        """fp8 trainer tensors whose 128-row blocks straddle ranks on some rank."""
-        keys, straddles = [], []
-        for key, value in model.state_dict().items():
-            fqn = resolve_fqn(model, key)
-            if fqn not in self.fp8:
-                continue
-            if self.wire_dtypes[key] != torch.bfloat16:
-                raise NotImplementedError(f"Sparse weight broadcast quantizes bf16 weights to fp8, not {key}")
-            shard = self._local_shard(key, value, is_master)
-            num_cols = self.fp8[fqn][1]
-            num_rows = value.numel() // num_cols
-            first = 0 if shard is None else shard[1] // num_cols
-            end = first if shard is None else first + shard[0].numel() // num_cols
-            keys.append(key)
-            straddles.append(first % FP8_BLOCK[0] != 0 or (end % FP8_BLOCK[0] != 0 and end != num_rows))
-        flags = torch.tensor(straddles, dtype=torch.int32, device=next(model.parameters()).device)
-        if dist.is_initialized():
-            dist.all_reduce(flags, op=dist.ReduceOp.MAX)
-        return {key for key, flag in zip(keys, flags.tolist()) if flag}
-
-    @torch.no_grad()
-    def local_shards(self, model: nn.Module, is_master: bool) -> list[tuple]:
-        """This rank's ``(key, fqn, flat shard, global flat offset, gathered)`` for ``encode``, taken on the
-        training thread: FSDP swaps module parameters during forward, but the sharded storage stays put.
-        ``gathered`` is ``(previous, current)`` bf16 full tensors of fp8 tensors whose blocks straddle ranks
-        (the all-gathers run here, as ``encode`` runs no collectives)."""
-        shards = []
-        for key, value in model.state_dict().items():
-            fqn = resolve_fqn(model, key)
-            if fqn not in self.segments:
-                continue
-            shard = self._local_shard(key, value, is_master)
-            empty = shard is None or shard[0].numel() == 0
-            gathered = None
-            if key in self.gather and self.initialized:
-                if empty:
-                    value.full_tensor()  # every rank joins the all-gathers
-                    value.full_tensor()
-                else:
-                    gathered = (self._gather_previous(key, value), value.full_tensor().to(torch.bfloat16))
-            if not empty:
-                shards.append((key, fqn, shard[0], shard[1], gathered))
-        return shards
-
-    def _gather_previous(self, key: str, value: DTensor) -> torch.Tensor:
-        previous = self.snapshot[key].to(f"cuda:{self.device}").view(torch.bfloat16)
-        local_shape = (previous.numel() // prod(value.shape[1:]), *value.shape[1:])
-        return DTensor.from_local(
-            previous.view(local_shape), value.device_mesh, value.placements, shape=value.shape, stride=value.stride()
-        ).full_tensor()
-
-    def _bits(self, key: str) -> torch.dtype:
-        return bits_dtype(str(self.wire_dtypes[key]).removeprefix("torch."))
-
-    def _allocate_host(self) -> None:
-        """One host buffer for all shards (the CPU snapshot, or anchor copies of a GPU snapshot)."""
-        sizes = {key: -(-numel * self.wire_dtypes[key].itemsize // 8) * 8 for key, (numel, _) in self.shards.items()}
-        self.host_buffer = torch.from_numpy(np.empty(sum(sizes.values()), np.uint8))
-        offset = 0
-        for key, (numel, _) in self.shards.items():
-            self.host[key] = self.host_buffer[offset : offset + numel * self.wire_dtypes[key].itemsize].view(
-                self._bits(key)
-            )
-            offset += sizes[key]
-
-    def _page_lock(self) -> None:
-        """Page-lock the host buffer, only while no training step runs: registering a large buffer can
-        stall the trainer's CUDA calls."""
-        start = time.perf_counter()
-        page_lock(self.host_buffer)
-        self.pinned = True
-        get_logger().info(
-            f"Page-locked the {self.host_buffer.numel() / 2**30:.1f} GiB sparse weight host buffer in "
-            f"{time.perf_counter() - start:.1f}s"
-        )
-
-    def _place(self, peak_reserved: int) -> None:
-        """Keep the snapshot on the GPU if it fits next to the first training step's peak, else on the host."""
-        num_bytes = sum(numel * self.wire_dtypes[key].itemsize for key, (numel, _) in self.shards.items())
-        total = torch.cuda.get_device_properties(self.device).total_memory
-        headroom = total - peak_reserved - num_bytes
-        self.on_gpu = self.snapshot_placement == "auto" and headroom >= max(4 * 2**30, total // 20)
-        get_logger().info(
-            f"Sparse weight snapshot ({num_bytes / 2**30:.1f} GiB) on the {'GPU' if self.on_gpu else 'host'} "
-            f"(training peak {peak_reserved / 2**30:.1f} GiB of {total / 2**30:.1f} GiB, "
-            f"snapshot={self.snapshot_placement})"
-        )
-
-    def _fill(self, shards: list[tuple]) -> None:
-        """The first broadcast only records the weights, in the host buffer (placed at the first delta,
-        against the measured training peak). The host buffer is filled in every mode, so its pages are
-        resident before training: faulting them in later stalls the trainer."""
-        self.shards = {key: (shard.numel(), offset) for key, _, shard, offset, _ in shards}
-        self.staging = torch.empty(1 << 28, dtype=torch.uint8, pin_memory=True)
-        self._allocate_host()
-        if self.snapshot_placement == "cpu":
-            self._page_lock()
-        for key, _, shard, _, _ in shards:
-            for start in range(0, shard.numel(), self.CHUNK):
-                end = start + self.CHUNK
-                self.host[key][start:end].copy_(as_bits(shard[start:end].to(self.wire_dtypes[key])), non_blocking=True)
-        self.snapshot = dict(self.host)
-
-    def copy_snapshot_to_host(self) -> None:
-        """Copy a GPU snapshot (the last broadcast's weights) into the host buffer, for an anchor. Callers wait
-        for earlier readers of the host buffer first."""
-        assert self.on_gpu
-        with torch.cuda.stream(self.stream):
-            for key, bits in self.snapshot.items():
-                self._to_host(bits, self.host[key])
-
-    def anchor_parts(self) -> list[tuple[str, list[int], int, torch.Tensor]]:
-        """This rank's slices of the full HF checkpoint of the last broadcast, from host memory (the CPU
-        snapshot, or ``copy_snapshot_to_host``)."""
-        parts = []
-        for key, bits in self.host.items():
-            values, (numel, offset) = bits.view(self.wire_dtypes[key]), self.shards[key]
-            for start, name in self.segments[self.key_fqn[key]]:
-                begin, end = max(start, offset), min(start + prod(self.hf_tensors[name][0]), offset + numel)
-                if begin < end:
-                    parts.append((name, self.hf_tensors[name][0], begin - start, values[begin - offset : end - offset]))
-        return parts
-
-    @torch.no_grad()
-    def prepare(self, peak_reserved: int) -> None:
-        """The training thread's part of an encode, before ``encode``: place the snapshot at the first delta."""
-        self.previous = self.snapshot
-        if not self.initialized or self.on_gpu is not None:
-            return
-        self._place(peak_reserved)
-        if self.on_gpu:
-            device = f"cuda:{self.device}"
-            self.snapshot = {key: torch.empty_like(bits, device=device) for key, bits in self.host.items()}
-        elif not self.pinned:
-            self._page_lock()
-
-    @torch.no_grad()
-    def encode(
-        self, shards: list[tuple], ready: torch.cuda.Event, weights_read: threading.Event, emit: Callable
-    ) -> None:
-        """Diff the shards against the snapshot and advance it, on the side stream once ``ready`` (recorded
-        after the optimizer step) has passed; the first call only records the snapshot. Hands the delta to
-        ``emit`` in parts ``(entries, gaps, value bytes, new-weight checksum per HF tensor part)`` on the
-        host as they are encoded. Sets ``weights_read`` once done reading the weights. ``scales_changed``:
-        whether an fp8 block scale changed on this rank."""
-        torch.cuda.set_device(self.device)
-        try:
-            with torch.cuda.stream(self.stream):
-                self.stream.wait_event(ready)
-                self.key_fqn = {key: fqn for key, fqn, *_ in shards}
-                if not self.initialized:
-                    self._fill(shards)
-                    self.stream.synchronize()
-                    self.initialized = True
-                    return
-                self.scales_changed = self._diff(shards, emit)
-                self.stream.synchronize()
-        finally:
-            weights_read.set()
-
-    def _to_host(self, tensor: torch.Tensor, host: torch.Tensor | None = None) -> torch.Tensor:
-        """Copy a GPU tensor to pageable host memory (``host``, or a new tensor) through the pinned staging
-        buffer: a copy straight into pageable memory holds the CUDA driver and stalls the training step's
-        own CUDA calls."""
-        host = torch.empty(tensor.shape, dtype=tensor.dtype) if host is None else host
-        data, host_bytes = tensor.view(torch.uint8), host.view(torch.uint8)
-        for begin in range(0, data.numel(), self.staging.numel()):
-            chunk = self.staging[: min(self.staging.numel(), data.numel() - begin)]
-            chunk.copy_(data[begin : begin + chunk.numel()], non_blocking=True)
-            torch.cuda.current_stream().synchronize()
-            host_bytes[begin : begin + chunk.numel()].copy_(chunk)
-        return host
-
-    def _segments(
-        self, key: str, fqn: str, offset: int, numel: int
-    ) -> tuple[list[str], list[int], list[int], list[int]]:
-        """The HF tensor parts in this rank's shard of ``key``: names, begin/end in the shard, and the shift
-        from a shard position to the HF flat index."""
-        if key not in self.shard_segments:
-            names, begins, ends, shifts = [], [], [], []
-            for start, name in self.segments[fqn]:
-                hf_numel = prod(self.hf_tensors[name][0])
-                if start < offset + numel and start + hf_numel > offset:
-                    names.append(name)
-                    begins.append(max(start - offset, 0))
-                    ends.append(min(start - offset + hf_numel, numel))
-                    shifts.append(offset - start)
-            self.shard_segments[key] = (names, begins, ends, shifts)
-        return self.shard_segments[key]
-
-    @staticmethod
-    def _segment_sums(bits: torch.Tensor, begins: list[int], ends: list[int], firsts: list[int]) -> torch.Tensor:
-        """``bits_checksum`` of each ``bits[begin:end]`` whose first element is at HF flat index ``first``:
-        Σ bits[k]·(first + k - begin + 1), wrapping int64, from cumulative sums over blocks of
-        ``SUM_BLOCK`` elements."""
-        device = bits.device
-        begins, ends, firsts = (torch.tensor(values, device=device) for values in (begins, ends, firsts))
-        plain = torch.zeros(len(begins), dtype=torch.int64, device=device)
-        weighted = torch.zeros_like(plain)
-        for start in range(0, bits.numel(), SUM_BLOCK):
-            values = bits[start : start + SUM_BLOCK].long()
-            first, last = (bound.clamp(start, start + values.numel()) - start for bound in (begins, ends))
-            positions = torch.arange(start + 1, start + 1 + values.numel(), device=device)
-            for total, summed in ((plain, values), (weighted, values * positions)):
-                cumulative = torch.cat([summed.new_zeros(1), summed.cumsum(0)])
-                total += cumulative[last] - cumulative[first]
-        return weighted + (firsts - begins) * plain
-
-    def _diff(self, shards: list[tuple], emit: Callable) -> bool:
-        """Diff ``shards`` against ``self.previous`` on the current stream, writing the new bits to the
-        snapshot, in chunks of ``CHUNK`` elements (fp8: whole 128-row blocks). Copies the delta to the
-        host and hands it to ``emit`` every ``PART`` changes. Returns whether an fp8 block scale changed."""
-        device = torch.device(f"cuda:{self.device}")
-        pending: list[tuple] = []
-        num_pending = 0
-        scales_changed = False
-
-        def flush() -> None:
-            names = [name for part in pending for name in part[0]]
-            counts = torch.cat([part[1] for part in pending]).tolist()
-            gaps = self._to_host(torch.cat([part[2] for part in pending]))
-            values = self._to_host(torch.cat([part[3] for part in pending]))
-            checksum_names = [name for part in pending for name in part[4]]
-            checksums = dict(zip(checksum_names, torch.cat([part[5] for part in pending]).tolist()))
-            emit(([(name, count) for name, count in zip(names, counts) if count], gaps, values, checksums))
-            pending.clear()
-
-        for key, fqn, shard, offset, gathered in shards:
-            previous, out, wire = self.previous[key], self.snapshot[key], self.wire_dtypes[key]
-            segment_names, begins, ends, shifts = self._segments(key, fqn, offset, shard.numel())
-            sums = torch.zeros(len(segment_names), dtype=torch.int64, device=device)
-            changed, changed_values = [], []
-            if fqn in self.fp8:
-                ue8m0, num_cols = self.fp8[fqn]
-                if gathered is not None:
-                    chunks = [(0, shard.numel())]
-                else:
-                    rows = max(FP8_BLOCK[0], self.CHUNK // num_cols // FP8_BLOCK[0] * FP8_BLOCK[0])
-                    chunks = [(start, start + rows * num_cols) for start in range(0, shard.numel(), rows * num_cols)]
-            else:
-                chunks = [(start, start + self.CHUNK) for start in range(0, shard.numel(), self.CHUNK)]
-            for start, end in chunks:
-                end = min(end, shard.numel())
-                current = as_bits(shard[start:end].to(device, non_blocking=True).to(wire))
-                previous_bits = previous[start:end].to(device, non_blocking=True)
-                if fqn in self.fp8:
-                    if gathered is not None:
-                        previous_weights, current_weights = gathered
-                    else:
-                        previous_weights, current_weights = (
-                            previous_bits.view(torch.bfloat16),
-                            current.view(torch.bfloat16),
-                        )
-                    (previous_fp8, previous_scales), (current_fp8, current_scales) = (
-                        quantize_fp8_blocks(weights.reshape(-1, num_cols), ue8m0, self.compiled_fp8)
-                        for weights in (previous_weights, current_weights)
-                    )
-                    scales_changed |= not torch.equal(previous_scales, current_scales)
-                    if gathered is not None:
-                        previous_fp8 = previous_fp8.view(-1)[offset : offset + shard.numel()]
-                        current_fp8 = current_fp8.view(-1)[offset : offset + shard.numel()]
-                    delta = current_fp8.view(-1) - previous_fp8.view(-1)
-                    new_bits = current_fp8.view(-1)
-                else:
-                    delta = current - previous_bits
-                    new_bits = current
-                out[start:end].copy_(current, non_blocking=True)
-                index = delta.nonzero().squeeze(1)
-                changed.append(index + start)
-                changed_values.append(zigzag(delta[index]))
-                first, last = bisect_right(ends, start), bisect_left(begins, end)
-                if first < last:
-                    sums[first:last] += self._segment_sums(
-                        new_bits,
-                        [max(begin, start) - start for begin in begins[first:last]],
-                        [min(stop, end) - start for stop in ends[first:last]],
-                        [max(begin, start) + shift for begin, shift in zip(begins[first:last], shifts[first:last])],
-                    )
-            changed, changed_values = torch.cat(changed), torch.cat(changed_values)
-            if fqn not in self.segment_starts:
-                self.segment_starts[fqn] = torch.tensor([start for start, _ in self.segments[fqn]], device=device)
-            starts = self.segment_starts[fqn]
-            position = changed + offset
-            segment = torch.searchsorted(starts, position, right=True) - 1
-            entry_starts = torch.ones_like(segment, dtype=torch.bool)
-            entry_starts[1:] = segment[1:] != segment[:-1]
-            gaps = encode_gaps(position - starts[segment], entry_starts)
-            counts = torch.bincount(segment, minlength=len(starts))
-            pending.append(
-                (
-                    [name for _, name in self.segments[fqn]],
-                    counts,
-                    gaps,
-                    changed_values.view(torch.uint8),
-                    segment_names,
-                    sums,
-                )
-            )
-            num_pending += len(gaps)
-            if num_pending >= self.PART:
-                flush()
-                num_pending = 0
-        if pending:
-            flush()
-        return bool(scales_changed)
-
-
-def _write_json(path: Path, value) -> None:
-    write_file(path, json.dumps(value).encode())
-
 
 class FileSystemWeightSender(WeightSender):
     """Broadcast weights by saving a HF-compatible checkpoint (or, for LoRA
-    runs, the PEFT-shaped adapter) to a shared filesystem. With ``sparse``,
-    steps only carry the delta against the previous broadcast (see
-    ``prime_rl.utils.sparse_weights``). The delta is encoded and written in the
-    background while the next step trains: the trainer only waits, before its
-    next optimizer step, until the encode has read the weights. Full
-    checkpoints (anchors) serve engines that fall off the delta chain."""
+    runs, the PEFT-shaped adapter) to a shared filesystem."""
 
     def __init__(
         self,
@@ -503,54 +30,9 @@ def __init__(
         lora_config: LoRAConfig | None = None,
     ):
         super().__init__(output_dir, config.timeout)
-        if config.sparse and lora_config is not None:
-            raise ValueError("Sparse weight broadcast does not support LoRA")
-        self.config = config
         self.lora_config = lora_config
-        self.encoder: SparseDeltaEncoder | None = None
-        self.last_step: int | None = None
-        self.anchor_step: int | None = None
-        # Background work, each in order: encodes, delta writes and commits, periodic anchor writes
-        self.encodes = ThreadPoolExecutor(1)
-        self.writes = ThreadPoolExecutor(1)
-        self.anchors = ThreadPoolExecutor(1)
-        self.host_readers: list[Future] = []
-        self.errors: list[BaseException] = []
-        self.weights_read: threading.Event | None = None
-        self.deferred_finish = False
-        self.part_bytes = 0
-        self.gathers_full_weights = not config.sparse
         self.logger.debug("Initialized filesystem weight broadcast")
 
-    def wait_for_weight_reads(self) -> None:
-        if self.weights_read is None:
-            return
-        start = time.perf_counter()
-        self.weights_read.wait()
-        self.weights_read = None
-        self.logger.debug(f"Waited {time.perf_counter() - start:.2f}s for the sparse weight encode")
-        self._raise_errors()
-
-    def close(self) -> None:
-        for pool in (self.encodes, self.writes, self.anchors):
-            pool.shutdown(wait=True)
-        self._raise_errors()
-
-    def _raise_errors(self) -> None:
-        if self.errors:
-            raise RuntimeError("Sparse weight broadcast failed in the background") from self.errors[0]
-
-    def _submit(self, pool: ThreadPoolExecutor, fn, *args) -> Future:
-        def run() -> None:
-            try:
-                fn(*args)
-            except BaseException as error:
-                self.logger.exception(f"Sparse weight broadcast failed in {fn.__name__}")
-                self.errors.append(error)
-                raise
-
-        return pool.submit(run)
-
     def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
         if self.lora_config is not None:
             # All ranks must participate in DTensor gathering, but only master saves
@@ -570,221 +52,11 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
                     alpha=self.lora_config.alpha,
                     dropout=self.lora_config.dropout,
                 )
-            return
-
-        if not self.config.sparse:
-            self._save_full(model, step_dir)
-            return
-        self.deferred_finish = False
-        if self.world.is_master:
-            # Engines that load a step with a delta dir report their weight layout next to it.
-            (step_dir / DELTA_DIR).mkdir(exist_ok=True)
-        assert self.encoder is not None
-        first = self.last_step is None
-        base_step, self.last_step = self.last_step, step
-        if first:
-            # The snapshot is recorded (``_prepare``) before the first training step.
-            self.wait_for_weight_reads()
-            if step == 0:
-                # The run's first broadcast is the checkpoint the engines booted from.
-                if self.world.is_master:
-                    manifest = {"step": step, "base_step": None, "ranks": 0, "tensors": self.encoder.hf_tensors}
-                    _write_json(step_dir / DELTA_DIR / MANIFEST, manifest)
-                self.anchor_step = step
-                return
-            # A resumed run's first broadcast is a full checkpoint.
-            self._submit(self.encodes, self._publish_full, step, step_dir)
-            self.deferred_finish = True
-            return
-        if self.encoder.fp8_layout is None:
-            # Every engine worker reports which tensors it holds in fp8 when it loads the first broadcast.
-            layout = [self._wait_for_engine_layout(step_dir.parent) if self.world.is_master else None]
-            dist.broadcast_object_list(layout, src=0)
-            self.encoder.set_fp8_layout(model, layout[0], self.world.is_master)
-        shards = self.encoder.local_shards(model, self.world.is_master)
-        self.encoder.prepare(torch.cuda.max_memory_reserved())
-        # Periodic anchors only serve engines that fall off the chain: written in the background, each after
-        # the previous one.
-        anchor = step % self.config.anchor_interval == 0
-        ready = torch.cuda.Event()
-        ready.record()
-        self.weights_read = threading.Event()
-        args = (step, base_step, step_dir, shards, ready, self.weights_read, anchor)
-        self._submit(self.encodes, self._encode, *args)
-        self.deferred_finish = True
-
-    def _wait_for_engine_layout(self, broadcast_dir: Path) -> dict[str, str]:
-        start = time.monotonic()
-        while (layout := read_engine_layout(broadcast_dir)) is None:
-            if time.monotonic() - start > self.timeout:
-                raise TimeoutError(f"Not every engine worker reported its weight layout within {self.timeout}s")
-            time.sleep(0.1)
-        self.logger.debug(f"Waited {time.monotonic() - start:.2f}s for the engine weight layout")
-        return layout
-
-    def _prepare(self, model: nn.Module, step: int) -> None:
-        if self.lora_config is not None or not self.config.sparse:
-            return
-        self.wait_for_weight_reads()
-        self._raise_errors()
-        if self.encoder is None:
-            # The first broadcast only records the snapshot, while the master waits for the receivers.
-            self.encoder = SparseDeltaEncoder(model, self.config.snapshot)
-            shards = self.encoder.local_shards(model, self.world.is_master)
-            ready = torch.cuda.Event()
-            ready.record()
-            self.weights_read = threading.Event()
-            self._submit(self.encodes, self._record, shards, ready, self.weights_read)
-
-    def _record(self, shards, ready, weights_read: threading.Event) -> None:
-        encoder = self.encoder
-        assert encoder is not None
-        start = time.perf_counter()
-        encoder.encode(shards, ready, weights_read, emit=None)
-        num_bytes = sum(numel * encoder.wire_dtypes[key].itemsize for key, (numel, _) in encoder.shards.items())
-        self.logger.info(
-            f"Recorded the {num_bytes / 2**30:.1f} GiB sparse weight snapshot in {time.perf_counter() - start:.1f}s"
-        )
-
-    def _wait_for_host_readers(self) -> None:
-        for reader in self.host_readers:
-            reader.exception()
-        self.host_readers = []
-
-    def _anchor_parts(self) -> list:
-        """This rank's slices of the last broadcast's full checkpoint, in host memory."""
-        assert self.encoder is not None
-        if self.encoder.on_gpu:
-            self._wait_for_host_readers()
-            self.encoder.copy_snapshot_to_host()
-        return self.encoder.anchor_parts()
-
-    def _publish_full(self, step: int, step_dir: Path) -> None:
-        self.host_readers.append(self._submit(self.writes, self._write_full, step, step_dir, self._anchor_parts()))
-
-    def _encode(
-        self,
-        step: int,
-        base_step: int,
-        step_dir: Path,
-        shards,
-        ready,
-        weights_read: threading.Event,
-        anchor: bool,
-    ) -> None:
-        """Encode thread: diff v{step}, handing the delta to the writer in parts as it is encoded. With fp8
-        deltas, the parts are held until the ranks agree that no block scale changed, else the step is a
-        full checkpoint."""
-        encoder = self.encoder
-        assert encoder is not None
-        if not encoder.on_gpu:
-            self._wait_for_host_readers()  # anchors written from the host snapshot, which the encode advances
-        start = time.perf_counter()
-        held: list[tuple] = []
-        num_parts = 0
-
-        def emit(part: tuple) -> None:
-            nonlocal num_parts
-            if num_parts == 0 and self.world.is_master:
-                self._submit(self.writes, self._start_delta, step, base_step, step_dir)
-            self._submit(self.writes, self._write_part, step_dir, num_parts, part)
-            num_parts += 1
-
-        encoder.encode(shards, ready, weights_read, held.append if encoder.fp8 else emit)
-        if encoder.fp8:
-            if self._scales_changed_anywhere(step_dir, encoder.scales_changed):
-                self.logger.info(f"fp8 block scales changed at v{step}: publishing a full checkpoint")
-                self._publish_full(step, step_dir)
-                return
-            for part in held:
-                emit(part)
-        if num_parts == 0 and self.world.is_master:
-            self._submit(self.writes, self._start_delta, step, base_step, step_dir)
-        self._submit(self.writes, self._finish_delta, step, base_step, step_dir, num_parts, time.perf_counter() - start)
-        if anchor:
-            self.host_readers.append(
-                self._submit(self.anchors, self._write_anchor, step, step_dir, self._anchor_parts())
-            )
-
-    def _scales_changed_anywhere(self, step_dir: Path, changed: bool) -> bool:
-        """Agree through the step dir whether any rank's fp8 block scales changed: the encode thread runs no
-        collectives, and the training thread may be waiting for a batch that needs this step."""
-        delta_dir = step_dir / DELTA_DIR
-        delta_dir.mkdir(exist_ok=True)
-        _write_json(delta_dir / f"rank_{self.world.rank}.scales", int(changed))
-        flags = [delta_dir / f"rank_{rank}.scales" for rank in range(self.world.world_size)]
-        while not all(flag.exists() for flag in flags):
-            time.sleep(0.05)
-            revalidate(delta_dir)
-        return any(flag.read_text() == "1" for flag in flags)
-
-    def _start_delta(self, step: int, base_step: int, step_dir: Path) -> None:
-        """The manifest comes first, so engines can read part files as they appear."""
-        assert self.encoder is not None
-        manifest = {
-            "step": step,
-            "base_step": base_step,
-            "ranks": self.world.world_size,
-            "tensors": self.encoder.hf_tensors,
-        }
-        _write_json(step_dir / DELTA_DIR / MANIFEST, manifest)
-
-    def _write_part(self, step_dir: Path, index: int, part: tuple) -> None:
-        assert self.encoder is not None
-        entries, gaps, values, checksums = part
-        path = step_dir / DELTA_DIR / rank_part(self.world.rank, index)
-        self.part_bytes += write_rank_file(
-            path, entries, gaps.numpy(), values.numpy(), self.encoder.hf_tensors, checksums
-        )
-
-    def _finish_delta(self, step: int, base_step: int, step_dir: Path, num_parts: int, encode_time: float) -> None:
-        """After this rank's last part: its parts marker. The master then waits for every rank's marker and
-        commits the step. Runs on the writer thread, without collectives."""
-        delta_dir = step_dir / DELTA_DIR
-        _write_json(delta_dir / rank_parts_marker(self.world.rank), {"parts": num_parts, "bytes": self.part_bytes})
-        self.part_bytes = 0
-        if not self.world.is_master:
-            return
-        markers = [delta_dir / rank_parts_marker(rank) for rank in range(self.world.world_size)]
-        while not all(marker.exists() for marker in markers):
-            time.sleep(0.05)
-            revalidate(delta_dir)
-        WeightSender._finish(self, step, step_dir)
-        num_bytes = sum(json.loads(marker.read_text())["bytes"] for marker in markers)
-        self.logger.info(
-            f"Sparse weight delta v{step} (base v{base_step}): {num_bytes / 2**20:.1f} MiB, "
-            f"encode {encode_time:.2f}s (parts written as encoded)"
-        )
-
-    def _write_full(self, step: int, step_dir: Path, parts: list) -> None:
-        """Publish v{step} as a full checkpoint: every rank writes its slices, the master commits."""
-        self._write_anchor(step, step_dir, parts)
-        if self.world.is_master:
-            WeightSender._finish(self, step, step_dir)
-
-    def _write_anchor(self, step: int, step_dir: Path, parts: list) -> None:
-        start = time.perf_counter()
-        write_anchor_shards(parts, step_dir, self.world.rank, self.world.world_size)
-        self.logger.debug(f"Wrote the {step_dir.name} anchor in {time.perf_counter() - start:.1f}s")
-        if self.world.is_master:
-            while not anchor_complete(step_dir):
-                time.sleep(0.1)
-            self.anchor_step = max(self.anchor_step or 0, step)
-
-    def _finish(self, step: int, step_dir: Path) -> None:
-        # A delta or full step is committed by the writer thread once every rank's files are on disk.
-        if not self.deferred_finish:
-            super()._finish(step, step_dir)
-
-    def _save_full(self, model: nn.Module, step_dir: Path) -> None:
-        state_dict = gather_weights_parallel(model)
-        state_dict = convert_state_dict_to_hf(model, state_dict)
-        self.logger.debug(f"Saving weights to {step_dir}")
-        save_state_dict_parallel(state_dict, step_dir)
-
-    def _clean(self, step: int) -> None:
-        # Sparse chains replay from their anchor, so keep everything since it.
-        super()._clean(step if self.anchor_step is None else min(step, self.anchor_step + 1))
+        else:
+            state_dict = gather_weights_parallel(model)
+            state_dict = convert_state_dict_to_hf(model, state_dict)
+            self.logger.debug(f"Saving weights to {step_dir}")
+            save_state_dict_parallel(state_dict, step_dir)
 
 
 class FileSystemWeightReceiver(WeightReceiver):
@@ -812,5 +84,5 @@ async def receive(self, step: int) -> None:
             revalidate(weights_dir)
         if (weights_dir / "adapter_config.json").exists():
             await load_lora_adapter(self.admin_plane, self.model_name, weights_dir)
-            return
-        await self.admin_plane.update_weights(weights_dir, transport="filesystem", step=step)
+        else:
+            await self.admin_plane.update_weights(weights_dir, transport="filesystem", step=step)
diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
new file mode 100644
index 00000000000..5540a795c62
--- /dev/null
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -0,0 +1,734 @@
+import json
+import threading
+import time
+from bisect import bisect_left, bisect_right
+from collections import defaultdict
+from collections.abc import Callable
+from concurrent.futures import Future, ThreadPoolExecutor
+from math import prod
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from torch.distributed.tensor import DTensor
+from torch.distributed.tensor._utils import compute_local_shape_and_global_offset
+
+from prime_rl.configs.shared import FileSystemWeightBroadcastConfig
+from prime_rl.transports.weights.base import WeightSender
+from prime_rl.transports.weights.filesystem import FileSystemWeightSender
+from prime_rl.utils.logger import get_logger
+from prime_rl.utils.sparse_weights import (
+    DELTA_DIR,
+    FP8_BLOCK,
+    MANIFEST,
+    anchor_complete,
+    as_bits,
+    bits_dtype,
+    compile_fp8_quantizer,
+    encode_gaps,
+    page_lock,
+    quantize_fp8_blocks,
+    rank_part,
+    rank_parts_marker,
+    read_engine_layout,
+    revalidate,
+    write_anchor_shards,
+    write_file,
+    write_rank_file,
+    zigzag,
+)
+from prime_rl.utils.weights import convert_state_dict_to_hf, resolve_fqn, resolve_wire_dtype
+
+# Elements per block of the position-weighted checksums (int64 temporaries of 8 B each)
+SUM_BLOCK = 1 << 22
+
+
+class SparseDeltaEncoder:
+    """Rank-local deltas of the trainer's weight shards, in HF coordinates.
+
+    Every rank keeps the wire-dtype bits of its own shards from the previous
+    broadcast (the snapshot) and diffs against them, so no weights are gathered.
+    The prime->HF conversion is replayed once on meta tensors: every HF tensor
+    must come out as a contiguous view of one trainer tensor (renames, expert
+    unstacking), which maps a changed trainer element to its HF name and index.
+
+    ``encode`` runs on a background thread and a side CUDA stream while the
+    next step trains, without collectives; it only reads the weights before
+    ``weights_read`` is set, and hands the delta to the writer in parts as it
+    goes. The snapshot lives on the GPU when it fits next to the first training
+    step's peak memory, else in one page-locked host buffer (locked only between
+    training steps: registering it can stall the trainer's CUDA calls). With a
+    GPU snapshot, the same host buffer (pageable) holds anchor copies instead, so
+    a rank has at most one shard-sized host copy.
+
+    Tensors the engines hold in fp8 (``set_fp8_layout``: HF name -> scale format,
+    from the engine layout) are quantized with vLLM's own per-block quantizer,
+    previous and current weights alike, and diffed as fp8 bytes. Their 128-row
+    blocks must not straddle trainer ranks, else the trainer tensor is gathered.
+    The quantizer is compiled ahead of time, so it can run on the encode thread
+    (columns a multiple of 128). A step that changes any block scale is
+    published as a full checkpoint.
+    """
+
+    # Elements per transient chunk of the diff (bf16: 128 MiB per temporary)
+    CHUNK = 1 << 26
+    # Changes per part handed to the writer
+    PART = 1 << 26
+
+    def __init__(self, model: nn.Module, snapshot: str):
+        keep_in_fp32 = getattr(model, "keep_in_fp32_for_weight_transfer", None)
+        self.wire_dtypes: dict[str, torch.dtype] = {}
+        meta: dict[str, torch.Tensor] = {}
+        for key, value in model.state_dict().items():
+            dtype = resolve_wire_dtype(keep_in_fp32, key, torch.bfloat16) if isinstance(value, DTensor) else value.dtype
+            self.wire_dtypes[key] = dtype
+            meta[resolve_fqn(model, key)] = torch.empty(value.shape, dtype=dtype, device="meta")
+        owner = {id(tensor): fqn for fqn, tensor in meta.items()}
+
+        segments = defaultdict(list)
+        self.hf_tensors: dict[str, tuple[list[int], str]] = {}
+        for name, tensor in convert_state_dict_to_hf(model, dict(meta)).items():
+            base = tensor if tensor._base is None else tensor._base
+            if id(base) not in owner or not tensor.is_contiguous():
+                raise NotImplementedError(f"Sparse weight broadcast cannot map {name} to a contiguous trainer slice")
+            if tensor.numel() >= 2**32:
+                raise NotImplementedError(f"Sparse weight broadcast indexes tensors below 2^32 elements ({name})")
+            segments[owner[id(base)]].append((tensor.storage_offset(), name))
+            self.hf_tensors[name] = (list(tensor.shape), str(tensor.dtype).removeprefix("torch."))
+        self.segments = {fqn: sorted(segs) for fqn, segs in segments.items()}
+        self.numel = {fqn: tensor.numel() for fqn, tensor in meta.items()}
+        # fqn -> (ue8m0 scales, columns) for trainer tensors held in fp8; quantized as one [rows, columns] matrix
+        self.fp8: dict[str, tuple[bool, int]] = {}
+        self.fp8_layout: dict[str, str] | None = None
+        self.compiled_fp8 = None
+        self.scales_changed = False
+        self.gather: set[str] = set()
+        self.segment_starts: dict[str, torch.Tensor] = {}
+        self.shard_segments: dict[str, tuple[list[str], list[int], list[int], list[int]]] = {}
+        self.snapshot_placement = snapshot
+        # key -> (local numel, global flat offset) of this rank's shard
+        self.shards: dict[str, tuple[int, int]] = {}
+        # key -> previous wire bits: GPU tensors, or views into the host buffer
+        self.snapshot: dict[str, torch.Tensor] = {}
+        self.previous: dict[str, torch.Tensor] = {}
+        self.on_gpu: bool | None = None  # None until placed (auto: at the first delta)
+        self.host_buffer: torch.Tensor | None = None
+        self.host: dict[str, torch.Tensor] = {}
+        self.pinned = False
+        self.initialized = False
+        self.device = torch.cuda.current_device()
+        self.stream = torch.cuda.Stream()
+
+    def set_fp8_layout(self, model: nn.Module, fp8: dict[str, str], is_master: bool) -> None:
+        """Diff the tensors the engines hold in fp8 (HF name -> scale format) as fp8 bytes from now on."""
+        self.fp8_layout = fp8
+        for fqn, segs in self.segments.items():
+            modes = {fp8.get(name) for _, name in segs}
+            if modes == {None}:
+                continue
+            shapes = [self.hf_tensors[name][0] for _, name in segs]
+            num_cols = shapes[0][-1]
+            num_rows = self.numel[fqn] // num_cols
+            boundaries = [start // num_cols for start, _ in segs] + [
+                start // num_cols + shape[0] for (start, _), shape in zip(segs, shapes)
+            ]
+            if (
+                len(modes) > 1
+                or num_cols % FP8_BLOCK[1]
+                or any(len(shape) != 2 or shape[1] != num_cols for shape in shapes)
+                or any(row % FP8_BLOCK[0] and row != num_rows for row in boundaries)
+            ):
+                raise NotImplementedError(f"Sparse weight broadcast cannot quantize {fqn} to fp8 as one matrix")
+            self.fp8[fqn] = (modes.pop() == "ue8m0", num_cols)
+            for _, name in segs:
+                self.hf_tensors[name] = (self.hf_tensors[name][0], "float8_e4m3fn")
+        self.gather = self._find_gathered(model, is_master)
+        if self.fp8:
+            start = time.perf_counter()
+            self.compiled_fp8 = compile_fp8_quantizer()
+            get_logger().info(f"Compiled the fp8 quantizer ahead of time in {time.perf_counter() - start:.1f}s")
+
+    def _local_shard(self, key: str, value: torch.Tensor, is_master: bool) -> tuple[torch.Tensor, int] | None:
+        """This rank's contiguous flat slice of ``value`` and its global flat offset,
+        or None when another rank owns the elements."""
+        if not isinstance(value, DTensor):
+            return (value.detach().reshape(-1), 0) if is_master else None
+        if all(placement.is_replicate() for placement in value.placements):
+            return (value.to_local().detach().reshape(-1), 0) if is_master else None
+        local_shape, global_offset = compute_local_shape_and_global_offset(
+            value.shape, value.device_mesh, value.placements
+        )
+        if tuple(local_shape[1:]) != tuple(value.shape[1:]) or any(global_offset[1:]):
+            raise NotImplementedError(f"Sparse weight broadcast needs dim-0 shards ({key}: {value.placements})")
+        local = value.to_local().detach()[: local_shape[0]]
+        return local.reshape(-1), global_offset[0] * prod(value.shape[1:])
+
+    def _find_gathered(self, model: nn.Module, is_master: bool) -> set[str]:
+        """fp8 trainer tensors whose 128-row blocks straddle ranks on some rank."""
+        keys, straddles = [], []
+        for key, value in model.state_dict().items():
+            fqn = resolve_fqn(model, key)
+            if fqn not in self.fp8:
+                continue
+            if self.wire_dtypes[key] != torch.bfloat16:
+                raise NotImplementedError(f"Sparse weight broadcast quantizes bf16 weights to fp8, not {key}")
+            shard = self._local_shard(key, value, is_master)
+            num_cols = self.fp8[fqn][1]
+            num_rows = value.numel() // num_cols
+            first = 0 if shard is None else shard[1] // num_cols
+            end = first if shard is None else first + shard[0].numel() // num_cols
+            keys.append(key)
+            straddles.append(first % FP8_BLOCK[0] != 0 or (end % FP8_BLOCK[0] != 0 and end != num_rows))
+        flags = torch.tensor(straddles, dtype=torch.int32, device=next(model.parameters()).device)
+        if dist.is_initialized():
+            dist.all_reduce(flags, op=dist.ReduceOp.MAX)
+        return {key for key, flag in zip(keys, flags.tolist()) if flag}
+
+    @torch.no_grad()
+    def local_shards(self, model: nn.Module, is_master: bool) -> list[tuple]:
+        """This rank's ``(key, fqn, flat shard, global flat offset, gathered)`` for ``encode``, taken on the
+        training thread: FSDP swaps module parameters during forward, but the sharded storage stays put.
+        ``gathered`` is ``(previous, current)`` bf16 full tensors of fp8 tensors whose blocks straddle ranks
+        (the all-gathers run here, as ``encode`` runs no collectives)."""
+        shards = []
+        for key, value in model.state_dict().items():
+            fqn = resolve_fqn(model, key)
+            if fqn not in self.segments:
+                continue
+            shard = self._local_shard(key, value, is_master)
+            empty = shard is None or shard[0].numel() == 0
+            gathered = None
+            if key in self.gather and self.initialized:
+                if empty:
+                    value.full_tensor()  # every rank joins the all-gathers
+                    value.full_tensor()
+                else:
+                    gathered = (self._gather_previous(key, value), value.full_tensor().to(torch.bfloat16))
+            if not empty:
+                shards.append((key, fqn, shard[0], shard[1], gathered))
+        return shards
+
+    def _gather_previous(self, key: str, value: DTensor) -> torch.Tensor:
+        previous = self.snapshot[key].to(f"cuda:{self.device}").view(torch.bfloat16)
+        local_shape = (previous.numel() // prod(value.shape[1:]), *value.shape[1:])
+        return DTensor.from_local(
+            previous.view(local_shape), value.device_mesh, value.placements, shape=value.shape, stride=value.stride()
+        ).full_tensor()
+
+    def _bits(self, key: str) -> torch.dtype:
+        return bits_dtype(str(self.wire_dtypes[key]).removeprefix("torch."))
+
+    def _allocate_host(self) -> None:
+        """One host buffer for all shards (the CPU snapshot, or anchor copies of a GPU snapshot)."""
+        sizes = {key: -(-numel * self.wire_dtypes[key].itemsize // 8) * 8 for key, (numel, _) in self.shards.items()}
+        self.host_buffer = torch.from_numpy(np.empty(sum(sizes.values()), np.uint8))
+        offset = 0
+        for key, (numel, _) in self.shards.items():
+            self.host[key] = self.host_buffer[offset : offset + numel * self.wire_dtypes[key].itemsize].view(
+                self._bits(key)
+            )
+            offset += sizes[key]
+
+    def _page_lock(self) -> None:
+        """Page-lock the host buffer, only while no training step runs: registering a large buffer can
+        stall the trainer's CUDA calls."""
+        start = time.perf_counter()
+        page_lock(self.host_buffer)
+        self.pinned = True
+        get_logger().info(
+            f"Page-locked the {self.host_buffer.numel() / 2**30:.1f} GiB sparse weight host buffer in "
+            f"{time.perf_counter() - start:.1f}s"
+        )
+
+    def _place(self, peak_reserved: int) -> None:
+        """Keep the snapshot on the GPU if it fits next to the first training step's peak, else on the host."""
+        num_bytes = sum(numel * self.wire_dtypes[key].itemsize for key, (numel, _) in self.shards.items())
+        total = torch.cuda.get_device_properties(self.device).total_memory
+        headroom = total - peak_reserved - num_bytes
+        self.on_gpu = self.snapshot_placement == "auto" and headroom >= max(4 * 2**30, total // 20)
+        get_logger().info(
+            f"Sparse weight snapshot ({num_bytes / 2**30:.1f} GiB) on the {'GPU' if self.on_gpu else 'host'} "
+            f"(training peak {peak_reserved / 2**30:.1f} GiB of {total / 2**30:.1f} GiB, "
+            f"snapshot={self.snapshot_placement})"
+        )
+
+    def _fill(self, shards: list[tuple]) -> None:
+        """The first broadcast only records the weights, in the host buffer (placed at the first delta,
+        against the measured training peak). The host buffer is filled in every mode, so its pages are
+        resident before training: faulting them in later stalls the trainer."""
+        self.shards = {key: (shard.numel(), offset) for key, _, shard, offset, _ in shards}
+        self.staging = torch.empty(1 << 28, dtype=torch.uint8, pin_memory=True)
+        self._allocate_host()
+        if self.snapshot_placement == "cpu":
+            self._page_lock()
+        for key, _, shard, _, _ in shards:
+            for start in range(0, shard.numel(), self.CHUNK):
+                end = start + self.CHUNK
+                self.host[key][start:end].copy_(as_bits(shard[start:end].to(self.wire_dtypes[key])), non_blocking=True)
+        self.snapshot = dict(self.host)
+
+    def copy_snapshot_to_host(self) -> None:
+        """Copy a GPU snapshot (the last broadcast's weights) into the host buffer, for an anchor. Callers wait
+        for earlier readers of the host buffer first. The pageable copy stalls the training step while it
+        runs; a copy through the pinned staging buffer lets the step run alongside it, but then raised the
+        step's peak memory by 14 GiB (Qwen3-30B)."""
+        assert self.on_gpu
+        with torch.cuda.stream(self.stream):
+            for key, bits in self.snapshot.items():
+                self.host[key].copy_(bits)
+
+    def anchor_parts(self) -> list[tuple[str, list[int], int, torch.Tensor]]:
+        """This rank's slices of the full HF checkpoint of the last broadcast, from host memory (the CPU
+        snapshot, or ``copy_snapshot_to_host``)."""
+        parts = []
+        for key, bits in self.host.items():
+            values, (numel, offset) = bits.view(self.wire_dtypes[key]), self.shards[key]
+            for start, name in self.segments[self.key_fqn[key]]:
+                begin, end = max(start, offset), min(start + prod(self.hf_tensors[name][0]), offset + numel)
+                if begin < end:
+                    parts.append((name, self.hf_tensors[name][0], begin - start, values[begin - offset : end - offset]))
+        return parts
+
+    @torch.no_grad()
+    def prepare(self, peak_reserved: int) -> None:
+        """The training thread's part of an encode, before ``encode``: place the snapshot at the first delta."""
+        self.previous = self.snapshot
+        if not self.initialized or self.on_gpu is not None:
+            return
+        self._place(peak_reserved)
+        if self.on_gpu:
+            device = f"cuda:{self.device}"
+            self.snapshot = {key: torch.empty_like(bits, device=device) for key, bits in self.host.items()}
+        elif not self.pinned:
+            self._page_lock()
+
+    @torch.no_grad()
+    def encode(
+        self, shards: list[tuple], ready: torch.cuda.Event, weights_read: threading.Event, emit: Callable
+    ) -> None:
+        """Diff the shards against the snapshot and advance it, on the side stream once ``ready`` (recorded
+        after the optimizer step) has passed; the first call only records the snapshot. Hands the delta to
+        ``emit`` in parts ``(entries, gaps, value bytes, new-weight checksum per HF tensor part)`` on the
+        host as they are encoded. Sets ``weights_read`` once done reading the weights. ``scales_changed``:
+        whether an fp8 block scale changed on this rank."""
+        torch.cuda.set_device(self.device)
+        try:
+            with torch.cuda.stream(self.stream):
+                self.stream.wait_event(ready)
+                self.key_fqn = {key: fqn for key, fqn, *_ in shards}
+                if not self.initialized:
+                    self._fill(shards)
+                    self.stream.synchronize()
+                    self.initialized = True
+                    return
+                self.scales_changed = self._diff(shards, emit)
+                self.stream.synchronize()
+        finally:
+            weights_read.set()
+
+    def _to_host(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Copy a delta part to pageable host memory through the pinned staging buffer: a copy straight
+        into pageable memory holds the CUDA driver and stalls the training step's own CUDA calls."""
+        data, host = tensor.view(torch.uint8), torch.empty(tensor.shape, dtype=tensor.dtype)
+        host_bytes = host.view(torch.uint8)
+        for begin in range(0, data.numel(), self.staging.numel()):
+            chunk = self.staging[: min(self.staging.numel(), data.numel() - begin)]
+            chunk.copy_(data[begin : begin + chunk.numel()], non_blocking=True)
+            torch.cuda.current_stream().synchronize()
+            host_bytes[begin : begin + chunk.numel()].copy_(chunk)
+        return host
+
+    def _segments(
+        self, key: str, fqn: str, offset: int, numel: int
+    ) -> tuple[list[str], list[int], list[int], list[int]]:
+        """The HF tensor parts in this rank's shard of ``key``: names, begin/end in the shard, and the shift
+        from a shard position to the HF flat index."""
+        if key not in self.shard_segments:
+            names, begins, ends, shifts = [], [], [], []
+            for start, name in self.segments[fqn]:
+                hf_numel = prod(self.hf_tensors[name][0])
+                if start < offset + numel and start + hf_numel > offset:
+                    names.append(name)
+                    begins.append(max(start - offset, 0))
+                    ends.append(min(start - offset + hf_numel, numel))
+                    shifts.append(offset - start)
+            self.shard_segments[key] = (names, begins, ends, shifts)
+        return self.shard_segments[key]
+
+    @staticmethod
+    def _segment_sums(bits: torch.Tensor, begins: list[int], ends: list[int], firsts: list[int]) -> torch.Tensor:
+        """``bits_checksum`` of each ``bits[begin:end]`` whose first element is at HF flat index ``first``:
+        Σ bits[k]·(first + k - begin + 1), wrapping int64, from cumulative sums over blocks of
+        ``SUM_BLOCK`` elements."""
+        device = bits.device
+        begins, ends, firsts = (torch.tensor(values, device=device) for values in (begins, ends, firsts))
+        plain = torch.zeros(len(begins), dtype=torch.int64, device=device)
+        weighted = torch.zeros_like(plain)
+        for start in range(0, bits.numel(), SUM_BLOCK):
+            values = bits[start : start + SUM_BLOCK].long()
+            first, last = (bound.clamp(start, start + values.numel()) - start for bound in (begins, ends))
+            positions = torch.arange(start + 1, start + 1 + values.numel(), device=device)
+            for total, summed in ((plain, values), (weighted, values * positions)):
+                cumulative = torch.cat([summed.new_zeros(1), summed.cumsum(0)])
+                total += cumulative[last] - cumulative[first]
+        return weighted + (firsts - begins) * plain
+
+    def _diff(self, shards: list[tuple], emit: Callable) -> bool:
+        """Diff ``shards`` against ``self.previous`` on the current stream, writing the new bits to the
+        snapshot, in chunks of ``CHUNK`` elements (fp8: whole 128-row blocks). Copies the delta to the
+        host and hands it to ``emit`` every ``PART`` changes. Returns whether an fp8 block scale changed."""
+        device = torch.device(f"cuda:{self.device}")
+        pending: list[tuple] = []
+        num_pending = 0
+        scales_changed = False
+
+        def flush() -> None:
+            names = [name for part in pending for name in part[0]]
+            counts = torch.cat([part[1] for part in pending]).tolist()
+            gaps = self._to_host(torch.cat([part[2] for part in pending]))
+            values = self._to_host(torch.cat([part[3] for part in pending]))
+            checksum_names = [name for part in pending for name in part[4]]
+            checksums = dict(zip(checksum_names, torch.cat([part[5] for part in pending]).tolist()))
+            emit(([(name, count) for name, count in zip(names, counts) if count], gaps, values, checksums))
+            pending.clear()
+
+        for key, fqn, shard, offset, gathered in shards:
+            previous, out, wire = self.previous[key], self.snapshot[key], self.wire_dtypes[key]
+            segment_names, begins, ends, shifts = self._segments(key, fqn, offset, shard.numel())
+            sums = torch.zeros(len(segment_names), dtype=torch.int64, device=device)
+            changed, changed_values = [], []
+            if fqn in self.fp8:
+                ue8m0, num_cols = self.fp8[fqn]
+                if gathered is not None:
+                    chunks = [(0, shard.numel())]
+                else:
+                    rows = max(FP8_BLOCK[0], self.CHUNK // num_cols // FP8_BLOCK[0] * FP8_BLOCK[0])
+                    chunks = [(start, start + rows * num_cols) for start in range(0, shard.numel(), rows * num_cols)]
+            else:
+                chunks = [(start, start + self.CHUNK) for start in range(0, shard.numel(), self.CHUNK)]
+            for start, end in chunks:
+                end = min(end, shard.numel())
+                current = as_bits(shard[start:end].to(device, non_blocking=True).to(wire))
+                previous_bits = previous[start:end].to(device, non_blocking=True)
+                if fqn in self.fp8:
+                    if gathered is not None:
+                        previous_weights, current_weights = gathered
+                    else:
+                        previous_weights, current_weights = (
+                            previous_bits.view(torch.bfloat16),
+                            current.view(torch.bfloat16),
+                        )
+                    (previous_fp8, previous_scales), (current_fp8, current_scales) = (
+                        quantize_fp8_blocks(weights.reshape(-1, num_cols), ue8m0, self.compiled_fp8)
+                        for weights in (previous_weights, current_weights)
+                    )
+                    scales_changed |= not torch.equal(previous_scales, current_scales)
+                    if gathered is not None:
+                        previous_fp8 = previous_fp8.view(-1)[offset : offset + shard.numel()]
+                        current_fp8 = current_fp8.view(-1)[offset : offset + shard.numel()]
+                    delta = current_fp8.view(-1) - previous_fp8.view(-1)
+                    new_bits = current_fp8.view(-1)
+                else:
+                    delta = current - previous_bits
+                    new_bits = current
+                out[start:end].copy_(current, non_blocking=True)
+                index = delta.nonzero().squeeze(1)
+                changed.append(index + start)
+                changed_values.append(zigzag(delta[index]))
+                first, last = bisect_right(ends, start), bisect_left(begins, end)
+                if first < last:
+                    sums[first:last] += self._segment_sums(
+                        new_bits,
+                        [max(begin, start) - start for begin in begins[first:last]],
+                        [min(stop, end) - start for stop in ends[first:last]],
+                        [max(begin, start) + shift for begin, shift in zip(begins[first:last], shifts[first:last])],
+                    )
+            changed, changed_values = torch.cat(changed), torch.cat(changed_values)
+            if fqn not in self.segment_starts:
+                self.segment_starts[fqn] = torch.tensor([start for start, _ in self.segments[fqn]], device=device)
+            starts = self.segment_starts[fqn]
+            position = changed + offset
+            segment = torch.searchsorted(starts, position, right=True) - 1
+            entry_starts = torch.ones_like(segment, dtype=torch.bool)
+            entry_starts[1:] = segment[1:] != segment[:-1]
+            gaps = encode_gaps(position - starts[segment], entry_starts)
+            counts = torch.bincount(segment, minlength=len(starts))
+            pending.append(
+                (
+                    [name for _, name in self.segments[fqn]],
+                    counts,
+                    gaps,
+                    changed_values.view(torch.uint8),
+                    segment_names,
+                    sums,
+                )
+            )
+            num_pending += len(gaps)
+            if num_pending >= self.PART:
+                flush()
+                num_pending = 0
+        if pending:
+            flush()
+        return bool(scales_changed)
+
+
+def _write_json(path: Path, value) -> None:
+    write_file(path, json.dumps(value).encode())
+
+
+class SparseFileSystemWeightSender(FileSystemWeightSender):
+    """Filesystem broadcast where steps only carry the delta against the
+    previous broadcast (see ``prime_rl.utils.sparse_weights``). The delta is
+    encoded and written in the background while the next step trains: the
+    trainer only waits, before its next optimizer step, until the encode has
+    read the weights. Full checkpoints (anchors) serve engines that fall off
+    the delta chain. Steps are committed by the writer thread once every rank's
+    files are on disk."""
+
+    gathers_full_weights = False
+
+    def __init__(self, output_dir: Path, config: FileSystemWeightBroadcastConfig):
+        super().__init__(output_dir, config)
+        self.config = config
+        self.encoder: SparseDeltaEncoder | None = None
+        self.last_step: int | None = None
+        self.anchor_step: int | None = None
+        # Background work, each in order: encodes, delta writes and commits, periodic anchor writes
+        self.encodes = ThreadPoolExecutor(1)
+        self.writes = ThreadPoolExecutor(1)
+        self.anchors = ThreadPoolExecutor(1)
+        self.host_readers: list[Future] = []
+        self.errors: list[BaseException] = []
+        self.weights_read: threading.Event | None = None
+        self.part_bytes = 0
+
+    def wait_for_weight_reads(self) -> None:
+        if self.weights_read is None:
+            return
+        start = time.perf_counter()
+        self.weights_read.wait()
+        self.weights_read = None
+        self.logger.debug(f"Waited {time.perf_counter() - start:.2f}s for the sparse weight encode")
+        self._raise_errors()
+
+    def close(self) -> None:
+        for pool in (self.encodes, self.writes, self.anchors):
+            pool.shutdown(wait=True)
+        self._raise_errors()
+
+    def _raise_errors(self) -> None:
+        if self.errors:
+            raise RuntimeError("Sparse weight broadcast failed in the background") from self.errors[0]
+
+    def _submit(self, pool: ThreadPoolExecutor, fn, *args) -> Future:
+        def run() -> None:
+            try:
+                fn(*args)
+            except BaseException as error:
+                self.logger.exception(f"Sparse weight broadcast failed in {fn.__name__}")
+                self.errors.append(error)
+                raise
+
+        return pool.submit(run)
+
+    def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
+        if self.world.is_master:
+            # Engines that load a step with a delta dir report their weight layout next to it.
+            (step_dir / DELTA_DIR).mkdir(exist_ok=True)
+        assert self.encoder is not None
+        first = self.last_step is None
+        base_step, self.last_step = self.last_step, step
+        if first:
+            # The snapshot is recorded (``_prepare``) before the first training step.
+            self.wait_for_weight_reads()
+            if step == 0:
+                # The run's first broadcast is the checkpoint the engines booted from.
+                self.anchor_step = step
+                if self.world.is_master:
+                    manifest = {"step": step, "base_step": None, "ranks": 0, "tensors": self.encoder.hf_tensors}
+                    _write_json(step_dir / DELTA_DIR / MANIFEST, manifest)
+                    WeightSender._finish(self, step, step_dir)
+                return
+            # A resumed run's first broadcast is a full checkpoint.
+            self._submit(self.encodes, self._publish_full, step, step_dir)
+            return
+        if self.encoder.fp8_layout is None:
+            # Every engine worker reports which tensors it holds in fp8 when it loads the first broadcast.
+            layout = [self._wait_for_engine_layout(step_dir.parent) if self.world.is_master else None]
+            dist.broadcast_object_list(layout, src=0)
+            self.encoder.set_fp8_layout(model, layout[0], self.world.is_master)
+        shards = self.encoder.local_shards(model, self.world.is_master)
+        self.encoder.prepare(torch.cuda.max_memory_reserved())
+        # Periodic anchors only serve engines that fall off the chain: written in the background, each after
+        # the previous one.
+        anchor = step % self.config.anchor_interval == 0
+        ready = torch.cuda.Event()
+        ready.record()
+        self.weights_read = threading.Event()
+        args = (step, base_step, step_dir, shards, ready, self.weights_read, anchor)
+        self._submit(self.encodes, self._encode, *args)
+
+    def _wait_for_engine_layout(self, broadcast_dir: Path) -> dict[str, str]:
+        start = time.monotonic()
+        while (layout := read_engine_layout(broadcast_dir)) is None:
+            if time.monotonic() - start > self.timeout:
+                raise TimeoutError(f"Not every engine worker reported its weight layout within {self.timeout}s")
+            time.sleep(0.1)
+        self.logger.debug(f"Waited {time.monotonic() - start:.2f}s for the engine weight layout")
+        return layout
+
+    def _prepare(self, model: nn.Module, step: int) -> None:
+        self.wait_for_weight_reads()
+        self._raise_errors()
+        if self.encoder is None:
+            # The first broadcast only records the snapshot, while the master waits for the receivers.
+            self.encoder = SparseDeltaEncoder(model, self.config.snapshot)
+            shards = self.encoder.local_shards(model, self.world.is_master)
+            ready = torch.cuda.Event()
+            ready.record()
+            self.weights_read = threading.Event()
+            self._submit(self.encodes, self._record, shards, ready, self.weights_read)
+
+    def _record(self, shards, ready, weights_read: threading.Event) -> None:
+        encoder = self.encoder
+        assert encoder is not None
+        start = time.perf_counter()
+        encoder.encode(shards, ready, weights_read, emit=None)
+        num_bytes = sum(numel * encoder.wire_dtypes[key].itemsize for key, (numel, _) in encoder.shards.items())
+        self.logger.info(
+            f"Recorded the {num_bytes / 2**30:.1f} GiB sparse weight snapshot in {time.perf_counter() - start:.1f}s"
+        )
+
+    def _wait_for_host_readers(self) -> None:
+        for reader in self.host_readers:
+            reader.exception()
+        self.host_readers = []
+
+    def _anchor_parts(self) -> list:
+        """This rank's slices of the last broadcast's full checkpoint, in host memory."""
+        assert self.encoder is not None
+        if self.encoder.on_gpu:
+            self._wait_for_host_readers()
+            self.encoder.copy_snapshot_to_host()
+        return self.encoder.anchor_parts()
+
+    def _publish_full(self, step: int, step_dir: Path) -> None:
+        self.host_readers.append(self._submit(self.writes, self._write_full, step, step_dir, self._anchor_parts()))
+
+    def _encode(
+        self,
+        step: int,
+        base_step: int,
+        step_dir: Path,
+        shards,
+        ready,
+        weights_read: threading.Event,
+        anchor: bool,
+    ) -> None:
+        """Encode thread: diff v{step}, handing the delta to the writer in parts as it is encoded. With fp8
+        deltas, the parts are held until the ranks agree that no block scale changed, else the step is a
+        full checkpoint."""
+        encoder = self.encoder
+        assert encoder is not None
+        if not encoder.on_gpu:
+            self._wait_for_host_readers()  # anchors written from the host snapshot, which the encode advances
+        start = time.perf_counter()
+        held: list[tuple] = []
+        num_parts = 0
+
+        def emit(part: tuple) -> None:
+            nonlocal num_parts
+            if num_parts == 0 and self.world.is_master:
+                self._submit(self.writes, self._start_delta, step, base_step, step_dir)
+            self._submit(self.writes, self._write_part, step_dir, num_parts, part)
+            num_parts += 1
+
+        encoder.encode(shards, ready, weights_read, held.append if encoder.fp8 else emit)
+        if encoder.fp8:
+            if self._scales_changed_anywhere(step_dir, encoder.scales_changed):
+                self.logger.info(f"fp8 block scales changed at v{step}: publishing a full checkpoint")
+                self._publish_full(step, step_dir)
+                return
+            for part in held:
+                emit(part)
+        if num_parts == 0 and self.world.is_master:
+            self._submit(self.writes, self._start_delta, step, base_step, step_dir)
+        self._submit(self.writes, self._finish_delta, step, base_step, step_dir, num_parts, time.perf_counter() - start)
+        if anchor:
+            self.host_readers.append(
+                self._submit(self.anchors, self._write_anchor, step, step_dir, self._anchor_parts())
+            )
+
+    def _scales_changed_anywhere(self, step_dir: Path, changed: bool) -> bool:
+        """Agree through the step dir whether any rank's fp8 block scales changed: the encode thread runs no
+        collectives, and the training thread may be waiting for a batch that needs this step."""
+        delta_dir = step_dir / DELTA_DIR
+        delta_dir.mkdir(exist_ok=True)
+        _write_json(delta_dir / f"rank_{self.world.rank}.scales", int(changed))
+        flags = [delta_dir / f"rank_{rank}.scales" for rank in range(self.world.world_size)]
+        while not all(flag.exists() for flag in flags):
+            time.sleep(0.05)
+            revalidate(delta_dir)
+        return any(flag.read_text() == "1" for flag in flags)
+
+    def _start_delta(self, step: int, base_step: int, step_dir: Path) -> None:
+        """The manifest comes first, so engines can read part files as they appear."""
+        assert self.encoder is not None
+        manifest = {
+            "step": step,
+            "base_step": base_step,
+            "ranks": self.world.world_size,
+            "tensors": self.encoder.hf_tensors,
+        }
+        _write_json(step_dir / DELTA_DIR / MANIFEST, manifest)
+
+    def _write_part(self, step_dir: Path, index: int, part: tuple) -> None:
+        assert self.encoder is not None
+        entries, gaps, values, checksums = part
+        path = step_dir / DELTA_DIR / rank_part(self.world.rank, index)
+        self.part_bytes += write_rank_file(
+            path, entries, gaps.numpy(), values.numpy(), self.encoder.hf_tensors, checksums
+        )
+
+    def _finish_delta(self, step: int, base_step: int, step_dir: Path, num_parts: int, encode_time: float) -> None:
+        """After this rank's last part: its parts marker. The master then waits for every rank's marker and
+        commits the step. Runs on the writer thread, without collectives."""
+        delta_dir = step_dir / DELTA_DIR
+        _write_json(delta_dir / rank_parts_marker(self.world.rank), {"parts": num_parts, "bytes": self.part_bytes})
+        self.part_bytes = 0
+        if not self.world.is_master:
+            return
+        markers = [delta_dir / rank_parts_marker(rank) for rank in range(self.world.world_size)]
+        while not all(marker.exists() for marker in markers):
+            time.sleep(0.05)
+            revalidate(delta_dir)
+        WeightSender._finish(self, step, step_dir)
+        num_bytes = sum(json.loads(marker.read_text())["bytes"] for marker in markers)
+        self.logger.info(
+            f"Sparse weight delta v{step} (base v{base_step}): {num_bytes / 2**20:.1f} MiB, "
+            f"encode {encode_time:.2f}s (parts written as encoded)"
+        )
+
+    def _write_full(self, step: int, step_dir: Path, parts: list) -> None:
+        """Publish v{step} as a full checkpoint: every rank writes its slices, the master commits."""
+        self._write_anchor(step, step_dir, parts)
+        if self.world.is_master:
+            WeightSender._finish(self, step, step_dir)
+
+    def _write_anchor(self, step: int, step_dir: Path, parts: list) -> None:
+        start = time.perf_counter()
+        write_anchor_shards(parts, step_dir, self.world.rank, self.world.world_size)
+        self.logger.debug(f"Wrote the {step_dir.name} anchor in {time.perf_counter() - start:.1f}s")
+        if self.world.is_master:
+            while not anchor_complete(step_dir):
+                time.sleep(0.1)
+            self.anchor_step = max(self.anchor_step or 0, step)
+
+    def _finish(self, step: int, step_dir: Path) -> None:
+        """Committed in ``_broadcast`` at v0, else in the background (``_finish_delta``, ``_write_full``)."""
+
+    def _clean(self, step: int) -> None:
+        # Sparse chains replay from their anchor, so keep everything since it.
+        super()._clean(step if self.anchor_step is None else min(step, self.anchor_step + 1))
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index e90672e7cd2..2b134298c2e 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -280,7 +280,7 @@ def write_anchor_shards(
     for index, file_parts in enumerate(files):
         header, offset = {"__metadata__": {"shapes": json.dumps({name: shape for name, shape, _, _ in file_parts})}}, 0
         for name, _, start, values in file_parts:
-            dtype = next(name for name, dtype in SAFETENSORS_DTYPES.items() if getattr(torch, dtype) == values.dtype)
+            dtype = next(key for key, dtype in SAFETENSORS_DTYPES.items() if getattr(torch, dtype) == values.dtype)
             header[f"{name}@{start}"] = {
                 "dtype": dtype,
                 "shape": [values.numel()],

From 1d731e3f849357d93a87579c048c46078744ecc2 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Fri, 16 Oct 2026 14:51:41 +0000
Subject: [PATCH 32/88] fix(weights): finish the encode's weight reads before
 the final backward; reload after a failed update; flatten shards where they
 are read

- The offloaded optimizer step overlaps the final backward and updates the weights layer by layer, so the trainer now waits for the encode's weight reads before the final backward instead of before optimizer.step(). A slow encode read partly updated weights (GLM-4.5-Air, layer 45 experts at v9).
- An engine's step is unknown until an update succeeds, so the update after a failed one reloads the last anchor and replays the deltas instead of patching already-patched weights.
- Trainer shards are flattened where the encode reads them: flattening a non-contiguous slice copies it, and those copies stayed alive until the encode (and an anchor copy) finished.
- Anchors from a GPU snapshot are copied to the host through the pinned staging buffer.
---
 docs/scaling.md                               |  2 +-
 .../inference/vllm/worker/filesystem.py       |  6 ++-
 src/prime_rl/trainer/rl/train.py              |  8 ++--
 src/prime_rl/trainer/sft/train.py             |  8 ++--
 src/prime_rl/transports/weights/base.py       |  2 +-
 src/prime_rl/transports/weights/sparse.py     | 41 ++++++++++---------
 6 files changed, 38 insertions(+), 29 deletions(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index 9fc1508cd61..add189217ab 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -276,7 +276,7 @@ By default, the trainer and inference worker each allocate one transfer arena. S
 
 ### Sparse filesystem weight broadcast
 
-With `[weight_broadcast] type = "filesystem"` and `sparse = true`, the trainer publishes only the elements of the weights that changed since the previous broadcast, in the engine's weight format (bf16, or fp8 for engines that quantize online to `fp8_per_block`). Each trainer rank diffs its own shards against a copy of the previous broadcast (`snapshot = "auto"` keeps it on the GPU when it fits next to the first training step's peak memory, else in pinned host memory), so no weights are gathered. The diff runs on a side CUDA stream while the next step trains; the trainer only waits for it before the next optimizer step. With AdamW at RL learning rates about 1% of bf16 elements change per step. A full checkpoint (anchor) is written in the background every `anchor_interval` steps (default 25); the first broadcast is the checkpoint the engines booted from. Engines read and decode a delta while they keep serving, patch their weights in place during the pause and verify every step against trainer checksums. An engine that is not at a delta's base step reloads the last anchor and replays the deltas since. Weights the engine transforms after loading other than online `fp8_per_block`, and LoRA, are not supported.
+With `[weight_broadcast] type = "filesystem"` and `sparse = true`, the trainer publishes only the elements of the weights that changed since the previous broadcast, in the engine's weight format (bf16, or fp8 for engines that quantize online to `fp8_per_block`). Each trainer rank diffs its own shards against a copy of the previous broadcast (`snapshot = "auto"` keeps it on the GPU when it fits next to the first training step's peak memory, else in pinned host memory), so no weights are gathered. The diff runs on a side CUDA stream while the next step trains; the trainer only waits for it before the next step's final backward (which may already run the optimizer step). With AdamW at RL learning rates about 1% of bf16 elements change per step. A full checkpoint (anchor) is written in the background every `anchor_interval` steps (default 25); the first broadcast is the checkpoint the engines booted from. Engines read and decode a delta while they keep serving, patch their weights in place during the pause and verify every step against trainer checksums. An engine that is not at a delta's base step reloads the last anchor and replays the deltas since. Weights the engine transforms after loading other than online `fp8_per_block`, and LoRA, are not supported.
 
 ### Custom Templates
 
diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 7b7a60ae2a5..b7d613c9d41 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -153,6 +153,8 @@ def update_weights_from_path(self, weight_path: str) -> None:
         """Move to the broadcast step in ``weight_path`` (a full HF checkpoint and/or a sparse delta)."""
         step_dir = Path(weight_path)
         step = int(step_dir.name.removeprefix("step_"))
+        # Unknown until this update succeeds: after a failure, the next update reloads and replays
+        held_step, self.weight_step = self.weight_step, -1
         if not (step_dir / DELTA_DIR / MANIFEST).exists():
             self._load_checkpoint(step_dir)
             if (step_dir / DELTA_DIR).is_dir():
@@ -162,10 +164,10 @@ def update_weights_from_path(self, weight_path: str) -> None:
             base_step = read_manifest(step_dir)["base_step"]
             if base_step is None:
                 # The run's first broadcast: the weights the engine booted with.
-                if self.weight_step is not None:
+                if held_step is not None:
                     self._load_checkpoint(self._boot_dir())
                 self._write_engine_layout(step_dir, self._boot_dir())
-            elif base_step == self.weight_step:
+            elif base_step == held_step:
                 self._apply_delta(step_dir)
             else:
                 self._replay(step_dir)
diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py
index f97a87af978..33e9b4f0e39 100644
--- a/src/prime_rl/trainer/rl/train.py
+++ b/src/prime_rl/trainer/rl/train.py
@@ -494,7 +494,11 @@ def train(config: TrainerConfig):
 
             # Backward pass
             with maybe_record_function("backward"):
-                begin_backward(gradient_manager, final_backward=micro_step == len(micro_batches) - 1)
+                final_backward = micro_step == len(micro_batches) - 1
+                if final_backward and weight_sender is not None:
+                    # The offloaded optimizer step overlaps the final backward and updates the weights
+                    weight_sender.wait_for_weight_reads()
+                begin_backward(gradient_manager, final_backward=final_backward)
                 loss.backward()
                 finish_backward(gradient_manager)
 
@@ -574,8 +578,6 @@ def train(config: TrainerConfig):
             grad_norm = clip_grad_norm_(gradient_manager, model, config.optim.max_norm)
 
         # Update the model parameters
-        if weight_sender is not None:
-            weight_sender.wait_for_weight_reads()
         optimizer.step()
         optimizer.zero_grad()
 
diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py
index bdc90d2fae9..178ab0d9f68 100644
--- a/src/prime_rl/trainer/sft/train.py
+++ b/src/prime_rl/trainer/sft/train.py
@@ -473,7 +473,11 @@ def is_online_eval_step(step: int) -> bool:
                 scaled_loss = local_loss_sum / grad_accum_steps
 
             with maybe_record_function("backward"):
-                begin_backward(gradient_manager, final_backward=micro_step == grad_accum_steps - 1)
+                final_backward = micro_step == grad_accum_steps - 1
+                if final_backward and weight_sender is not None:
+                    # The offloaded optimizer step overlaps the final backward and updates the weights
+                    weight_sender.wait_for_weight_reads()
+                begin_backward(gradient_manager, final_backward=final_backward)
                 scaled_loss.backward()
                 finish_backward(gradient_manager)
 
@@ -519,8 +523,6 @@ def is_online_eval_step(step: int) -> bool:
             logger.debug(f"Clipping gradients with max norm {config.optim.max_norm}")
             grad_norm = clip_grad_norm_(gradient_manager, model, config.optim.max_norm)
         logger.debug("Optimizer step")
-        if weight_sender is not None:
-            weight_sender.wait_for_weight_reads()
         optimizer.step()
         optimizer.zero_grad()
 
diff --git a/src/prime_rl/transports/weights/base.py b/src/prime_rl/transports/weights/base.py
index baeeb8ef6f9..154b8b87b5a 100644
--- a/src/prime_rl/transports/weights/base.py
+++ b/src/prime_rl/transports/weights/base.py
@@ -107,7 +107,7 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
 
     def wait_for_weight_reads(self) -> None:
         """Block until background work of the last broadcast no longer reads the trainer weights. Called
-        before every optimizer step."""
+        before the final backward of every step, which may already run the (overlapped) optimizer step."""
 
     def close(self) -> None:
         """Finish the background work of the last broadcast."""
diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 5540a795c62..d759358c991 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -151,19 +151,18 @@ def set_fp8_layout(self, model: nn.Module, fp8: dict[str, str], is_master: bool)
             get_logger().info(f"Compiled the fp8 quantizer ahead of time in {time.perf_counter() - start:.1f}s")
 
     def _local_shard(self, key: str, value: torch.Tensor, is_master: bool) -> tuple[torch.Tensor, int] | None:
-        """This rank's contiguous flat slice of ``value`` and its global flat offset,
-        or None when another rank owns the elements."""
+        """This rank's dim-0 slice of ``value`` and its global flat offset, or None when another rank owns the
+        elements. Flattened only where it is read (``reshape(-1)``), as that copies a non-contiguous slice."""
         if not isinstance(value, DTensor):
-            return (value.detach().reshape(-1), 0) if is_master else None
+            return (value.detach(), 0) if is_master else None
         if all(placement.is_replicate() for placement in value.placements):
-            return (value.to_local().detach().reshape(-1), 0) if is_master else None
+            return (value.to_local().detach(), 0) if is_master else None
         local_shape, global_offset = compute_local_shape_and_global_offset(
             value.shape, value.device_mesh, value.placements
         )
         if tuple(local_shape[1:]) != tuple(value.shape[1:]) or any(global_offset[1:]):
             raise NotImplementedError(f"Sparse weight broadcast needs dim-0 shards ({key}: {value.placements})")
-        local = value.to_local().detach()[: local_shape[0]]
-        return local.reshape(-1), global_offset[0] * prod(value.shape[1:])
+        return value.to_local().detach()[: local_shape[0]], global_offset[0] * prod(value.shape[1:])
 
     def _find_gathered(self, model: nn.Module, is_master: bool) -> set[str]:
         """fp8 trainer tensors whose 128-row blocks straddle ranks on some rank."""
@@ -188,7 +187,7 @@ def _find_gathered(self, model: nn.Module, is_master: bool) -> set[str]:
 
     @torch.no_grad()
     def local_shards(self, model: nn.Module, is_master: bool) -> list[tuple]:
-        """This rank's ``(key, fqn, flat shard, global flat offset, gathered)`` for ``encode``, taken on the
+        """This rank's ``(key, fqn, shard, global flat offset, gathered)`` for ``encode``, taken on the
         training thread: FSDP swaps module parameters during forward, but the sharded storage stays put.
         ``gathered`` is ``(previous, current)`` bf16 full tensors of fp8 tensors whose blocks straddle ranks
         (the all-gathers run here, as ``encode`` runs no collectives)."""
@@ -264,20 +263,19 @@ def _fill(self, shards: list[tuple]) -> None:
         if self.snapshot_placement == "cpu":
             self._page_lock()
         for key, _, shard, _, _ in shards:
+            shard = shard.reshape(-1)
             for start in range(0, shard.numel(), self.CHUNK):
                 end = start + self.CHUNK
                 self.host[key][start:end].copy_(as_bits(shard[start:end].to(self.wire_dtypes[key])), non_blocking=True)
         self.snapshot = dict(self.host)
 
     def copy_snapshot_to_host(self) -> None:
-        """Copy a GPU snapshot (the last broadcast's weights) into the host buffer, for an anchor. Callers wait
-        for earlier readers of the host buffer first. The pageable copy stalls the training step while it
-        runs; a copy through the pinned staging buffer lets the step run alongside it, but then raised the
-        step's peak memory by 14 GiB (Qwen3-30B)."""
+        """Copy a GPU snapshot (the last broadcast's weights) into the host buffer, for an anchor, through the
+        pinned staging buffer. Callers wait for earlier readers of the host buffer first."""
         assert self.on_gpu
         with torch.cuda.stream(self.stream):
             for key, bits in self.snapshot.items():
-                self.host[key].copy_(bits)
+                self._to_host(bits, self.host[key])
 
     def anchor_parts(self) -> list[tuple[str, list[int], int, torch.Tensor]]:
         """This rank's slices of the full HF checkpoint of the last broadcast, from host memory (the CPU
@@ -328,11 +326,12 @@ def encode(
         finally:
             weights_read.set()
 
-    def _to_host(self, tensor: torch.Tensor) -> torch.Tensor:
-        """Copy a delta part to pageable host memory through the pinned staging buffer: a copy straight
-        into pageable memory holds the CUDA driver and stalls the training step's own CUDA calls."""
-        data, host = tensor.view(torch.uint8), torch.empty(tensor.shape, dtype=tensor.dtype)
-        host_bytes = host.view(torch.uint8)
+    def _to_host(self, tensor: torch.Tensor, host: torch.Tensor | None = None) -> torch.Tensor:
+        """Copy a GPU tensor to pageable host memory (``host``, or a new tensor) through the pinned staging
+        buffer: a copy straight into pageable memory holds the CUDA driver and stalls the training step's
+        own CUDA calls."""
+        host = torch.empty(tensor.shape, dtype=tensor.dtype) if host is None else host
+        data, host_bytes = tensor.view(torch.uint8), host.view(torch.uint8)
         for begin in range(0, data.numel(), self.staging.numel()):
             chunk = self.staging[: min(self.staging.numel(), data.numel() - begin)]
             chunk.copy_(data[begin : begin + chunk.numel()], non_blocking=True)
@@ -395,6 +394,7 @@ def flush() -> None:
             pending.clear()
 
         for key, fqn, shard, offset, gathered in shards:
+            shard = shard.reshape(-1)
             previous, out, wire = self.previous[key], self.snapshot[key], self.wire_dtypes[key]
             segment_names, begins, ends, shifts = self._segments(key, fqn, offset, shard.numel())
             sums = torch.zeros(len(segment_names), dtype=torch.int64, device=device)
@@ -482,8 +482,8 @@ class SparseFileSystemWeightSender(FileSystemWeightSender):
     """Filesystem broadcast where steps only carry the delta against the
     previous broadcast (see ``prime_rl.utils.sparse_weights``). The delta is
     encoded and written in the background while the next step trains: the
-    trainer only waits, before its next optimizer step, until the encode has
-    read the weights. Full checkpoints (anchors) serve engines that fall off
+    trainer only waits, before its next step's final backward (which may run
+    the optimizer step), until the encode has read the weights. Full checkpoints (anchors) serve engines that fall off
     the delta chain. Steps are committed by the writer thread once every rank's
     files are on disk."""
 
@@ -646,6 +646,9 @@ def emit(part: tuple) -> None:
             num_parts += 1
 
         encoder.encode(shards, ready, weights_read, held.append if encoder.fp8 else emit)
+        # Free the gathered fp8 tensors now, not after the anchor copy below, which may overlap the next step's
+        # memory peak.
+        shards.clear()
         if encoder.fp8:
             if self._scales_changed_anywhere(step_dir, encoder.scales_changed):
                 self.logger.info(f"fp8 block scales changed at v{step}: publishing a full checkpoint")

From 31ca464eb0bb3a36d713430f5f08cd993878633e Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Fri, 16 Oct 2026 15:28:19 +0000
Subject: [PATCH 33/88] fix(weights): wait for every anchor write before the
 host buffer is reused

`self.host_readers.append(...)` bound the list before `_anchor_parts()` reset it, so the new anchor write was never waited for. With anchors closer together than an anchor write takes, the next anchor copy overwrote the host buffer while the previous anchor was still being written (30B, anchor every 2 steps: anchors held later weights).
---
 src/prime_rl/transports/weights/sparse.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index d759358c991..b3ceea2d62a 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -615,7 +615,8 @@ def _anchor_parts(self) -> list:
         return self.encoder.anchor_parts()
 
     def _publish_full(self, step: int, step_dir: Path) -> None:
-        self.host_readers.append(self._submit(self.writes, self._write_full, step, step_dir, self._anchor_parts()))
+        parts = self._anchor_parts()  # before reading self.host_readers, which this resets
+        self.host_readers.append(self._submit(self.writes, self._write_full, step, step_dir, parts))
 
     def _encode(
         self,
@@ -660,9 +661,8 @@ def emit(part: tuple) -> None:
             self._submit(self.writes, self._start_delta, step, base_step, step_dir)
         self._submit(self.writes, self._finish_delta, step, base_step, step_dir, num_parts, time.perf_counter() - start)
         if anchor:
-            self.host_readers.append(
-                self._submit(self.anchors, self._write_anchor, step, step_dir, self._anchor_parts())
-            )
+            parts = self._anchor_parts()  # before reading self.host_readers, which this resets
+            self.host_readers.append(self._submit(self.anchors, self._write_anchor, step, step_dir, parts))
 
     def _scales_changed_anywhere(self, step_dir: Path, changed: bool) -> bool:
         """Agree through the step dir whether any rank's fp8 block scales changed: the encode thread runs no

From 54e241a36441bfae74dd8787fad2bf2005c58102 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Fri, 16 Oct 2026 17:11:48 +0000
Subject: [PATCH 34/88] perf(weights): position-weighted checksums with one
 reduction per run of equal-length segments

---
 src/prime_rl/transports/weights/sparse.py | 32 +++++++++++------------
 1 file changed, 15 insertions(+), 17 deletions(-)

diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index b3ceea2d62a..59ae768c01f 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -41,9 +41,6 @@
 )
 from prime_rl.utils.weights import convert_state_dict_to_hf, resolve_fqn, resolve_wire_dtype
 
-# Elements per block of the position-weighted checksums (int64 temporaries of 8 B each)
-SUM_BLOCK = 1 << 22
-
 
 class SparseDeltaEncoder:
     """Rank-local deltas of the trainer's weight shards, in HF coordinates.
@@ -359,20 +356,21 @@ def _segments(
     @staticmethod
     def _segment_sums(bits: torch.Tensor, begins: list[int], ends: list[int], firsts: list[int]) -> torch.Tensor:
         """``bits_checksum`` of each ``bits[begin:end]`` whose first element is at HF flat index ``first``:
-        Σ bits[k]·(first + k - begin + 1), wrapping int64, from cumulative sums over blocks of
-        ``SUM_BLOCK`` elements."""
-        device = bits.device
-        begins, ends, firsts = (torch.tensor(values, device=device) for values in (begins, ends, firsts))
-        plain = torch.zeros(len(begins), dtype=torch.int64, device=device)
-        weighted = torch.zeros_like(plain)
-        for start in range(0, bits.numel(), SUM_BLOCK):
-            values = bits[start : start + SUM_BLOCK].long()
-            first, last = (bound.clamp(start, start + values.numel()) - start for bound in (begins, ends))
-            positions = torch.arange(start + 1, start + 1 + values.numel(), device=device)
-            for total, summed in ((plain, values), (weighted, values * positions)):
-                cumulative = torch.cat([summed.new_zeros(1), summed.cumsum(0)])
-                total += cumulative[last] - cumulative[first]
-        return weighted + (firsts - begins) * plain
+        Σ bits[k]·(first + k - begin + 1), wrapping int64. One reduction per run of adjacent equal-length
+        segments (e.g. unstacked experts): few kernels, as the encode thread shares the GIL with training."""
+        sums, i = [], 0
+        while i < len(begins):
+            length, j = ends[i] - begins[i], i + 1
+            while j < len(begins) and begins[j] == ends[j - 1] and ends[j] - begins[j] == length:
+                j += 1
+            rows = bits[begins[i] : begins[i] + (j - i) * length].view(j - i, length)
+            positions = torch.arange(1, length + 1, device=bits.device)
+            sums.append(
+                (rows * positions).sum(1)
+                + rows.sum(1, dtype=torch.int64) * torch.tensor(firsts[i:j], device=bits.device)
+            )
+            i = j
+        return torch.cat(sums)
 
     def _diff(self, shards: list[tuple], emit: Callable) -> bool:
         """Diff ``shards`` against ``self.previous`` on the current stream, writing the new bits to the

From 44b0c837135d8b0835d7bb5dafad67b1be7881c9 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Fri, 16 Oct 2026 17:33:03 +0000
Subject: [PATCH 35/88] fix(weights): engines recognize a sparse run's full
 checkpoint by its anchor markers, not by an empty delta dir

---
 src/prime_rl/inference/vllm/worker/filesystem.py | 2 +-
 src/prime_rl/transports/weights/sparse.py        | 1 -
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index b7d613c9d41..c8927278a2e 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -157,7 +157,7 @@ def update_weights_from_path(self, weight_path: str) -> None:
         held_step, self.weight_step = self.weight_step, -1
         if not (step_dir / DELTA_DIR / MANIFEST).exists():
             self._load_checkpoint(step_dir)
-            if (step_dir / DELTA_DIR).is_dir():
+            if anchor_complete(step_dir):
                 # A sparse run's full checkpoint
                 self._write_engine_layout(step_dir, step_dir)
         else:
diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 59ae768c01f..228b1463889 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -533,7 +533,6 @@ def run() -> None:
 
     def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
         if self.world.is_master:
-            # Engines that load a step with a delta dir report their weight layout next to it.
             (step_dir / DELTA_DIR).mkdir(exist_ok=True)
         assert self.encoder is not None
         first = self.last_step is None

From f2bea83d599374470a19f4b5879f93490b1444c1 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Fri, 16 Oct 2026 17:53:11 +0000
Subject: [PATCH 36/88] fix(weights): revalidate NFS dirs for anchor markers
 and engine layouts; index tensors below 2^31 elements (int32 gaps)

---
 src/prime_rl/transports/weights/sparse.py | 4 ++--
 src/prime_rl/utils/sparse_weights.py      | 2 ++
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 228b1463889..6d090053901 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -90,8 +90,8 @@ def __init__(self, model: nn.Module, snapshot: str):
             base = tensor if tensor._base is None else tensor._base
             if id(base) not in owner or not tensor.is_contiguous():
                 raise NotImplementedError(f"Sparse weight broadcast cannot map {name} to a contiguous trainer slice")
-            if tensor.numel() >= 2**32:
-                raise NotImplementedError(f"Sparse weight broadcast indexes tensors below 2^32 elements ({name})")
+            if tensor.numel() >= 2**31:
+                raise NotImplementedError(f"Sparse weight broadcast indexes tensors below 2^31 elements ({name})")
             segments[owner[id(base)]].append((tensor.storage_offset(), name))
             self.hf_tensors[name] = (list(tensor.shape), str(tensor.dtype).removeprefix("torch."))
         self.segments = {fqn: sorted(segs) for fqn, segs in segments.items()}
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index 2b134298c2e..db3fc3adaff 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -336,6 +336,7 @@ def anchor_tensors(
 
 def anchor_complete(step_dir: Path) -> bool:
     """Whether ``step_dir`` holds a complete full checkpoint written by ``write_anchor_shards``."""
+    revalidate(step_dir)
     markers = list(step_dir.glob("anchor_rank*.done"))
     return bool(markers) and len(markers) == int(markers[0].read_text())
 
@@ -343,6 +344,7 @@ def anchor_complete(step_dir: Path) -> bool:
 def read_engine_layout(broadcast_dir: Path) -> dict[str, str] | None:
     """Checkpoint tensor -> fp8 scale format (``"fp32"``/``"ue8m0"``) the engines hold it in,
     or None until every engine worker has reported."""
+    revalidate(broadcast_dir / ENGINE_LAYOUT)
     layouts = [json.loads(path.read_text()) for path in (broadcast_dir / ENGINE_LAYOUT).glob("*.json")]
     if not layouts or len(layouts) < layouts[0]["world_size"]:
         return None

From 8501f911c9a193a206a367616d606086f96bc245 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Fri, 16 Oct 2026 19:07:30 +0000
Subject: [PATCH 37/88] refactor(weights): record the v0 snapshot inside the
 first broadcast (the handshake overlap saved a few seconds once)

---
 src/prime_rl/transports/weights/base.py   |  4 ---
 src/prime_rl/transports/weights/sparse.py | 42 +++++++++--------------
 2 files changed, 16 insertions(+), 30 deletions(-)

diff --git a/src/prime_rl/transports/weights/base.py b/src/prime_rl/transports/weights/base.py
index 154b8b87b5a..41d57d10886 100644
--- a/src/prime_rl/transports/weights/base.py
+++ b/src/prime_rl/transports/weights/base.py
@@ -59,7 +59,6 @@ def broadcast(self, model: nn.Module, step: int) -> None:
         """Broadcast policy v{step} to the inference pool."""
         start_time = time.perf_counter()
         step_dir = self.step_dir(step)
-        self._prepare(model, step)
         if self.world.is_master:
             # Reset per attempt so a re-broadcast (e.g. on resume) never trips
             # the consumer or the trainer on stale markers of a previous run.
@@ -97,9 +96,6 @@ def _wait_for_receiver_ready(self, step_dir: Path) -> None:
             time.sleep(0.1)
         self.logger.debug("Receiver ready, starting the transfer")
 
-    def _prepare(self, model: nn.Module, step: int) -> None:
-        """Work on v{step} that needs no consumer, started before the handshake."""
-
     @abstractmethod
     def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
         """Move v{step}'s weights to the consumer. All ranks enter only after
diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 6d090053901..d42d639850c 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -503,12 +503,11 @@ def __init__(self, output_dir: Path, config: FileSystemWeightBroadcastConfig):
         self.part_bytes = 0
 
     def wait_for_weight_reads(self) -> None:
-        if self.weights_read is None:
-            return
-        start = time.perf_counter()
-        self.weights_read.wait()
-        self.weights_read = None
-        self.logger.debug(f"Waited {time.perf_counter() - start:.2f}s for the sparse weight encode")
+        if self.weights_read is not None:
+            start = time.perf_counter()
+            self.weights_read.wait()
+            self.weights_read = None
+            self.logger.debug(f"Waited {time.perf_counter() - start:.2f}s for the sparse weight encode")
         self._raise_errors()
 
     def close(self) -> None:
@@ -532,14 +531,13 @@ def run() -> None:
         return pool.submit(run)
 
     def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
+        self.wait_for_weight_reads()
         if self.world.is_master:
             (step_dir / DELTA_DIR).mkdir(exist_ok=True)
-        assert self.encoder is not None
         first = self.last_step is None
         base_step, self.last_step = self.last_step, step
         if first:
-            # The snapshot is recorded (``_prepare``) before the first training step.
-            self.wait_for_weight_reads()
+            self._record(model)
             if step == 0:
                 # The run's first broadcast is the checkpoint the engines booted from.
                 self.anchor_step = step
@@ -576,24 +574,16 @@ def _wait_for_engine_layout(self, broadcast_dir: Path) -> dict[str, str]:
         self.logger.debug(f"Waited {time.monotonic() - start:.2f}s for the engine weight layout")
         return layout
 
-    def _prepare(self, model: nn.Module, step: int) -> None:
-        self.wait_for_weight_reads()
-        self._raise_errors()
-        if self.encoder is None:
-            # The first broadcast only records the snapshot, while the master waits for the receivers.
-            self.encoder = SparseDeltaEncoder(model, self.config.snapshot)
-            shards = self.encoder.local_shards(model, self.world.is_master)
-            ready = torch.cuda.Event()
-            ready.record()
-            self.weights_read = threading.Event()
-            self._submit(self.encodes, self._record, shards, ready, self.weights_read)
-
-    def _record(self, shards, ready, weights_read: threading.Event) -> None:
-        encoder = self.encoder
-        assert encoder is not None
+    def _record(self, model: nn.Module) -> None:
+        """The first broadcast only records the snapshot."""
         start = time.perf_counter()
-        encoder.encode(shards, ready, weights_read, emit=None)
-        num_bytes = sum(numel * encoder.wire_dtypes[key].itemsize for key, (numel, _) in encoder.shards.items())
+        self.encoder = SparseDeltaEncoder(model, self.config.snapshot)
+        ready = torch.cuda.Event()
+        ready.record()
+        self.encoder.encode(self.encoder.local_shards(model, self.world.is_master), ready, threading.Event(), emit=None)
+        num_bytes = sum(
+            numel * self.encoder.wire_dtypes[key].itemsize for key, (numel, _) in self.encoder.shards.items()
+        )
         self.logger.info(
             f"Recorded the {num_bytes / 2**30:.1f} GiB sparse weight snapshot in {time.perf_counter() - start:.1f}s"
         )

From 2ecd0c95be95b36d9637daed4c4c4fea2b569bda Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Fri, 16 Oct 2026 19:44:28 +0000
Subject: [PATCH 38/88] perf(weights): no host snapshot buffer for a GPU
 snapshot: v1 is a full checkpoint written from the snapshot through pinned
 staging; the host buffer for anchors is allocated at the first anchor

---
 src/prime_rl/transports/weights/sparse.py | 213 ++++++++++++----------
 src/prime_rl/utils/sparse_weights.py      |  21 ++-
 2 files changed, 130 insertions(+), 104 deletions(-)

diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index d42d639850c..57dbccf19f7 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -109,9 +109,9 @@ def __init__(self, model: nn.Module, snapshot: str):
         self.shards: dict[str, tuple[int, int]] = {}
         # key -> previous wire bits: GPU tensors, or views into the host buffer
         self.snapshot: dict[str, torch.Tensor] = {}
-        self.previous: dict[str, torch.Tensor] = {}
-        self.on_gpu: bool | None = None  # None until placed (auto: at the first delta)
+        self.on_gpu: bool | None = None  # None until placed, at the first broadcast after a training step
         self.host_buffer: torch.Tensor | None = None
+        self.write_staging: torch.Tensor | None = None
         self.host: dict[str, torch.Tensor] = {}
         self.pinned = False
         self.initialized = False
@@ -239,7 +239,7 @@ def _page_lock(self) -> None:
         )
 
     def _place(self, peak_reserved: int) -> None:
-        """Keep the snapshot on the GPU if it fits next to the first training step's peak, else on the host."""
+        """Keep the snapshot on the GPU if it fits next to the training step's peak, else on the host."""
         num_bytes = sum(numel * self.wire_dtypes[key].itemsize for key, (numel, _) in self.shards.items())
         total = torch.cuda.get_device_properties(self.device).total_memory
         headroom = total - peak_reserved - num_bytes
@@ -250,35 +250,26 @@ def _place(self, peak_reserved: int) -> None:
             f"snapshot={self.snapshot_placement})"
         )
 
-    def _fill(self, shards: list[tuple]) -> None:
-        """The first broadcast only records the weights, in the host buffer (placed at the first delta,
-        against the measured training peak). The host buffer is filled in every mode, so its pages are
-        resident before training: faulting them in later stalls the trainer."""
+    def track(self, shards: list[tuple]) -> None:
+        """Record which elements of each tensor this rank's shards hold."""
         self.shards = {key: (shard.numel(), offset) for key, _, shard, offset, _ in shards}
-        self.staging = torch.empty(1 << 28, dtype=torch.uint8, pin_memory=True)
-        self._allocate_host()
-        if self.snapshot_placement == "cpu":
-            self._page_lock()
-        for key, _, shard, _, _ in shards:
-            shard = shard.reshape(-1)
-            for start in range(0, shard.numel(), self.CHUNK):
-                end = start + self.CHUNK
-                self.host[key][start:end].copy_(as_bits(shard[start:end].to(self.wire_dtypes[key])), non_blocking=True)
-        self.snapshot = dict(self.host)
+        self.key_fqn = {key: fqn for key, fqn, *_ in shards}
 
     def copy_snapshot_to_host(self) -> None:
-        """Copy a GPU snapshot (the last broadcast's weights) into the host buffer, for an anchor, through the
-        pinned staging buffer. Callers wait for earlier readers of the host buffer first."""
+        """Copy a GPU snapshot (the last broadcast's weights) into the host buffer (allocated at the first
+        anchor), for an anchor, through the pinned staging buffer. Callers wait for earlier readers of the
+        host buffer first."""
         assert self.on_gpu
+        if self.host_buffer is None:
+            self._allocate_host()
         with torch.cuda.stream(self.stream):
             for key, bits in self.snapshot.items():
                 self._to_host(bits, self.host[key])
 
-    def anchor_parts(self) -> list[tuple[str, list[int], int, torch.Tensor]]:
-        """This rank's slices of the full HF checkpoint of the last broadcast, from host memory (the CPU
-        snapshot, or ``copy_snapshot_to_host``)."""
+    def anchor_parts(self, source: dict[str, torch.Tensor]) -> list[tuple[str, list[int], int, torch.Tensor]]:
+        """This rank's slices of a full HF checkpoint, from ``source`` (key -> flat wire bits of its shard)."""
         parts = []
-        for key, bits in self.host.items():
+        for key, bits in source.items():
             values, (numel, offset) = bits.view(self.wire_dtypes[key]), self.shards[key]
             for start, name in self.segments[self.key_fqn[key]]:
                 begin, end = max(start, offset), min(start + prod(self.hf_tensors[name][0]), offset + numel)
@@ -287,24 +278,30 @@ def anchor_parts(self) -> list[tuple[str, list[int], int, torch.Tensor]]:
         return parts
 
     @torch.no_grad()
-    def prepare(self, peak_reserved: int) -> None:
-        """The training thread's part of an encode, before ``encode``: place the snapshot at the first delta."""
-        self.previous = self.snapshot
-        if not self.initialized or self.on_gpu is not None:
-            return
-        self._place(peak_reserved)
-        if self.on_gpu:
-            device = f"cuda:{self.device}"
-            self.snapshot = {key: torch.empty_like(bits, device=device) for key, bits in self.host.items()}
-        elif not self.pinned:
-            self._page_lock()
+    def prepare(self, shards: list[tuple], peak_reserved: int) -> None:
+        """The training thread's part of an encode, before ``encode``. The first one places and allocates
+        the snapshot, against the measured training peak; the host buffer only holds a CPU snapshot."""
+        if not self.initialized:
+            self.track(shards)
+            self._place(peak_reserved)
+            self.staging = torch.empty(1 << 28, dtype=torch.uint8, pin_memory=True)
+            if self.on_gpu:
+                device = f"cuda:{self.device}"
+                self.snapshot = {
+                    key: torch.empty(numel, dtype=self._bits(key), device=device)
+                    for key, (numel, _) in self.shards.items()
+                }
+            else:
+                self._allocate_host()
+                self._page_lock()
+                self.snapshot = dict(self.host)
 
     @torch.no_grad()
     def encode(
         self, shards: list[tuple], ready: torch.cuda.Event, weights_read: threading.Event, emit: Callable
     ) -> None:
         """Diff the shards against the snapshot and advance it, on the side stream once ``ready`` (recorded
-        after the optimizer step) has passed; the first call only records the snapshot. Hands the delta to
+        after the optimizer step) has passed; the first call only fills the snapshot. Hands the delta to
         ``emit`` in parts ``(entries, gaps, value bytes, new-weight checksum per HF tensor part)`` on the
         host as they are encoded. Sets ``weights_read`` once done reading the weights. ``scales_changed``:
         whether an fp8 block scale changed on this rank."""
@@ -312,17 +309,35 @@ def encode(
         try:
             with torch.cuda.stream(self.stream):
                 self.stream.wait_event(ready)
-                self.key_fqn = {key: fqn for key, fqn, *_ in shards}
                 if not self.initialized:
-                    self._fill(shards)
-                    self.stream.synchronize()
+                    for key, _, shard, _, _ in shards:
+                        self.snapshot[key].copy_(
+                            as_bits(shard.reshape(-1).to(self.wire_dtypes[key])), non_blocking=True
+                        )
                     self.initialized = True
-                    return
-                self.scales_changed = self._diff(shards, emit)
+                else:
+                    self.scales_changed = self._diff(shards, emit)
                 self.stream.synchronize()
         finally:
             weights_read.set()
 
+    def stream_to_host(self, tensor: torch.Tensor):
+        """A tensor's bytes on the host; for a GPU tensor, in chunks through a pinned staging buffer, each to be
+        written before the next is taken."""
+        data = tensor.reshape(-1).view(torch.uint8)
+        if not data.is_cuda:
+            yield data.numpy()
+            return
+        if self.write_staging is None:
+            self.write_staging = torch.empty(1 << 28, dtype=torch.uint8, pin_memory=True)
+            self.write_stream = torch.cuda.Stream(self.device)
+        with torch.cuda.stream(self.write_stream):
+            for begin in range(0, data.numel(), self.write_staging.numel()):
+                chunk = self.write_staging[: min(self.write_staging.numel(), data.numel() - begin)]
+                chunk.copy_(data[begin : begin + chunk.numel()], non_blocking=True)
+                self.write_stream.synchronize()
+                yield chunk.numpy()
+
     def _to_host(self, tensor: torch.Tensor, host: torch.Tensor | None = None) -> torch.Tensor:
         """Copy a GPU tensor to pageable host memory (``host``, or a new tensor) through the pinned staging
         buffer: a copy straight into pageable memory holds the CUDA driver and stalls the training step's
@@ -393,7 +408,8 @@ def flush() -> None:
 
         for key, fqn, shard, offset, gathered in shards:
             shard = shard.reshape(-1)
-            previous, out, wire = self.previous[key], self.snapshot[key], self.wire_dtypes[key]
+            previous = out = self.snapshot[key]
+            wire = self.wire_dtypes[key]
             segment_names, begins, ends, shifts = self._segments(key, fqn, offset, shard.numel())
             sums = torch.zeros(len(segment_names), dtype=torch.int64, device=device)
             changed, changed_values = [], []
@@ -497,7 +513,10 @@ def __init__(self, output_dir: Path, config: FileSystemWeightBroadcastConfig):
         self.encodes = ThreadPoolExecutor(1)
         self.writes = ThreadPoolExecutor(1)
         self.anchors = ThreadPoolExecutor(1)
-        self.host_readers: list[Future] = []
+        # Background writes that read the snapshot (the next encode waits for them), and anchor writes from the
+        # host buffer of a GPU snapshot (the next anchor copy waits for them)
+        self.snapshot_writes: list[Future] = []
+        self.anchor_writes: list[Future] = []
         self.errors: list[BaseException] = []
         self.weights_read: threading.Event | None = None
         self.part_bytes = 0
@@ -534,35 +553,40 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
         self.wait_for_weight_reads()
         if self.world.is_master:
             (step_dir / DELTA_DIR).mkdir(exist_ok=True)
-        first = self.last_step is None
+        if self.encoder is None:
+            self.encoder = SparseDeltaEncoder(model, self.config.snapshot)
+        encoder = self.encoder
         base_step, self.last_step = self.last_step, step
-        if first:
-            self._record(model)
-            if step == 0:
-                # The run's first broadcast is the checkpoint the engines booted from.
-                self.anchor_step = step
-                if self.world.is_master:
-                    manifest = {"step": step, "base_step": None, "ranks": 0, "tensors": self.encoder.hf_tensors}
-                    _write_json(step_dir / DELTA_DIR / MANIFEST, manifest)
-                    WeightSender._finish(self, step, step_dir)
-                return
-            # A resumed run's first broadcast is a full checkpoint.
-            self._submit(self.encodes, self._publish_full, step, step_dir)
+        if base_step is None and step == 0:
+            # The run's first broadcast is the checkpoint the engines booted from.
+            self.anchor_step = step
+            if self.world.is_master:
+                manifest = {"step": step, "base_step": None, "ranks": 0, "tensors": encoder.hf_tensors}
+                _write_json(step_dir / DELTA_DIR / MANIFEST, manifest)
+                WeightSender._finish(self, step, step_dir)
             return
-        if self.encoder.fp8_layout is None:
+        if encoder.initialized and encoder.fp8_layout is None:
             # Every engine worker reports which tensors it holds in fp8 when it loads the first broadcast.
             layout = [self._wait_for_engine_layout(step_dir.parent) if self.world.is_master else None]
             dist.broadcast_object_list(layout, src=0)
-            self.encoder.set_fp8_layout(model, layout[0], self.world.is_master)
-        shards = self.encoder.local_shards(model, self.world.is_master)
-        self.encoder.prepare(torch.cuda.max_memory_reserved())
+            encoder.set_fp8_layout(model, layout[0], self.world.is_master)
+        shards = encoder.local_shards(model, self.world.is_master)
+        if base_step is None:
+            # A resumed run's startup broadcast: a full checkpoint, written from the weights before training.
+            encoder.track(shards)
+            source = {key: as_bits(shard.reshape(-1).to(encoder.wire_dtypes[key])) for key, _, shard, _, _ in shards}
+            torch.cuda.synchronize()
+            self._write_full(step, step_dir, encoder.anchor_parts(source))
+            return
+        full = not encoder.initialized
+        encoder.prepare(shards, torch.cuda.max_memory_reserved())
         # Periodic anchors only serve engines that fall off the chain: written in the background, each after
         # the previous one.
         anchor = step % self.config.anchor_interval == 0
         ready = torch.cuda.Event()
         ready.record()
         self.weights_read = threading.Event()
-        args = (step, base_step, step_dir, shards, ready, self.weights_read, anchor)
+        args = (step, base_step, step_dir, shards, ready, self.weights_read, full, anchor)
         self._submit(self.encodes, self._encode, *args)
 
     def _wait_for_engine_layout(self, broadcast_dir: Path) -> dict[str, str]:
@@ -574,36 +598,26 @@ def _wait_for_engine_layout(self, broadcast_dir: Path) -> dict[str, str]:
         self.logger.debug(f"Waited {time.monotonic() - start:.2f}s for the engine weight layout")
         return layout
 
-    def _record(self, model: nn.Module) -> None:
-        """The first broadcast only records the snapshot."""
-        start = time.perf_counter()
-        self.encoder = SparseDeltaEncoder(model, self.config.snapshot)
-        ready = torch.cuda.Event()
-        ready.record()
-        self.encoder.encode(self.encoder.local_shards(model, self.world.is_master), ready, threading.Event(), emit=None)
-        num_bytes = sum(
-            numel * self.encoder.wire_dtypes[key].itemsize for key, (numel, _) in self.encoder.shards.items()
-        )
-        self.logger.info(
-            f"Recorded the {num_bytes / 2**30:.1f} GiB sparse weight snapshot in {time.perf_counter() - start:.1f}s"
-        )
-
-    def _wait_for_host_readers(self) -> None:
-        for reader in self.host_readers:
-            reader.exception()
-        self.host_readers = []
-
-    def _anchor_parts(self) -> list:
-        """This rank's slices of the last broadcast's full checkpoint, in host memory."""
-        assert self.encoder is not None
-        if self.encoder.on_gpu:
-            self._wait_for_host_readers()
-            self.encoder.copy_snapshot_to_host()
-        return self.encoder.anchor_parts()
-
-    def _publish_full(self, step: int, step_dir: Path) -> None:
-        parts = self._anchor_parts()  # before reading self.host_readers, which this resets
-        self.host_readers.append(self._submit(self.writes, self._write_full, step, step_dir, parts))
+    @staticmethod
+    def _wait(writes: list[Future]) -> None:
+        for write in writes:
+            write.exception()
+        writes.clear()
+
+    def _write_from_snapshot(self, pool: ThreadPoolExecutor, write, step: int, step_dir: Path, full: bool) -> None:
+        """Write this rank's slices of the last broadcast's full checkpoint in the background. A periodic anchor
+        of a GPU snapshot goes through the host buffer, so the next encode does not wait for the write; other
+        writes read the snapshot itself (a CPU snapshot, or a GPU one through pinned staging chunks), and the
+        next encode waits for them."""
+        encoder = self.encoder
+        assert encoder is not None
+        if encoder.on_gpu and not full:
+            self._wait(self.anchor_writes)
+            encoder.copy_snapshot_to_host()
+            parts, writes = encoder.anchor_parts(encoder.host), self.anchor_writes
+        else:
+            parts, writes = encoder.anchor_parts(encoder.snapshot), self.snapshot_writes
+        writes.append(self._submit(pool, write, step, step_dir, parts))
 
     def _encode(
         self,
@@ -613,15 +627,21 @@ def _encode(
         shards,
         ready,
         weights_read: threading.Event,
+        full: bool,
         anchor: bool,
     ) -> None:
         """Encode thread: diff v{step}, handing the delta to the writer in parts as it is encoded. With fp8
         deltas, the parts are held until the ranks agree that no block scale changed, else the step is a
-        full checkpoint."""
+        full checkpoint. The first broadcast after a training step only fills the snapshot and is published
+        as a full checkpoint (the engines reload it once)."""
         encoder = self.encoder
         assert encoder is not None
-        if not encoder.on_gpu:
-            self._wait_for_host_readers()  # anchors written from the host snapshot, which the encode advances
+        self._wait(self.snapshot_writes)
+        if full:
+            encoder.encode(shards, ready, weights_read, emit=None)
+            shards.clear()
+            self._write_from_snapshot(self.writes, self._write_full, step, step_dir, full=True)
+            return
         start = time.perf_counter()
         held: list[tuple] = []
         num_parts = 0
@@ -640,7 +660,7 @@ def emit(part: tuple) -> None:
         if encoder.fp8:
             if self._scales_changed_anywhere(step_dir, encoder.scales_changed):
                 self.logger.info(f"fp8 block scales changed at v{step}: publishing a full checkpoint")
-                self._publish_full(step, step_dir)
+                self._write_from_snapshot(self.writes, self._write_full, step, step_dir, full=True)
                 return
             for part in held:
                 emit(part)
@@ -648,8 +668,7 @@ def emit(part: tuple) -> None:
             self._submit(self.writes, self._start_delta, step, base_step, step_dir)
         self._submit(self.writes, self._finish_delta, step, base_step, step_dir, num_parts, time.perf_counter() - start)
         if anchor:
-            parts = self._anchor_parts()  # before reading self.host_readers, which this resets
-            self.host_readers.append(self._submit(self.anchors, self._write_anchor, step, step_dir, parts))
+            self._write_from_snapshot(self.anchors, self._write_anchor, step, step_dir, full=False)
 
     def _scales_changed_anywhere(self, step_dir: Path, changed: bool) -> bool:
         """Agree through the step dir whether any rank's fp8 block scales changed: the encode thread runs no
@@ -709,7 +728,7 @@ def _write_full(self, step: int, step_dir: Path, parts: list) -> None:
 
     def _write_anchor(self, step: int, step_dir: Path, parts: list) -> None:
         start = time.perf_counter()
-        write_anchor_shards(parts, step_dir, self.world.rank, self.world.world_size)
+        write_anchor_shards(parts, step_dir, self.world.rank, self.world.world_size, self.encoder.stream_to_host)
         self.logger.debug(f"Wrote the {step_dir.name} anchor in {time.perf_counter() - start:.1f}s")
         if self.world.is_master:
             while not anchor_complete(step_dir):
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index db3fc3adaff..94e513ddfc1 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -12,13 +12,14 @@
 checksums of the new weights catch it if not.
 """
 
+import itertools
 import json
 import os
 import re
 import threading
 import time
 from collections import defaultdict
-from collections.abc import Iterable, Iterator
+from collections.abc import Callable, Iterable, Iterator
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import ExitStack
 from itertools import groupby
@@ -252,7 +253,7 @@ def _read_frame(
     return kept, np.concatenate(kept_gaps), np.concatenate(kept_values)
 
 
-def write_file(path: Path, data: bytes | list) -> None:
+def write_file(path: Path, data: bytes | Iterable) -> None:
     """Write ``data`` (bytes, or buffers written back to back) under a temporary name, then rename, so
     readers never see a partial file. Plain file I/O releases the GIL during the (possibly slow) write,
     so background writes don't stall training."""
@@ -265,11 +266,16 @@ def write_file(path: Path, data: bytes | list) -> None:
 
 
 def write_anchor_shards(
-    parts: list[tuple[str, list[int], int, torch.Tensor]], step_dir: Path, rank: int, world_size: int
+    parts: list[tuple[str, list[int], int, torch.Tensor]],
+    step_dir: Path,
+    rank: int,
+    world_size: int,
+    to_host: Callable[[torch.Tensor], Iterable] = lambda values: [values.view(torch.uint8).numpy()],
 ) -> None:
     """Write one trainer rank's share of a full checkpoint: ``(hf_name, shape, flat start, flat values)``
-    slices, stored as tensors ``"{hf_name}@{start}"`` in ~1 GiB safetensors files, straight from host
-    memory and without collectives. The ``anchor_rank{rank}.done`` marker (the world size) comes last."""
+    slices, stored as tensors ``"{hf_name}@{start}"`` in ~1 GiB safetensors files, without collectives.
+    ``to_host`` yields each slice's bytes on the host, each written before the next is taken. The
+    ``anchor_rank{rank}.done`` marker (the world size) comes last."""
     files, size = [[]], 0
     for part in parts:
         if size > 2**30:
@@ -289,8 +295,9 @@ def write_anchor_shards(
             offset += values.nbytes
         encoded = json.dumps(header).encode()
         encoded += b" " * (-len(encoded) % 8)
-        buffers = [len(encoded).to_bytes(8, "little") + encoded]
-        buffers.extend(values.view(torch.uint8).numpy() for *_, values in file_parts)
+        buffers = itertools.chain(
+            [len(encoded).to_bytes(8, "little") + encoded], (chunk for *_, v in file_parts for chunk in to_host(v))
+        )
         write_file(step_dir / f"model-rank{rank:05d}-{index:05d}.safetensors", buffers)
     write_file(step_dir / f"anchor_rank{rank}.done", str(world_size).encode())
 

From 427b749a07fd3c3f0c60bc7922700f35289e553a Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Fri, 16 Oct 2026 19:45:57 +0000
Subject: [PATCH 39/88] docs(weights): snapshot placement at v1 and the host
 anchor buffer

---
 docs/scaling.md                           |  2 +-
 src/prime_rl/transports/weights/sparse.py | 10 +++++-----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index add189217ab..c51a1aff7b3 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -276,7 +276,7 @@ By default, the trainer and inference worker each allocate one transfer arena. S
 
 ### Sparse filesystem weight broadcast
 
-With `[weight_broadcast] type = "filesystem"` and `sparse = true`, the trainer publishes only the elements of the weights that changed since the previous broadcast, in the engine's weight format (bf16, or fp8 for engines that quantize online to `fp8_per_block`). Each trainer rank diffs its own shards against a copy of the previous broadcast (`snapshot = "auto"` keeps it on the GPU when it fits next to the first training step's peak memory, else in pinned host memory), so no weights are gathered. The diff runs on a side CUDA stream while the next step trains; the trainer only waits for it before the next step's final backward (which may already run the optimizer step). With AdamW at RL learning rates about 1% of bf16 elements change per step. A full checkpoint (anchor) is written in the background every `anchor_interval` steps (default 25); the first broadcast is the checkpoint the engines booted from. Engines read and decode a delta while they keep serving, patch their weights in place during the pause and verify every step against trainer checksums. An engine that is not at a delta's base step reloads the last anchor and replays the deltas since. Weights the engine transforms after loading other than online `fp8_per_block`, and LoRA, are not supported.
+With `[weight_broadcast] type = "filesystem"` and `sparse = true`, the trainer publishes only the elements of the weights that changed since the previous broadcast, in the engine's weight format (bf16, or fp8 for engines that quantize online to `fp8_per_block`). Each trainer rank diffs its own shards against a copy of the previous broadcast (`snapshot = "auto"` keeps it on the GPU when it fits next to the first training step's peak memory, else in pinned host memory), so no weights are gathered. The diff runs on a side CUDA stream while the next step trains; the trainer only waits for it before the next step's final backward (which may already run the optimizer step). With AdamW at RL learning rates about 1% of bf16 elements change per step. A full checkpoint (anchor) is written in the background every `anchor_interval` steps (default 25); the first broadcast is the checkpoint the engines booted from, and the first broadcast after a training step is a full checkpoint, which the engines reload once. With a GPU snapshot, anchors are copied to a host buffer of the snapshot's size (2 B per parameter the rank holds), allocated at the first anchor; on hosts without that much free memory (e.g. next to a fully offloaded Muon optimizer), use a larger `anchor_interval`. Engines read and decode a delta while they keep serving, patch their weights in place during the pause and verify every step against trainer checksums. An engine that is not at a delta's base step reloads the last anchor and replays the deltas since. Weights the engine transforms after loading other than online `fp8_per_block`, and LoRA, are not supported.
 
 ### Custom Templates
 
diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 57dbccf19f7..19ff60eff56 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -54,11 +54,11 @@ class SparseDeltaEncoder:
     ``encode`` runs on a background thread and a side CUDA stream while the
     next step trains, without collectives; it only reads the weights before
     ``weights_read`` is set, and hands the delta to the writer in parts as it
-    goes. The snapshot lives on the GPU when it fits next to the first training
-    step's peak memory, else in one page-locked host buffer (locked only between
-    training steps: registering it can stall the trainer's CUDA calls). With a
-    GPU snapshot, the same host buffer (pageable) holds anchor copies instead, so
-    a rank has at most one shard-sized host copy.
+    goes. The snapshot is placed at the first broadcast after a training step: on
+    the GPU when it fits next to that step's peak memory, else in one page-locked
+    host buffer (locked only between training steps: registering it can stall the
+    trainer's CUDA calls). With a GPU snapshot, a pageable host buffer holds anchor
+    copies, allocated at the first periodic anchor.
 
     Tensors the engines hold in fp8 (``set_fp8_layout``: HF name -> scale format,
     from the engine layout) are quantized with vLLM's own per-block quantizer,

From 4b9e0d963bf064ed07343f911ac6d71690c2a328 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Fri, 16 Oct 2026 20:45:18 +0000
Subject: [PATCH 40/88] refactor(weights): drop the engine reader's NUMA
 binding (no gain in a same-node GLM-4.5-Air A/B)

---
 src/prime_rl/inference/vllm/worker/filesystem.py | 11 ++---------
 src/prime_rl/utils/sparse_weights.py             | 14 --------------
 2 files changed, 2 insertions(+), 23 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index c8927278a2e..819935cc87e 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -1,6 +1,5 @@
 import itertools
 import json
-import os
 import time
 from concurrent.futures import Future, ThreadPoolExecutor
 from pathlib import Path
@@ -29,7 +28,6 @@
     anchor_tensors,
     bits_dtype,
     decode_gaps,
-    numa_cpus,
     pin,
     read_delta,
     read_manifest,
@@ -86,15 +84,10 @@ def _read_delta(
     ) -> tuple[list[tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]], dict[str, int]]:
         """Read this worker's part of a delta into a pinned host buffer that is reused across steps
         (allocating pinned memory every step is slow on a loaded host), in chunks of up to 2^26 changes
-        of consecutive frames. Reads and copies run on the CPUs of the GPU's NUMA node."""
+        of consecutive frames."""
         global _reader_pool
-        cpus = numa_cpus(torch.cuda.current_device())
-        if cpus is not None:
-            os.sched_setaffinity(0, cpus)
         if _reader_pool is None:
-            _reader_pool = ThreadPoolExecutor(
-                16, initializer=None if cpus is None else lambda: os.sched_setaffinity(0, cpus)
-            )
+            _reader_pool = ThreadPoolExecutor(16)
         start = time.perf_counter()
         frames, checksums = read_delta(step_dir, _reader_pool, None if self.index_map is None else self.index_map.names)
         read_time = time.perf_counter() - start
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index 94e513ddfc1..567b984cbf2 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -418,20 +418,6 @@ def read_delta(
     return frames, checksums
 
 
-def numa_cpus(device: int) -> set[int] | None:
-    """The CPUs of the NUMA node the CUDA ``device`` hangs off, or None if unknown."""
-    props = torch.cuda.get_device_properties(device)
-    sysfs = Path(f"/sys/bus/pci/devices/{props.pci_domain_id:04x}:{props.pci_bus_id:02x}:{props.pci_device_id:02x}.0")
-    node = int((sysfs / "numa_node").read_text()) if (sysfs / "numa_node").exists() else -1
-    if node < 0:
-        return None
-    cpus = set()
-    for part in Path(f"/sys/devices/system/node/node{node}/cpulist").read_text().strip().split(","):
-        first, _, last = part.partition("-")
-        cpus.update(range(int(first), int(last or first) + 1))
-    return cpus
-
-
 def pin(num_bytes: int) -> torch.Tensor:
     """Page-locked host memory of exactly ``num_bytes``, outside torch's caching host allocator
     (which rounds sizes up and keeps freed blocks). Release it with ``unpin``."""

From 1a429711a57b35cd4ae9e544eb8cfcd5fcf6ed15 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Fri, 16 Oct 2026 21:08:46 +0000
Subject: [PATCH 41/88] perf(weights): v1 is a delta against the boot
 checkpoint read back from storage, not a full checkpoint

---
 docs/scaling.md                           |   2 +-
 src/prime_rl/transports/weights/sparse.py | 137 ++++++++++++++--------
 src/prime_rl/utils/sparse_weights.py      |  16 +++
 3 files changed, 105 insertions(+), 50 deletions(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index c51a1aff7b3..d722c2b2198 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -276,7 +276,7 @@ By default, the trainer and inference worker each allocate one transfer arena. S
 
 ### Sparse filesystem weight broadcast
 
-With `[weight_broadcast] type = "filesystem"` and `sparse = true`, the trainer publishes only the elements of the weights that changed since the previous broadcast, in the engine's weight format (bf16, or fp8 for engines that quantize online to `fp8_per_block`). Each trainer rank diffs its own shards against a copy of the previous broadcast (`snapshot = "auto"` keeps it on the GPU when it fits next to the first training step's peak memory, else in pinned host memory), so no weights are gathered. The diff runs on a side CUDA stream while the next step trains; the trainer only waits for it before the next step's final backward (which may already run the optimizer step). With AdamW at RL learning rates about 1% of bf16 elements change per step. A full checkpoint (anchor) is written in the background every `anchor_interval` steps (default 25); the first broadcast is the checkpoint the engines booted from, and the first broadcast after a training step is a full checkpoint, which the engines reload once. With a GPU snapshot, anchors are copied to a host buffer of the snapshot's size (2 B per parameter the rank holds), allocated at the first anchor; on hosts without that much free memory (e.g. next to a fully offloaded Muon optimizer), use a larger `anchor_interval`. Engines read and decode a delta while they keep serving, patch their weights in place during the pause and verify every step against trainer checksums. An engine that is not at a delta's base step reloads the last anchor and replays the deltas since. Weights the engine transforms after loading other than online `fp8_per_block`, and LoRA, are not supported.
+With `[weight_broadcast] type = "filesystem"` and `sparse = true`, the trainer publishes only the elements of the weights that changed since the previous broadcast, in the engine's weight format (bf16, or fp8 for engines that quantize online to `fp8_per_block`). Each trainer rank diffs its own shards against a copy of the previous broadcast (`snapshot = "auto"` keeps it on the GPU when it fits next to the first training step's peak memory, else in pinned host memory), so no weights are gathered. The diff runs on a side CUDA stream while the next step trains; the trainer only waits for it before the next step's final backward (which may already run the optimizer step). With AdamW at RL learning rates about 1% of bf16 elements change per step. A full checkpoint (anchor) is written in the background every `anchor_interval` steps (default 25); the first broadcast is the checkpoint the engines booted from, and the trainer reads it back (its own shard of each tensor) to encode the first delta. With a GPU snapshot, anchors are copied to a host buffer of the snapshot's size (2 B per parameter the rank holds), allocated at the first anchor; on hosts without that much free memory (e.g. next to a fully offloaded Muon optimizer), use a larger `anchor_interval`. Engines read and decode a delta while they keep serving, patch their weights in place during the pause and verify every step against trainer checksums. An engine that is not at a delta's base step reloads the last anchor and replays the deltas since. Weights the engine transforms after loading other than online `fp8_per_block`, and LoRA, are not supported.
 
 ### Custom Templates
 
diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 19ff60eff56..629e60cc067 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -1,4 +1,5 @@
 import json
+import os
 import threading
 import time
 from bisect import bisect_left, bisect_right
@@ -12,6 +13,7 @@
 import torch
 import torch.distributed as dist
 import torch.nn as nn
+from huggingface_hub import snapshot_download
 from torch.distributed.tensor import DTensor
 from torch.distributed.tensor._utils import compute_local_shape_and_global_offset
 
@@ -34,6 +36,7 @@
     rank_parts_marker,
     read_engine_layout,
     revalidate,
+    safetensors_index,
     write_anchor_shards,
     write_file,
     write_rank_file,
@@ -57,8 +60,9 @@ class SparseDeltaEncoder:
     goes. The snapshot is placed at the first broadcast after a training step: on
     the GPU when it fits next to that step's peak memory, else in one page-locked
     host buffer (locked only between training steps: registering it can stall the
-    trainer's CUDA calls). With a GPU snapshot, a pageable host buffer holds anchor
-    copies, allocated at the first periodic anchor.
+    trainer's CUDA calls). It starts as the checkpoint the engines booted from,
+    read back from storage. With a GPU snapshot, a pageable host buffer holds
+    anchor copies, allocated at the first periodic anchor.
 
     Tensors the engines hold in fp8 (``set_fp8_layout``: HF name -> scale format,
     from the engine layout) are quantized with vLLM's own per-block quantizer,
@@ -112,6 +116,9 @@ def __init__(self, model: nn.Module, snapshot: str):
         self.on_gpu: bool | None = None  # None until placed, at the first broadcast after a training step
         self.host_buffer: torch.Tensor | None = None
         self.write_staging: torch.Tensor | None = None
+        # HF name -> (file, dtype, data offset) in the boot checkpoint, and keys whose v0 is still to read from it
+        self.boot_files: dict[str, list[tuple[int, Path, torch.dtype, int]]] = {}
+        self.boot_pending: set[str] = set()
         self.host: dict[str, torch.Tensor] = {}
         self.pinned = False
         self.initialized = False
@@ -278,45 +285,67 @@ def anchor_parts(self, source: dict[str, torch.Tensor]) -> list[tuple[str, list[
         return parts
 
     @torch.no_grad()
-    def prepare(self, shards: list[tuple], peak_reserved: int) -> None:
-        """The training thread's part of an encode, before ``encode``. The first one places and allocates
-        the snapshot, against the measured training peak; the host buffer only holds a CPU snapshot."""
-        if not self.initialized:
-            self.track(shards)
-            self._place(peak_reserved)
-            self.staging = torch.empty(1 << 28, dtype=torch.uint8, pin_memory=True)
-            if self.on_gpu:
-                device = f"cuda:{self.device}"
-                self.snapshot = {
-                    key: torch.empty(numel, dtype=self._bits(key), device=device)
-                    for key, (numel, _) in self.shards.items()
-                }
-            else:
-                self._allocate_host()
-                self._page_lock()
-                self.snapshot = dict(self.host)
+    def initialize(self, model: nn.Module, is_master: bool, peak_reserved: int, boot_dir: Path) -> None:
+        """At the first delta (v1), on the training thread: place the snapshot against the measured training
+        peak and allocate it (the host buffer only for a CPU snapshot). The snapshot starts as v0, the
+        checkpoint in ``boot_dir`` the engines booted from, read as the encode goes (tensors gathered for fp8
+        are read here, before ``local_shards`` gathers them)."""
+        self.track(self.local_shards(model, is_master))
+        self._place(peak_reserved)
+        self.staging = torch.empty(1 << 28, dtype=torch.uint8, pin_memory=True)
+        if self.on_gpu:
+            device = f"cuda:{self.device}"
+            self.snapshot = {
+                key: torch.empty(numel, dtype=self._bits(key), device=device) for key, (numel, _) in self.shards.items()
+            }
+        else:
+            self._allocate_host()
+            self._page_lock()
+            self.snapshot = dict(self.host)
+        self.boot_files = safetensors_index(boot_dir)
+        self.boot_pending = set(self.shards) - self.gather
+        with torch.cuda.stream(self.stream):
+            for key in self.gather & set(self.shards):
+                self._read_boot(key)
+            self.stream.synchronize()
+        self.initialized = True
+
+    def _read_boot(self, key: str) -> None:
+        """Fill the snapshot of ``key`` with this rank's slice of the boot checkpoint, through the pinned
+        staging buffer, on the current stream."""
+        out = self.snapshot[key].view(self.wire_dtypes[key])
+        numel, offset = self.shards[key]
+        staging = self.staging.numpy()
+        for start, name in self.segments[self.key_fqn[key]]:
+            begin, end = max(start, offset), min(start + prod(self.hf_tensors[name][0]), offset + numel)
+            if begin >= end:
+                continue
+            first, path, dtype, data_start = max(s for s in self.boot_files[name] if s[0] <= begin - start)
+            itemsize = dtype.itemsize
+            with open(path, "rb") as file:
+                for chunk in range(begin, end, self.staging.numel() // itemsize):
+                    count = min(self.staging.numel() // itemsize, end - chunk)
+                    file_offset = data_start + (chunk - start - first) * itemsize
+                    assert os.preadv(file.fileno(), [staging[: count * itemsize]], file_offset) == count * itemsize
+                    os.posix_fadvise(file.fileno(), file_offset, count * itemsize, os.POSIX_FADV_DONTNEED)
+                    out[chunk - offset : chunk - offset + count].copy_(
+                        self.staging[: count * itemsize].view(dtype), non_blocking=True
+                    )
+                    torch.cuda.current_stream().synchronize()
 
     @torch.no_grad()
     def encode(
         self, shards: list[tuple], ready: torch.cuda.Event, weights_read: threading.Event, emit: Callable
     ) -> None:
         """Diff the shards against the snapshot and advance it, on the side stream once ``ready`` (recorded
-        after the optimizer step) has passed; the first call only fills the snapshot. Hands the delta to
-        ``emit`` in parts ``(entries, gaps, value bytes, new-weight checksum per HF tensor part)`` on the
-        host as they are encoded. Sets ``weights_read`` once done reading the weights. ``scales_changed``:
-        whether an fp8 block scale changed on this rank."""
+        after the optimizer step) has passed. Hands the delta to ``emit`` in parts ``(entries, gaps, value
+        bytes, new-weight checksum per HF tensor part)`` on the host as they are encoded. Sets ``weights_read``
+        once done reading the weights. ``scales_changed``: whether an fp8 block scale changed on this rank."""
         torch.cuda.set_device(self.device)
         try:
             with torch.cuda.stream(self.stream):
                 self.stream.wait_event(ready)
-                if not self.initialized:
-                    for key, _, shard, _, _ in shards:
-                        self.snapshot[key].copy_(
-                            as_bits(shard.reshape(-1).to(self.wire_dtypes[key])), non_blocking=True
-                        )
-                    self.initialized = True
-                else:
-                    self.scales_changed = self._diff(shards, emit)
+                self.scales_changed = self._diff(shards, emit)
                 self.stream.synchronize()
         finally:
             weights_read.set()
@@ -408,6 +437,9 @@ def flush() -> None:
 
         for key, fqn, shard, offset, gathered in shards:
             shard = shard.reshape(-1)
+            if key in self.boot_pending:
+                self._read_boot(key)
+                self.boot_pending.discard(key)
             previous = out = self.snapshot[key]
             wire = self.wire_dtypes[key]
             segment_names, begins, ends, shifts = self._segments(key, fqn, offset, shard.numel())
@@ -488,6 +520,12 @@ def flush() -> None:
         return bool(scales_changed)
 
 
+def _model_dir(model: nn.Module) -> Path:
+    """The local dir of the checkpoint the model was loaded from, as the trainer resolves it."""
+    name = model.config.name_or_path
+    return Path(name) if Path(name).exists() else Path(snapshot_download(repo_id=name, repo_type="model"))
+
+
 def _write_json(path: Path, value) -> None:
     write_file(path, json.dumps(value).encode())
 
@@ -520,6 +558,9 @@ def __init__(self, output_dir: Path, config: FileSystemWeightBroadcastConfig):
         self.errors: list[BaseException] = []
         self.weights_read: threading.Event | None = None
         self.part_bytes = 0
+        # The checkpoint the engines hold before the first delta: the model's (None), or a resumed run's first
+        # broadcast
+        self.boot_dir: Path | None = None
 
     def wait_for_weight_reads(self) -> None:
         if self.weights_read is not None:
@@ -565,28 +606,33 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
                 _write_json(step_dir / DELTA_DIR / MANIFEST, manifest)
                 WeightSender._finish(self, step, step_dir)
             return
-        if encoder.initialized and encoder.fp8_layout is None:
-            # Every engine worker reports which tensors it holds in fp8 when it loads the first broadcast.
-            layout = [self._wait_for_engine_layout(step_dir.parent) if self.world.is_master else None]
-            dist.broadcast_object_list(layout, src=0)
-            encoder.set_fp8_layout(model, layout[0], self.world.is_master)
-        shards = encoder.local_shards(model, self.world.is_master)
         if base_step is None:
-            # A resumed run's startup broadcast: a full checkpoint, written from the weights before training.
+            # A resumed run's startup broadcast: a full checkpoint, written from the weights before training. The
+            # next broadcast's delta is against it.
+            shards = encoder.local_shards(model, self.world.is_master)
             encoder.track(shards)
             source = {key: as_bits(shard.reshape(-1).to(encoder.wire_dtypes[key])) for key, _, shard, _, _ in shards}
             torch.cuda.synchronize()
             self._write_full(step, step_dir, encoder.anchor_parts(source))
+            self.boot_dir = step_dir
             return
-        full = not encoder.initialized
-        encoder.prepare(shards, torch.cuda.max_memory_reserved())
+        if encoder.fp8_layout is None:
+            # Every engine worker reports which tensors it holds in fp8 when it loads the first broadcast.
+            layout = [self._wait_for_engine_layout(step_dir.parent) if self.world.is_master else None]
+            dist.broadcast_object_list(layout, src=0)
+            encoder.set_fp8_layout(model, layout[0], self.world.is_master)
+        if not encoder.initialized:
+            encoder.initialize(
+                model, self.world.is_master, torch.cuda.max_memory_reserved(), self.boot_dir or _model_dir(model)
+            )
+        shards = encoder.local_shards(model, self.world.is_master)
         # Periodic anchors only serve engines that fall off the chain: written in the background, each after
         # the previous one.
         anchor = step % self.config.anchor_interval == 0
         ready = torch.cuda.Event()
         ready.record()
         self.weights_read = threading.Event()
-        args = (step, base_step, step_dir, shards, ready, self.weights_read, full, anchor)
+        args = (step, base_step, step_dir, shards, ready, self.weights_read, anchor)
         self._submit(self.encodes, self._encode, *args)
 
     def _wait_for_engine_layout(self, broadcast_dir: Path) -> dict[str, str]:
@@ -627,21 +673,14 @@ def _encode(
         shards,
         ready,
         weights_read: threading.Event,
-        full: bool,
         anchor: bool,
     ) -> None:
         """Encode thread: diff v{step}, handing the delta to the writer in parts as it is encoded. With fp8
         deltas, the parts are held until the ranks agree that no block scale changed, else the step is a
-        full checkpoint. The first broadcast after a training step only fills the snapshot and is published
-        as a full checkpoint (the engines reload it once)."""
+        full checkpoint."""
         encoder = self.encoder
         assert encoder is not None
         self._wait(self.snapshot_writes)
-        if full:
-            encoder.encode(shards, ready, weights_read, emit=None)
-            shards.clear()
-            self._write_from_snapshot(self.writes, self._write_full, step, step_dir, full=True)
-            return
         start = time.perf_counter()
         held: list[tuple] = []
         num_parts = 0
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index 567b984cbf2..d4ba40bbb75 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -302,6 +302,22 @@ def write_anchor_shards(
     write_file(step_dir / f"anchor_rank{rank}.done", str(world_size).encode())
 
 
+def safetensors_index(directory: Path) -> dict[str, list[tuple[int, Path, torch.dtype, int]]]:
+    """Tensor name -> its slices ``(flat start, file, dtype, data offset)`` in the safetensors files in
+    ``directory``: a HF checkpoint (one slice at 0), or an anchor (``"{name}@{start}"`` slices)."""
+    index = defaultdict(list)
+    for path in directory.glob("*.safetensors"):
+        with open(path, "rb") as file:
+            header_size = int.from_bytes(file.read(8), "little")
+            header = json.loads(file.read(header_size))
+        header.pop("__metadata__", None)
+        for key, info in header.items():
+            name, _, start = key.partition("@")
+            dtype = getattr(torch, SAFETENSORS_DTYPES[info["dtype"]])
+            index[name].append((int(start or 0), path, dtype, 8 + header_size + info["data_offsets"][0]))
+    return dict(index)
+
+
 def anchor_index(step_dir: Path) -> tuple[dict[str, list[tuple[Path, str, int]]], dict[str, tuple[list[int], str]]]:
     """Where each tensor's slices are in an anchor written by ``write_anchor_shards``, and its shape and dtype."""
     slices: dict[str, list[tuple[Path, str, int]]] = defaultdict(list)

From 515cfc243948ea870f2d431e52effdd2c644ca7e Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Fri, 16 Oct 2026 23:04:59 +0000
Subject: [PATCH 42/88] fix(weights): order the gathered fp8 tensors v0 read
 after the snapshot allocation

---
 src/prime_rl/transports/weights/sparse.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 629e60cc067..5e5f67329f2 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -304,6 +304,8 @@ def initialize(self, model: nn.Module, is_master: bool, peak_reserved: int, boot
             self.snapshot = dict(self.host)
         self.boot_files = safetensors_index(boot_dir)
         self.boot_pending = set(self.shards) - self.gather
+        # The snapshot was just allocated on this stream, whose earlier work may still use that memory
+        self.stream.wait_stream(torch.cuda.current_stream())
         with torch.cuda.stream(self.stream):
             for key in self.gather & set(self.shards):
                 self._read_boot(key)

From 54ad1b47ca5b9e7616497395c9469daf0e66b8b7 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sat, 17 Oct 2026 01:40:11 +0000
Subject: [PATCH 43/88] feat(weights): fp8 block scales in the delta;
 slice-local full checkpoint loads

fp8 tensors' block scales are diffed as "{name}_scale_inv" tensors in the normal delta and patched in place by the engines, so a scale change no longer publishes a full checkpoint.
Anchor and boot checkpoint loads in sparse runs read only the rows each engine worker holds, straight into its params through the checkpoint index map (the set counterpart of the delta apply), instead of every worker reading the whole checkpoint through the layerwise loader. The engine layout and index map are built at the first update, from the boot checkpoint.
---
 .../inference/vllm/worker/filesystem.py       | 112 ++++-----
 .../inference/vllm/worker/weight_transfer.py  | 109 +++++++--
 src/prime_rl/transports/weights/sparse.py     | 215 ++++++++++--------
 src/prime_rl/utils/sparse_weights.py          |  83 +++----
 4 files changed, 304 insertions(+), 215 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 819935cc87e..90702cd89ff 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -2,12 +2,12 @@
 import json
 import time
 from concurrent.futures import Future, ThreadPoolExecutor
+from math import prod
 from pathlib import Path
 from typing import TYPE_CHECKING
 
 import numpy as np
 import torch
-from safetensors import safe_open
 from torch.nn import Module
 from vllm.distributed.parallel_state import get_ep_group, get_tp_group, get_world_group
 from vllm.logger import init_logger
@@ -22,13 +22,12 @@
     DELTA_DIR,
     ENGINE_LAYOUT,
     MANIFEST,
-    SAFETENSORS_DTYPES,
     anchor_complete,
-    anchor_index,
-    anchor_tensors,
     bits_dtype,
+    checkpoint_index,
     decode_gaps,
     pin,
+    read_checkpoint_range,
     read_delta,
     read_manifest,
     unpin,
@@ -146,20 +145,19 @@ def update_weights_from_path(self, weight_path: str) -> None:
         """Move to the broadcast step in ``weight_path`` (a full HF checkpoint and/or a sparse delta)."""
         step_dir = Path(weight_path)
         step = int(step_dir.name.removeprefix("step_"))
+        has_delta = (step_dir / DELTA_DIR / MANIFEST).exists()
+        if has_delta or anchor_complete(step_dir):
+            self._build_index_map(step_dir)
         # Unknown until this update succeeds: after a failure, the next update reloads and replays
         held_step, self.weight_step = self.weight_step, -1
-        if not (step_dir / DELTA_DIR / MANIFEST).exists():
+        if not has_delta:
             self._load_checkpoint(step_dir)
-            if anchor_complete(step_dir):
-                # A sparse run's full checkpoint
-                self._write_engine_layout(step_dir, step_dir)
         else:
             base_step = read_manifest(step_dir)["base_step"]
             if base_step is None:
                 # The run's first broadcast: the weights the engine booted with.
                 if held_step is not None:
                     self._load_checkpoint(self._boot_dir())
-                self._write_engine_layout(step_dir, self._boot_dir())
             elif base_step == held_step:
                 self._apply_delta(step_dir)
             else:
@@ -194,46 +192,37 @@ def _apply_chunk(
             offset += num_changes
             num_bytes += num_changes * itemsize
 
-    def _write_engine_layout(self, step_dir: Path, checkpoint_dir: Path) -> None:
-        """Tell the trainer which checkpoint tensors this engine holds in fp8 (and with which
-        scales), so it can send their deltas in the fp8 domain. Once per engine, while its weights
-        are those of ``checkpoint_dir`` (an anchor or the boot checkpoint): it reads one checkpoint
-        tensor per fp8 param from there."""
-        if self.wrote_layout:
+    def _build_index_map(self, step_dir: Path) -> None:
+        """At a sparse run's first update, while the engine still holds its boot weights: map checkpoint
+        tensors to the params, and tell the trainer which ones this engine holds in fp8 (and with which
+        scales), so it can send their deltas in the fp8 domain. Reads one boot checkpoint tensor per fp8
+        param."""
+        if self.index_map is not None:
             return
-        self.wrote_layout = True
-        if anchor_complete(checkpoint_dir):
-            index = anchor_index(checkpoint_dir)
-            tensors = index[1]
-
-            def read_tensor(name: str) -> torch.Tensor:
-                return next(anchor_tensors(checkpoint_dir, [name], index))[1]
-
-        else:
-            files = {}
-            tensors = {}
-            for path in checkpoint_dir.glob("*.safetensors"):
-                with safe_open(path, framework="pt") as f:
-                    for name in f.keys():
-                        tensor_slice = f.get_slice(name)
-                        files[name] = path
-                        tensors[name] = (tensor_slice.get_shape(), SAFETENSORS_DTYPES[tensor_slice.get_dtype()])
-
-            def read_tensor(name: str) -> torch.Tensor:
-                with safe_open(files[name], framework="pt") as f:
-                    return f.get_tensor(name)
-
         if (step_dir / DELTA_DIR / MANIFEST).exists():
             tensors = read_manifest(step_dir)["tensors"]
-        if self.index_map is None:
-            self.index_map = CheckpointIndexMap(self._model(), tensors)
+        else:
+            tensors = {
+                name: (shape, str(dtype).removeprefix("torch."))
+                for name, (shape, dtype, _) in checkpoint_index(step_dir).items()
+            }
+        model = self._model()
+        index_map = CheckpointIndexMap(model, tensors)
+        boot = checkpoint_index(self._boot_dir())
+
+        def read_tensor(name: str) -> torch.Tensor:
+            shape, dtype, _ = boot[name]
+            data = np.empty(prod(shape) * dtype.itemsize, np.uint8)
+            read_checkpoint_range(boot[name], 0, prod(shape), data)
+            return torch.from_numpy(data).view(dtype).view(shape)
 
         # Each worker reports the tensors it holds; the trainer waits for all of them.
         if self.vllm_config.parallel_config.enable_expert_parallel:
             group = get_ep_group()
         else:
             group = get_world_group()
-        layout = {"world_size": group.world_size, "fp8": self.index_map.fp8_layout(read_tensor)}
+        layout = {"world_size": group.world_size, "fp8": index_map.fp8_layout(model, read_tensor)}
+        self.index_map = index_map
         path = step_dir.parent / ENGINE_LAYOUT / f"rank_{group.rank_in_group}.json"
         path.parent.mkdir(exist_ok=True)
         write_file(path, json.dumps(layout).encode())
@@ -250,9 +239,7 @@ def _replay(self, step_dir: Path) -> None:
             chain.append(source)
             source = source.parent / f"step_{base_step}"
         logger.info(f"Replaying {len(chain)} sparse weight deltas on top of {source.name}")
-        source = source if anchor_complete(source) else self._boot_dir()
-        self._load_checkpoint(source)
-        self._write_engine_layout(step_dir, source)
+        self._load_checkpoint(source if anchor_complete(source) else self._boot_dir())
         for delta_dir in reversed(chain):
             self._apply_delta(delta_dir)
 
@@ -268,9 +255,6 @@ def _apply_delta(self, step_dir: Path) -> None:
         tensors = read_manifest(step_dir)["tensors"]
         model = self._model()
         device = next(model.parameters()).device
-        if self.index_map is None:
-            self.index_map = CheckpointIndexMap(model, tensors)
-            logger.info(f"Built the checkpoint index map in {time.perf_counter() - start - read_time:.2f}s")
         for entries, gaps, values in delta:
             self._apply_chunk(entries, gaps, values, tensors, device)
         update_mla_absorbed_weights(model)
@@ -305,11 +289,10 @@ def _verify(self, step_dir: Path, checksums: dict[str, int]) -> None:
             raise RuntimeError(f"Sparse weight delta {step_dir.name} diverged from the trainer: {mismatched[:8]}")
 
     def _load_checkpoint(self, step_dir: Path) -> None:
-        """Reload every weight from the HF-compatible checkpoint (or sparse-run anchor) in ``step_dir``."""
+        """Reload every weight from the HF-compatible checkpoint (or sparse-run anchor) in ``step_dir``. In a
+        sparse run, each worker reads only the rows it holds, straight into its params."""
         model = self._model()
-        if anchor_complete(step_dir):
-            weights_iterator = anchor_tensors(step_dir)
-        else:
+        if self.index_map is None:
             model_loader = get_model_loader(self.load_config)
             assert isinstance(model_loader, DefaultModelLoader)
             local_source = DefaultModelLoader.Source(
@@ -320,9 +303,28 @@ def _load_checkpoint(self, step_dir: Path) -> None:
                 allow_patterns_overrides=getattr(model, "allow_patterns_overrides", None),
             )
             weights_iterator = model_loader._get_weights_iterator(local_source)
-        load_weights_checkpoint_layerwise(
-            model,
-            weights_iterator,
-            self.model_runner.model_config,
-            self.vllm_config,
+            load_weights_checkpoint_layerwise(model, weights_iterator, self.model_runner.model_config, self.vllm_config)
+            return
+        start = time.perf_counter()
+        index_map, index = self.index_map, checkpoint_index(step_dir)
+        device = next(model.parameters()).device
+
+        def read(name: str, begin: int, end: int) -> torch.Tensor:
+            num_cols, dtype = index_map.num_cols[name], index[name][1]
+            data = np.empty((end - begin) * num_cols * dtype.itemsize, np.uint8)
+            read_checkpoint_range(index[name], begin * num_cols, end * num_cols, data)
+            return torch.from_numpy(data).view(dtype).view(end - begin, num_cols)
+
+        ranges = list(index_map.row_ranges().items())
+        num_bytes = 0
+        with ThreadPoolExecutor(16) as pool:
+            # A bounded window of reads ahead of the GPU writes
+            for window in itertools.batched(ranges, 64):
+                for (name, (begin, _)), rows in zip(window, pool.map(lambda item: read(item[0], *item[1]), window)):
+                    index_map.load_rows(name, rows.to(device), begin)
+                    num_bytes += rows.nbytes
+        update_mla_absorbed_weights(model)
+        torch.cuda.synchronize()
+        logger.info(
+            f"Loaded this worker's {num_bytes / 2**30:.2f} GiB of {step_dir.name} in {time.perf_counter() - start:.1f}s"
         )
diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py
index 36b10bfbbea..05984415f49 100644
--- a/src/prime_rl/inference/vllm/worker/weight_transfer.py
+++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py
@@ -15,7 +15,7 @@
 from vllm.model_executor.model_loader.reload.meta import restore_layer_on_meta
 from vllm.model_executor.model_loader.reload.utils import get_layer_params_buffers, get_layer_tensors
 
-from prime_rl.utils.sparse_weights import as_bits, quantize_fp8_blocks
+from prime_rl.utils.sparse_weights import FP8_BLOCK, as_bits, quantize_fp8_blocks
 
 logger = init_logger("vllm.inference.vllm.worker_weight_transfer")
 
@@ -85,6 +85,8 @@ def record_copy_(destination: torch.Tensor, source: torch.Tensor, non_blocking:
 
         def placeholders():
             for name, (shape, dtype) in tensors.items():
+                if name.endswith("_scale_inv"):
+                    continue  # fp8 block scales of a delta: recorded in fp8_layout
                 # fp8 tensors of a delta are quantized from bf16 checkpoint tensors
                 checkpoint_dtype = torch.bfloat16 if dtype == "float8_e4m3fn" else getattr(torch, dtype)
                 tensor = torch.empty(shape, dtype=checkpoint_dtype, device=device)
@@ -111,6 +113,7 @@ def placeholders():
             if live[name].shape != shapes[name] or not live[name].is_contiguous() or live[name].storage_offset():
                 raise NotImplementedError(f"Sparse weight deltas cannot patch {name}: its layout changes after loading")
         self.params: list[torch.Tensor] = [live[name] for name in used]
+        self.param_names = used
         param_index = {name: index for index, name in enumerate(used)}
         self.records = {
             name: [(*write[:4], param_index[write[4]], *write[5:]) for write in writes]
@@ -119,6 +122,7 @@ def placeholders():
         self.num_cols = {name: (shape[-1] if shape else 1) for name, (shape, _) in tensors.items()}
         self.numel = {name: prod(shape) for name, (shape, _) in tensors.items()}
         self.checksum_plan = None
+        self.fp8: dict[str, str] = {}
         self.unpatchable = unpatchable
         # Checkpoint tensors this worker holds; unpatchable ones too, so a change to them is caught in apply_delta
         self.names = set(self.records) | unpatchable
@@ -186,22 +190,58 @@ def _plan_checksums(self):
         partial = torch.tensor([h < self.numel[name] for h, name in zip(held, names)], device=device)
         return names, partial, row_ranges, other_writes
 
-    def fp8_layout(self, read_tensor) -> dict[str, str]:
+    def fp8_layout(self, model: Module, read_tensor) -> dict[str, str]:
         """The vLLM per-block quantization (``"fp32"`` or ``"ue8m0"`` scales) that reproduces each
-        fp8 param from the checkpoint, per checkpoint tensor stored in one. Checks one checkpoint
-        tensor per param; ``read_tensor(name)`` returns it."""
+        fp8 param from the checkpoint, per checkpoint tensor stored in one. Also records where each such
+        tensor's block scales (``"{name}_scale_inv"``, one fp32 per 128x128 block) live, so deltas and loads
+        patch them in place. Checks one checkpoint tensor per param (weights and scales); ``read_tensor(name)``
+        returns it."""
         modes: dict[int, str] = {}
-        layout = {}
-        for name, writes in self.records.items():
+        live = dict(model.named_parameters())
+        scale_params: dict[str, int] = {}
+        for name, writes in list(self.records.items()):
             for row0, rows, col0, cols, index, offset, *strides in writes:
-                if self.params[index].dtype != torch.float8_e4m3fn:
+                param = self.params[index]
+                if param.dtype != torch.float8_e4m3fn:
                     continue
+                scale = live.get(f"{self.param_names[index]}_scale_inv")
+                num_rows, num_cols = param.shape[-2:]
+                row, col = offset // num_cols % num_rows, offset % num_cols
+                if (
+                    scale is None
+                    or scale.dtype != torch.float32
+                    or scale.shape[:-2] != param.shape[:-2]
+                    or strides != [num_cols, 1]
+                    or any(x % FP8_BLOCK[0] for x in (row0, col0, row, col))
+                ):
+                    raise NotImplementedError(f"Sparse weight deltas cannot patch the fp8 block scales of {name}")
+                if scale.untyped_storage()._cdata not in scale_params:
+                    scale_params[scale.untyped_storage()._cdata] = len(self.params)
+                    # The whole storage, flat: kernels may keep scales in a strided (e.g. MN-major) layout
+                    self.params.append(scale.data.as_strided((scale.untyped_storage().nbytes() // 4,), (1,), 0))
+                matrix, scale_offset = offset // (num_rows * num_cols), scale.storage_offset()
+                for dim in reversed(range(param.dim() - 2)):
+                    matrix, position = divmod(matrix, param.shape[dim])
+                    scale_offset += position * scale.stride(dim)
+                scale_offset += row // FP8_BLOCK[0] * scale.stride(-2) + col // FP8_BLOCK[1] * scale.stride(-1)
+                block = (row0 // FP8_BLOCK[0], -(-rows // FP8_BLOCK[0]), col0 // FP8_BLOCK[1], -(-cols // FP8_BLOCK[1]))
+                scale_name = f"{name}_scale_inv"
+                self.records.setdefault(scale_name, []).append(
+                    (*block, scale_params[scale.untyped_storage()._cdata], scale_offset, *scale.stride()[-2:])
+                )
+                self.num_cols[scale_name] = -(-self.num_cols[name] // FP8_BLOCK[1])
+                self.numel[scale_name] = (
+                    -(-self.numel[name] // self.num_cols[name] // FP8_BLOCK[0]) * self.num_cols[scale_name]
+                )
                 if index not in modes:
-                    live = as_bits(self.params[index].data).view(-1).as_strided((rows, cols), strides, offset)
-                    weight = read_tensor(name).to(self.params[index].device)
+                    weight = read_tensor(name).to(param.device)
+                    weight = weight.view(-1, self.num_cols[name])
                     for mode in ("ue8m0", "fp32"):
-                        quantized, _ = quantize_fp8_blocks(weight, mode == "ue8m0")
-                        if torch.equal(quantized[row0 : row0 + rows, col0 : col0 + cols], live):
+                        quantized, scales = quantize_fp8_blocks(weight, mode == "ue8m0")
+                        if torch.equal(
+                            quantized[row0 : row0 + rows, col0 : col0 + cols],
+                            self._view(index, rows, cols, offset, strides),
+                        ):
                             modes[index] = mode
                             break
                     else:
@@ -209,9 +249,52 @@ def fp8_layout(self, read_tensor) -> dict[str, str]:
                             f"Sparse weight deltas cannot patch {name}: the engine's fp8 weights are not vLLM's "
                             "per-block quantization of the checkpoint (e.g. a weight-shuffling kernel)"
                         )
-                if layout.setdefault(name, modes[index]) != modes[index]:
+                    _, block_rows, _, block_cols, scale_index, scale_offset, *scale_strides = self.records[scale_name][
+                        -1
+                    ]
+                    expected = scales[block[0] : block[0] + block_rows, block[2] : block[2] + block_cols]
+                    if not torch.equal(
+                        as_bits(expected), self._view(scale_index, block_rows, block_cols, scale_offset, scale_strides)
+                    ):
+                        raise NotImplementedError(f"Sparse weight deltas cannot patch the fp8 block scales of {name}")
+                if self.fp8.setdefault(name, modes[index]) != modes[index]:
                     raise NotImplementedError(f"Sparse weight deltas cannot patch {name}: mixed fp8 layouts")
-        return layout
+        self.names |= set(self.records)
+        self.checksum_plan = None
+        return self.fp8
+
+    def _view(self, index: int, rows: int, cols: int, offset: int, strides) -> torch.Tensor:
+        """The raw bits of one write of a checkpoint block into param ``index``."""
+        return as_bits(self.params[index].data).view(-1).as_strided((rows, cols), tuple(strides), offset)
+
+    def row_ranges(self) -> dict[str, tuple[int, int]]:
+        """Checkpoint tensor -> the rows ``[begin, end)`` of its ``[rows, num_cols]`` view this worker holds
+        (fp8 tensors: whole 128-row blocks), for ``load_rows``. Tensors the engine casts on load are left out:
+        they never change."""
+        ranges = {}
+        for name, writes in self.records.items():
+            if name.endswith("_scale_inv") or name in self.unpatchable:
+                continue  # scales are set with their weights
+            begin, end = min(write[0] for write in writes), max(write[0] + write[1] for write in writes)
+            if name in self.fp8:
+                num_rows = self.numel[name] // self.num_cols[name]
+                begin, end = begin // FP8_BLOCK[0] * FP8_BLOCK[0], min(-(-end // FP8_BLOCK[0]) * FP8_BLOCK[0], num_rows)
+            ranges[name] = (begin, end)
+        return ranges
+
+    @torch.no_grad()
+    def load_rows(self, name: str, rows: torch.Tensor, begin: int) -> None:
+        """Set the params from rows ``begin:`` of checkpoint tensor ``name`` (the ``row_ranges`` rows, on the
+        params' device): the counterpart of ``apply_delta`` for full checkpoints. fp8 tensors are quantized
+        like the engine does, block scales included."""
+        sources = [(name, rows, begin)]
+        if name in self.fp8:
+            quantized, scales = quantize_fp8_blocks(rows, self.fp8[name] == "ue8m0")
+            sources = [(name, quantized, begin), (f"{name}_scale_inv", scales, begin // FP8_BLOCK[0])]
+        for source_name, source, first in sources:
+            for row0, rows_, col0, cols, index, offset, *strides in self.records[source_name]:
+                block = source[row0 - first : row0 - first + rows_, col0 : col0 + cols]
+                self._view(index, rows_, cols, offset, strides).copy_(as_bits(block))
 
     @torch.no_grad()
     def apply_delta(
diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 5e5f67329f2..13fa0bbf5aa 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -1,5 +1,4 @@
 import json
-import os
 import threading
 import time
 from bisect import bisect_left, bisect_right
@@ -28,15 +27,16 @@
     anchor_complete,
     as_bits,
     bits_dtype,
+    checkpoint_index,
     compile_fp8_quantizer,
     encode_gaps,
     page_lock,
     quantize_fp8_blocks,
     rank_part,
     rank_parts_marker,
+    read_checkpoint_range,
     read_engine_layout,
     revalidate,
-    safetensors_index,
     write_anchor_shards,
     write_file,
     write_rank_file,
@@ -69,8 +69,9 @@ class SparseDeltaEncoder:
     previous and current weights alike, and diffed as fp8 bytes. Their 128-row
     blocks must not straddle trainer ranks, else the trainer tensor is gathered.
     The quantizer is compiled ahead of time, so it can run on the encode thread
-    (columns a multiple of 128). A step that changes any block scale is
-    published as a full checkpoint.
+    (columns a multiple of 128). Their block scales are diffed too, as tensors
+    ``"{hf_name}_scale_inv"`` (fp32, one element per 128x128 block) in the same
+    delta, so engines patch scales in place.
     """
 
     # Elements per transient chunk of the diff (bf16: 128 MiB per temporary)
@@ -104,7 +105,9 @@ def __init__(self, model: nn.Module, snapshot: str):
         self.fp8: dict[str, tuple[bool, int]] = {}
         self.fp8_layout: dict[str, str] | None = None
         self.compiled_fp8 = None
-        self.scales_changed = False
+        # fqn -> segments of its fp8 block scales, like ``segments``
+        self.scale_segments: dict[str, list[tuple[int, str]]] = {}
+        self.scale_changes = 0
         self.gather: set[str] = set()
         self.segment_starts: dict[str, torch.Tensor] = {}
         self.shard_segments: dict[str, tuple[list[str], list[int], list[int], list[int]]] = {}
@@ -116,8 +119,8 @@ def __init__(self, model: nn.Module, snapshot: str):
         self.on_gpu: bool | None = None  # None until placed, at the first broadcast after a training step
         self.host_buffer: torch.Tensor | None = None
         self.write_staging: torch.Tensor | None = None
-        # HF name -> (file, dtype, data offset) in the boot checkpoint, and keys whose v0 is still to read from it
-        self.boot_files: dict[str, list[tuple[int, Path, torch.dtype, int]]] = {}
+        # The boot checkpoint's ``checkpoint_index``, and keys whose v0 is still to read from it
+        self.boot_files: dict[str, tuple] = {}
         self.boot_pending: set[str] = set()
         self.host: dict[str, torch.Tensor] = {}
         self.pinned = False
@@ -146,8 +149,12 @@ def set_fp8_layout(self, model: nn.Module, fp8: dict[str, str], is_master: bool)
             ):
                 raise NotImplementedError(f"Sparse weight broadcast cannot quantize {fqn} to fp8 as one matrix")
             self.fp8[fqn] = (modes.pop() == "ue8m0", num_cols)
-            for _, name in segs:
-                self.hf_tensors[name] = (self.hf_tensors[name][0], "float8_e4m3fn")
+            scale_cols = num_cols // FP8_BLOCK[1]
+            self.scale_segments[fqn] = []
+            for (start, name), shape in zip(segs, shapes):
+                self.hf_tensors[name] = (shape, "float8_e4m3fn")
+                self.hf_tensors[f"{name}_scale_inv"] = ([-(-shape[0] // FP8_BLOCK[0]), scale_cols], "float32")
+                self.scale_segments[fqn].append((start // num_cols // FP8_BLOCK[0] * scale_cols, f"{name}_scale_inv"))
         self.gather = self._find_gathered(model, is_master)
         if self.fp8:
             start = time.perf_counter()
@@ -302,7 +309,7 @@ def initialize(self, model: nn.Module, is_master: bool, peak_reserved: int, boot
             self._allocate_host()
             self._page_lock()
             self.snapshot = dict(self.host)
-        self.boot_files = safetensors_index(boot_dir)
+        self.boot_files = checkpoint_index(boot_dir)
         self.boot_pending = set(self.shards) - self.gather
         # The snapshot was just allocated on this stream, whose earlier work may still use that memory
         self.stream.wait_stream(torch.cuda.current_stream())
@@ -322,18 +329,15 @@ def _read_boot(self, key: str) -> None:
             begin, end = max(start, offset), min(start + prod(self.hf_tensors[name][0]), offset + numel)
             if begin >= end:
                 continue
-            first, path, dtype, data_start = max(s for s in self.boot_files[name] if s[0] <= begin - start)
-            itemsize = dtype.itemsize
-            with open(path, "rb") as file:
-                for chunk in range(begin, end, self.staging.numel() // itemsize):
-                    count = min(self.staging.numel() // itemsize, end - chunk)
-                    file_offset = data_start + (chunk - start - first) * itemsize
-                    assert os.preadv(file.fileno(), [staging[: count * itemsize]], file_offset) == count * itemsize
-                    os.posix_fadvise(file.fileno(), file_offset, count * itemsize, os.POSIX_FADV_DONTNEED)
-                    out[chunk - offset : chunk - offset + count].copy_(
-                        self.staging[: count * itemsize].view(dtype), non_blocking=True
-                    )
-                    torch.cuda.current_stream().synchronize()
+            dtype = self.boot_files[name][1]
+            step = self.staging.numel() // dtype.itemsize
+            for chunk in range(begin, end, step):
+                count = min(step, end - chunk)
+                read_checkpoint_range(self.boot_files[name], chunk - start, chunk - start + count, staging)
+                out[chunk - offset : chunk - offset + count].copy_(
+                    self.staging[: count * dtype.itemsize].view(dtype), non_blocking=True
+                )
+                torch.cuda.current_stream().synchronize()
 
     @torch.no_grad()
     def encode(
@@ -342,12 +346,12 @@ def encode(
         """Diff the shards against the snapshot and advance it, on the side stream once ``ready`` (recorded
         after the optimizer step) has passed. Hands the delta to ``emit`` in parts ``(entries, gaps, value
         bytes, new-weight checksum per HF tensor part)`` on the host as they are encoded. Sets ``weights_read``
-        once done reading the weights. ``scales_changed``: whether an fp8 block scale changed on this rank."""
+        once done reading the weights."""
         torch.cuda.set_device(self.device)
         try:
             with torch.cuda.stream(self.stream):
                 self.stream.wait_event(ready)
-                self.scales_changed = self._diff(shards, emit)
+                self._diff(shards, emit)
                 self.stream.synchronize()
         finally:
             weights_read.set()
@@ -383,21 +387,21 @@ def _to_host(self, tensor: torch.Tensor, host: torch.Tensor | None = None) -> to
         return host
 
     def _segments(
-        self, key: str, fqn: str, offset: int, numel: int
+        self, cache_key, segments: list[tuple[int, str]], offset: int, numel: int
     ) -> tuple[list[str], list[int], list[int], list[int]]:
-        """The HF tensor parts in this rank's shard of ``key``: names, begin/end in the shard, and the shift
-        from a shard position to the HF flat index."""
-        if key not in self.shard_segments:
+        """The HF tensor parts of ``segments`` in this rank's ``numel`` elements from ``offset``: names,
+        begin/end in the shard, and the shift from a shard position to the HF flat index."""
+        if cache_key not in self.shard_segments:
             names, begins, ends, shifts = [], [], [], []
-            for start, name in self.segments[fqn]:
+            for start, name in segments:
                 hf_numel = prod(self.hf_tensors[name][0])
                 if start < offset + numel and start + hf_numel > offset:
                     names.append(name)
                     begins.append(max(start - offset, 0))
                     ends.append(min(start - offset + hf_numel, numel))
                     shifts.append(offset - start)
-            self.shard_segments[key] = (names, begins, ends, shifts)
-        return self.shard_segments[key]
+            self.shard_segments[cache_key] = (names, begins, ends, shifts)
+        return self.shard_segments[cache_key]
 
     @staticmethod
     def _segment_sums(bits: torch.Tensor, begins: list[int], ends: list[int], firsts: list[int]) -> torch.Tensor:
@@ -418,14 +422,42 @@ def _segment_sums(bits: torch.Tensor, begins: list[int], ends: list[int], firsts
             i = j
         return torch.cat(sums)
 
-    def _diff(self, shards: list[tuple], emit: Callable) -> bool:
-        """Diff ``shards`` against ``self.previous`` on the current stream, writing the new bits to the
-        snapshot, in chunks of ``CHUNK`` elements (fp8: whole 128-row blocks). Copies the delta to the
-        host and hands it to ``emit`` every ``PART`` changes. Returns whether an fp8 block scale changed."""
+    def _part(self, cache_key, segments: list[tuple[int, str]], position, values, checksums) -> tuple:
+        """A part entry for the changes at global flat ``position`` of a trainer tensor made of ``segments``:
+        ``(names, counts, gaps, value bytes, checksum names, checksums)``."""
+        if cache_key not in self.segment_starts:
+            self.segment_starts[cache_key] = torch.tensor([start for start, _ in segments], device=position.device)
+        starts = self.segment_starts[cache_key]
+        segment = torch.searchsorted(starts, position, right=True) - 1
+        entry_starts = torch.ones_like(segment, dtype=torch.bool)
+        entry_starts[1:] = segment[1:] != segment[:-1]
+        gaps = encode_gaps(position - starts[segment], entry_starts)
+        counts = torch.bincount(segment, minlength=len(starts))
+        return [name for _, name in segments], counts, gaps, values.view(torch.uint8), *checksums
+
+    def _scale_part(self, fqn: str, first_block: int, previous: torch.Tensor, current: torch.Tensor) -> tuple:
+        """The part entry of an fp8 tensor's changed block scales: block rows from ``first_block`` on."""
+        scale_cols = current.shape[1]
+        delta = as_bits(current).reshape(-1) - as_bits(previous).reshape(-1)
+        index = delta.nonzero().squeeze(1)
+        self.scale_changes += len(index)
+        offset = first_block * scale_cols
+        names, begins, ends, shifts = self._segments((fqn, "scales"), self.scale_segments[fqn], offset, current.numel())
+        sums = self._segment_sums(
+            as_bits(current).reshape(-1), begins, ends, [begin + shift for begin, shift in zip(begins, shifts)]
+        )
+        return self._part(
+            (fqn, "scales"), self.scale_segments[fqn], index + offset, zigzag(delta[index]), (names, sums)
+        )
+
+    def _diff(self, shards: list[tuple], emit: Callable) -> None:
+        """Diff ``shards`` against the snapshot on the current stream, writing the new bits to the
+        snapshot, in chunks of ``CHUNK`` elements (fp8: whole 128-row blocks, and their block scales).
+        Copies the delta to the host and hands it to ``emit`` every ``PART`` changes."""
         device = torch.device(f"cuda:{self.device}")
         pending: list[tuple] = []
         num_pending = 0
-        scales_changed = False
+        self.scale_changes = 0
 
         def flush() -> None:
             names = [name for part in pending for name in part[0]]
@@ -444,9 +476,9 @@ def flush() -> None:
                 self.boot_pending.discard(key)
             previous = out = self.snapshot[key]
             wire = self.wire_dtypes[key]
-            segment_names, begins, ends, shifts = self._segments(key, fqn, offset, shard.numel())
+            segment_names, begins, ends, shifts = self._segments(key, self.segments[fqn], offset, shard.numel())
             sums = torch.zeros(len(segment_names), dtype=torch.int64, device=device)
-            changed, changed_values = [], []
+            changed, changed_values, scales = [], [], []
             if fqn in self.fp8:
                 ue8m0, num_cols = self.fp8[fqn]
                 if gathered is not None:
@@ -472,7 +504,7 @@ def flush() -> None:
                         quantize_fp8_blocks(weights.reshape(-1, num_cols), ue8m0, self.compiled_fp8)
                         for weights in (previous_weights, current_weights)
                     )
-                    scales_changed |= not torch.equal(previous_scales, current_scales)
+                    scales.append((previous_scales, current_scales))
                     if gathered is not None:
                         previous_fp8 = previous_fp8.view(-1)[offset : offset + shard.numel()]
                         current_fp8 = current_fp8.view(-1)[offset : offset + shard.numel()]
@@ -493,33 +525,36 @@ def flush() -> None:
                         [min(stop, end) - start for stop in ends[first:last]],
                         [max(begin, start) + shift for begin, shift in zip(begins[first:last], shifts[first:last])],
                     )
-            changed, changed_values = torch.cat(changed), torch.cat(changed_values)
-            if fqn not in self.segment_starts:
-                self.segment_starts[fqn] = torch.tensor([start for start, _ in self.segments[fqn]], device=device)
-            starts = self.segment_starts[fqn]
-            position = changed + offset
-            segment = torch.searchsorted(starts, position, right=True) - 1
-            entry_starts = torch.ones_like(segment, dtype=torch.bool)
-            entry_starts[1:] = segment[1:] != segment[:-1]
-            gaps = encode_gaps(position - starts[segment], entry_starts)
-            counts = torch.bincount(segment, minlength=len(starts))
             pending.append(
-                (
-                    [name for _, name in self.segments[fqn]],
-                    counts,
-                    gaps,
-                    changed_values.view(torch.uint8),
-                    segment_names,
-                    sums,
+                self._part(
+                    fqn,
+                    self.segments[fqn],
+                    torch.cat(changed) + offset,
+                    torch.cat(changed_values),
+                    (segment_names, sums),
                 )
             )
-            num_pending += len(gaps)
+            num_pending += len(pending[-1][2])
+            if scales:
+                # This rank sends the scales of the blocks that start in its rows
+                first_row = offset // num_cols
+                first_block, end_block = (
+                    -(-row // FP8_BLOCK[0]) for row in (first_row, first_row + len(shard) // num_cols)
+                )
+                previous_scales, current_scales = (torch.cat(side) for side in zip(*scales))
+                if gathered is None:
+                    first_block = first_row // FP8_BLOCK[0]
+                else:
+                    previous_scales, current_scales = (
+                        s[first_block:end_block] for s in (previous_scales, current_scales)
+                    )
+                if current_scales.numel():
+                    pending.append(self._scale_part(fqn, first_block, previous_scales, current_scales))
             if num_pending >= self.PART:
                 flush()
                 num_pending = 0
         if pending:
             flush()
-        return bool(scales_changed)
 
 
 def _model_dir(model: nn.Module) -> Path:
@@ -652,20 +687,19 @@ def _wait(writes: list[Future]) -> None:
             write.exception()
         writes.clear()
 
-    def _write_from_snapshot(self, pool: ThreadPoolExecutor, write, step: int, step_dir: Path, full: bool) -> None:
-        """Write this rank's slices of the last broadcast's full checkpoint in the background. A periodic anchor
-        of a GPU snapshot goes through the host buffer, so the next encode does not wait for the write; other
-        writes read the snapshot itself (a CPU snapshot, or a GPU one through pinned staging chunks), and the
-        next encode waits for them."""
+    def _write_periodic_anchor(self, step: int, step_dir: Path) -> None:
+        """Write this rank's slices of the last broadcast's full checkpoint in the background. A GPU snapshot
+        goes through the host buffer, so the next encode does not wait for the write; a CPU snapshot is
+        written from itself, and the next encode waits for that."""
         encoder = self.encoder
         assert encoder is not None
-        if encoder.on_gpu and not full:
+        if encoder.on_gpu:
             self._wait(self.anchor_writes)
             encoder.copy_snapshot_to_host()
             parts, writes = encoder.anchor_parts(encoder.host), self.anchor_writes
         else:
             parts, writes = encoder.anchor_parts(encoder.snapshot), self.snapshot_writes
-        writes.append(self._submit(pool, write, step, step_dir, parts))
+        writes.append(self._submit(self.anchors, self._write_anchor, step, step_dir, parts))
 
     def _encode(
         self,
@@ -677,14 +711,11 @@ def _encode(
         weights_read: threading.Event,
         anchor: bool,
     ) -> None:
-        """Encode thread: diff v{step}, handing the delta to the writer in parts as it is encoded. With fp8
-        deltas, the parts are held until the ranks agree that no block scale changed, else the step is a
-        full checkpoint."""
+        """Encode thread: diff v{step}, handing the delta to the writer in parts as it is encoded."""
         encoder = self.encoder
         assert encoder is not None
         self._wait(self.snapshot_writes)
         start = time.perf_counter()
-        held: list[tuple] = []
         num_parts = 0
 
         def emit(part: tuple) -> None:
@@ -694,34 +725,17 @@ def emit(part: tuple) -> None:
             self._submit(self.writes, self._write_part, step_dir, num_parts, part)
             num_parts += 1
 
-        encoder.encode(shards, ready, weights_read, held.append if encoder.fp8 else emit)
+        encoder.encode(shards, ready, weights_read, emit)
         # Free the gathered fp8 tensors now, not after the anchor copy below, which may overlap the next step's
         # memory peak.
         shards.clear()
-        if encoder.fp8:
-            if self._scales_changed_anywhere(step_dir, encoder.scales_changed):
-                self.logger.info(f"fp8 block scales changed at v{step}: publishing a full checkpoint")
-                self._write_from_snapshot(self.writes, self._write_full, step, step_dir, full=True)
-                return
-            for part in held:
-                emit(part)
         if num_parts == 0 and self.world.is_master:
             self._submit(self.writes, self._start_delta, step, base_step, step_dir)
-        self._submit(self.writes, self._finish_delta, step, base_step, step_dir, num_parts, time.perf_counter() - start)
+        encode_time = time.perf_counter() - start
+        args = (step, base_step, step_dir, num_parts, encode_time, encoder.scale_changes)
+        self._submit(self.writes, self._finish_delta, *args)
         if anchor:
-            self._write_from_snapshot(self.anchors, self._write_anchor, step, step_dir, full=False)
-
-    def _scales_changed_anywhere(self, step_dir: Path, changed: bool) -> bool:
-        """Agree through the step dir whether any rank's fp8 block scales changed: the encode thread runs no
-        collectives, and the training thread may be waiting for a batch that needs this step."""
-        delta_dir = step_dir / DELTA_DIR
-        delta_dir.mkdir(exist_ok=True)
-        _write_json(delta_dir / f"rank_{self.world.rank}.scales", int(changed))
-        flags = [delta_dir / f"rank_{rank}.scales" for rank in range(self.world.world_size)]
-        while not all(flag.exists() for flag in flags):
-            time.sleep(0.05)
-            revalidate(delta_dir)
-        return any(flag.read_text() == "1" for flag in flags)
+            self._write_periodic_anchor(step, step_dir)
 
     def _start_delta(self, step: int, base_step: int, step_dir: Path) -> None:
         """The manifest comes first, so engines can read part files as they appear."""
@@ -742,11 +756,16 @@ def _write_part(self, step_dir: Path, index: int, part: tuple) -> None:
             path, entries, gaps.numpy(), values.numpy(), self.encoder.hf_tensors, checksums
         )
 
-    def _finish_delta(self, step: int, base_step: int, step_dir: Path, num_parts: int, encode_time: float) -> None:
+    def _finish_delta(
+        self, step: int, base_step: int, step_dir: Path, num_parts: int, encode_time: float, scale_changes: int
+    ) -> None:
         """After this rank's last part: its parts marker. The master then waits for every rank's marker and
         commits the step. Runs on the writer thread, without collectives."""
         delta_dir = step_dir / DELTA_DIR
-        _write_json(delta_dir / rank_parts_marker(self.world.rank), {"parts": num_parts, "bytes": self.part_bytes})
+        _write_json(
+            delta_dir / rank_parts_marker(self.world.rank),
+            {"parts": num_parts, "bytes": self.part_bytes, "scale_changes": scale_changes},
+        )
         self.part_bytes = 0
         if not self.world.is_master:
             return
@@ -755,10 +774,16 @@ def _finish_delta(self, step: int, base_step: int, step_dir: Path, num_parts: in
             time.sleep(0.05)
             revalidate(delta_dir)
         WeightSender._finish(self, step, step_dir)
-        num_bytes = sum(json.loads(marker.read_text())["bytes"] for marker in markers)
+        written = [json.loads(marker.read_text()) for marker in markers]
+        num_bytes = sum(marker["bytes"] for marker in written)
+        scales = (
+            f", {sum(marker['scale_changes'] for marker in written)} fp8 block scales changed"
+            if self.encoder.fp8
+            else ""
+        )
         self.logger.info(
             f"Sparse weight delta v{step} (base v{base_step}): {num_bytes / 2**20:.1f} MiB, "
-            f"encode {encode_time:.2f}s (parts written as encoded)"
+            f"encode {encode_time:.2f}s (parts written as encoded){scales}"
         )
 
     def _write_full(self, step: int, step_dir: Path, parts: list) -> None:
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index d4ba40bbb75..ad37c5327f1 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -15,22 +15,18 @@
 import itertools
 import json
 import os
-import re
 import threading
 import time
-from collections import defaultdict
-from collections.abc import Callable, Iterable, Iterator
+from collections.abc import Callable, Iterable
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import ExitStack
 from itertools import groupby
-from math import prod
 from pathlib import Path
 from typing import BinaryIO
 
 import numpy as np
 import pyzstd
 import torch
-from safetensors import safe_open
 from safetensors.numpy import save
 
 DELTA_DIR = "delta"
@@ -302,59 +298,42 @@ def write_anchor_shards(
     write_file(step_dir / f"anchor_rank{rank}.done", str(world_size).encode())
 
 
-def safetensors_index(directory: Path) -> dict[str, list[tuple[int, Path, torch.dtype, int]]]:
-    """Tensor name -> its slices ``(flat start, file, dtype, data offset)`` in the safetensors files in
-    ``directory``: a HF checkpoint (one slice at 0), or an anchor (``"{name}@{start}"`` slices)."""
-    index = defaultdict(list)
-    for path in directory.glob("*.safetensors"):
+def checkpoint_index(directory: Path) -> dict[str, tuple[list[int], torch.dtype, list[tuple[int, int, Path, int]]]]:
+    """Tensor name -> ``(shape, dtype, slices)`` of the safetensors files in ``directory``, with slices
+    ``(flat start, numel, file, data offset)``: a HF checkpoint (one slice at 0), or an anchor
+    (``"{name}@{start}"`` slices, full shapes in the metadata)."""
+    index = {}
+    for path in sorted(directory.glob("*.safetensors")):
         with open(path, "rb") as file:
             header_size = int.from_bytes(file.read(8), "little")
             header = json.loads(file.read(header_size))
-        header.pop("__metadata__", None)
+        shapes = json.loads(header.pop("__metadata__", {}).get("shapes", "{}"))
         for key, info in header.items():
             name, _, start = key.partition("@")
             dtype = getattr(torch, SAFETENSORS_DTYPES[info["dtype"]])
-            index[name].append((int(start or 0), path, dtype, 8 + header_size + info["data_offsets"][0]))
-    return dict(index)
-
-
-def anchor_index(step_dir: Path) -> tuple[dict[str, list[tuple[Path, str, int]]], dict[str, tuple[list[int], str]]]:
-    """Where each tensor's slices are in an anchor written by ``write_anchor_shards``, and its shape and dtype."""
-    slices: dict[str, list[tuple[Path, str, int]]] = defaultdict(list)
-    tensors: dict[str, tuple[list[int], str]] = {}
-    for path in sorted(step_dir.glob("model-rank*.safetensors")):
-        with safe_open(path, framework="pt") as file:
-            shapes = json.loads(file.metadata()["shapes"])
-            for key in file.keys():
-                name, start = key.rsplit("@", 1)
-                slices[name].append((path, key, int(start)))
-                dtype = str(file.get_slice(key).get_dtype())
-                tensors[name] = (shapes[name], SAFETENSORS_DTYPES[dtype])
-    return dict(slices), tensors
-
-
-def _numeric_order(name: str) -> list:
-    return [int(part) if part.isdigit() else part for part in re.split(r"(\d+)", name)]
-
-
-def anchor_tensors(
-    step_dir: Path, names: Iterable[str] | None = None, index=None
-) -> Iterator[tuple[str, torch.Tensor]]:
-    """Assemble the full tensors of an anchor (all, or ``names``)."""
-    slices, tensors = index or anchor_index(step_dir)
-    # Layer by layer (numeric order), so the engine's layerwise reload holds few layers at a time
-    for name in sorted(slices, key=_numeric_order) if names is None else names:
-        shape, dtype = tensors[name]
-        tensor = torch.empty(prod(shape), dtype=getattr(torch, dtype))
-        for path, key, start in slices[name]:
-            with safe_open(path, framework="pt") as file:
-                values = file.get_tensor(key)
-            tensor[start : start + values.numel()] = values
-        yield name, tensor.view(shape)
-    if names is None:
-        for path in step_dir.glob("model-rank*.safetensors"):
-            with open(path, "rb") as file:
-                drop_page_cache(file.fileno())
+            _, _, slices = index.setdefault(name, (shapes.get(name, info["shape"]), dtype, []))
+            numel = (info["data_offsets"][1] - info["data_offsets"][0]) // dtype.itemsize
+            slices.append((int(start or 0), numel, path, 8 + header_size + info["data_offsets"][0]))
+    return index
+
+
+def read_checkpoint_range(entry: tuple, begin: int, end: int, out: np.ndarray) -> None:
+    """Read flat elements ``[begin, end)`` of a ``checkpoint_index`` entry into the bytes ``out``, and drop
+    them from the page cache (checkpoints are read once)."""
+    _, dtype, slices = entry
+    itemsize = dtype.itemsize
+    for start, numel, path, data_offset in slices:
+        low, high = max(begin, start), min(end, start + numel)
+        if low >= high:
+            continue
+        with open(path, "rb") as file:
+            done, size, offset = 0, (high - low) * itemsize, data_offset + (low - start) * itemsize
+            target = out[(low - begin) * itemsize : (high - begin) * itemsize]
+            while done < size:
+                read = os.preadv(file.fileno(), [target[done:]], offset + done)
+                assert read > 0, f"short read of {path}"
+                done += read
+            os.posix_fadvise(file.fileno(), offset, size, os.POSIX_FADV_DONTNEED)
 
 
 def anchor_complete(step_dir: Path) -> bool:

From 255b0d0956c8065f38f69698fdebe75c168f3c9b Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sat, 17 Oct 2026 01:41:09 +0000
Subject: [PATCH 44/88] docs(weights): fp8 scales in the delta, slice-local
 anchor reloads

---
 docs/scaling.md                                  | 2 +-
 src/prime_rl/inference/vllm/worker/filesystem.py | 1 -
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index d722c2b2198..7bd421ec3fb 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -276,7 +276,7 @@ By default, the trainer and inference worker each allocate one transfer arena. S
 
 ### Sparse filesystem weight broadcast
 
-With `[weight_broadcast] type = "filesystem"` and `sparse = true`, the trainer publishes only the elements of the weights that changed since the previous broadcast, in the engine's weight format (bf16, or fp8 for engines that quantize online to `fp8_per_block`). Each trainer rank diffs its own shards against a copy of the previous broadcast (`snapshot = "auto"` keeps it on the GPU when it fits next to the first training step's peak memory, else in pinned host memory), so no weights are gathered. The diff runs on a side CUDA stream while the next step trains; the trainer only waits for it before the next step's final backward (which may already run the optimizer step). With AdamW at RL learning rates about 1% of bf16 elements change per step. A full checkpoint (anchor) is written in the background every `anchor_interval` steps (default 25); the first broadcast is the checkpoint the engines booted from, and the trainer reads it back (its own shard of each tensor) to encode the first delta. With a GPU snapshot, anchors are copied to a host buffer of the snapshot's size (2 B per parameter the rank holds), allocated at the first anchor; on hosts without that much free memory (e.g. next to a fully offloaded Muon optimizer), use a larger `anchor_interval`. Engines read and decode a delta while they keep serving, patch their weights in place during the pause and verify every step against trainer checksums. An engine that is not at a delta's base step reloads the last anchor and replays the deltas since. Weights the engine transforms after loading other than online `fp8_per_block`, and LoRA, are not supported.
+With `[weight_broadcast] type = "filesystem"` and `sparse = true`, the trainer publishes only the elements of the weights that changed since the previous broadcast, in the engine's weight format (bf16, or fp8 for engines that quantize online to `fp8_per_block`). Each trainer rank diffs its own shards against a copy of the previous broadcast (`snapshot = "auto"` keeps it on the GPU when it fits next to the first training step's peak memory, else in pinned host memory), so no weights are gathered. The diff runs on a side CUDA stream while the next step trains; the trainer only waits for it before the next step's final backward (which may already run the optimizer step). With AdamW at RL learning rates about 1% of bf16 elements change per step. A full checkpoint (anchor) is written in the background every `anchor_interval` steps (default 25); the first broadcast is the checkpoint the engines booted from, and the trainer reads it back (its own shard of each tensor) to encode the first delta. With a GPU snapshot, anchors are copied to a host buffer of the snapshot's size (2 B per parameter the rank holds), allocated at the first anchor; on hosts without that much free memory (e.g. next to a fully offloaded Muon optimizer), use a larger `anchor_interval`. Engines read and decode a delta while they keep serving, patch their weights in place during the pause and verify every step against trainer checksums. For fp8 engines, changed block scales are part of the delta. An engine that is not at a delta's base step reloads the last anchor, each worker reading only the rows it holds, and replays the deltas since. Weights the engine transforms after loading other than online `fp8_per_block`, and LoRA, are not supported.
 
 ### Custom Templates
 
diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 90702cd89ff..6863fcd4597 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -59,7 +59,6 @@ class FileSystemWeightUpdateWorker(Worker):
     weight_step: int | None = None
     prefetched: tuple[str, Future] | None = None
     index_map: CheckpointIndexMap | None = None
-    wrote_layout = False
     host_buffer: torch.Tensor | None = None
 
     def liveness_probe(self) -> None:

From 0e90805416e7f52d4f5044ad6ee1743fcb78c7f2 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sat, 17 Oct 2026 02:46:00 +0000
Subject: [PATCH 45/88] perf(weights): checksum same-shape writes into a param
 in one batch (fp8 block scales of every expert)

---
 .../inference/vllm/worker/weight_transfer.py  | 26 +++++++++++++------
 1 file changed, 18 insertions(+), 8 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py
index 05984415f49..a8672442612 100644
--- a/src/prime_rl/inference/vllm/worker/weight_transfer.py
+++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py
@@ -151,11 +151,16 @@ def held_checksums(self) -> tuple[list[str], torch.Tensor, torch.Tensor]:
                 )
             )
             totals.index_add_(0, ids, num_cols * by_row + firsts * plain + by_col)
-        for name_id, num_cols, first, rows, cols, index, offset, *strides in other_writes:
-            view = as_bits(self.params[index].data).view(-1).as_strided((rows, cols), strides, offset)
-            row_weights = first + num_cols * torch.arange(rows, device=view.device)
-            totals[name_id] += (row_weights * view.sum(1, dtype=torch.int64)).sum()
-            totals[name_id] += (torch.arange(cols, device=view.device) * view.sum(0, dtype=torch.int64)).sum()
+        # Writes of one shape into one param at once, e.g. the block scales of every expert
+        for (index, rows, cols, row_stride, col_stride), (ids, num_cols, firsts, offsets) in other_writes.items():
+            flat = as_bits(self.params[index].data).view(-1)
+            row = torch.arange(rows, device=flat.device)[:, None]
+            col = torch.arange(cols, device=flat.device)
+            for batch in range(0, len(ids), max(1, (1 << 24) // (rows * cols))):
+                part = slice(batch, batch + max(1, (1 << 24) // (rows * cols)))
+                values = flat[offsets[part, None, None] + row * row_stride + col * col_stride].long()
+                weights = firsts[part, None, None] + num_cols[part, None, None] * row + col
+                totals.index_add_(0, ids[part], (values * weights).sum((1, 2)))
         return names, totals, partial
 
     def _plan_checksums(self):
@@ -169,7 +174,7 @@ def _plan_checksums(self):
         ]
         device = self.params[0].device
         ranges: dict[int, list[list[int]]] = defaultdict(lambda: [[], [], [], [], []])
-        other_writes = []
+        other_writes: dict[tuple, list[list[int]]] = defaultdict(lambda: [[], [], [], []])
         for name_id, name in enumerate(names):
             num_cols = self.num_cols[name]
             for row0, rows, col0, cols, index, offset, row_stride, col_stride in self.records[name]:
@@ -181,13 +186,18 @@ def _plan_checksums(self):
                     ):
                         values.append(value)
                 else:
-                    first = row0 * num_cols + col0 + 1
-                    other_writes.append((name_id, num_cols, first, rows, cols, index, offset, row_stride, col_stride))
+                    batch = other_writes[(index, rows, cols, row_stride, col_stride)]
+                    for values, value in zip(batch, (name_id, num_cols, row0 * num_cols + col0 + 1, offset)):
+                        values.append(value)
         row_ranges = {
             index: tuple(torch.tensor(values, device=device) for values in lists) for index, lists in ranges.items()
         }
         held = [sum(rows * cols for _, rows, _, cols, *_ in self.records[name]) for name in names]
         partial = torch.tensor([h < self.numel[name] for h, name in zip(held, names)], device=device)
+        other_writes = {
+            shape: tuple(torch.tensor(values, device=device) for values in lists)
+            for shape, lists in other_writes.items()
+        }
         return names, partial, row_ranges, other_writes
 
     def fp8_layout(self, model: Module, read_tensor) -> dict[str, str]:

From 48de997528187c4ab2931a7aa2420d398b2ce683 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sat, 17 Oct 2026 08:44:13 +0000
Subject: [PATCH 46/88] fix(weights): fault the sparse host buffer in before
 training; skip an anchor whose copy holds the trainer

---
 docs/scaling.md                           |  2 +-
 src/prime_rl/transports/weights/sparse.py | 69 +++++++++++++++++------
 2 files changed, 52 insertions(+), 19 deletions(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index 7bd421ec3fb..c3a8b83a496 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -276,7 +276,7 @@ By default, the trainer and inference worker each allocate one transfer arena. S
 
 ### Sparse filesystem weight broadcast
 
-With `[weight_broadcast] type = "filesystem"` and `sparse = true`, the trainer publishes only the elements of the weights that changed since the previous broadcast, in the engine's weight format (bf16, or fp8 for engines that quantize online to `fp8_per_block`). Each trainer rank diffs its own shards against a copy of the previous broadcast (`snapshot = "auto"` keeps it on the GPU when it fits next to the first training step's peak memory, else in pinned host memory), so no weights are gathered. The diff runs on a side CUDA stream while the next step trains; the trainer only waits for it before the next step's final backward (which may already run the optimizer step). With AdamW at RL learning rates about 1% of bf16 elements change per step. A full checkpoint (anchor) is written in the background every `anchor_interval` steps (default 25); the first broadcast is the checkpoint the engines booted from, and the trainer reads it back (its own shard of each tensor) to encode the first delta. With a GPU snapshot, anchors are copied to a host buffer of the snapshot's size (2 B per parameter the rank holds), allocated at the first anchor; on hosts without that much free memory (e.g. next to a fully offloaded Muon optimizer), use a larger `anchor_interval`. Engines read and decode a delta while they keep serving, patch their weights in place during the pause and verify every step against trainer checksums. For fp8 engines, changed block scales are part of the delta. An engine that is not at a delta's base step reloads the last anchor, each worker reading only the rows it holds, and replays the deltas since. Weights the engine transforms after loading other than online `fp8_per_block`, and LoRA, are not supported.
+With `[weight_broadcast] type = "filesystem"` and `sparse = true`, the trainer publishes only the elements of the weights that changed since the previous broadcast, in the engine's weight format (bf16, or fp8 for engines that quantize online to `fp8_per_block`). Each trainer rank diffs its own shards against a copy of the previous broadcast (`snapshot = "auto"` keeps it on the GPU when it fits next to the first training step's peak memory, else in pinned host memory), so no weights are gathered. The diff runs on a side CUDA stream while the next step trains; the trainer only waits for it before the next step's final backward (which may already run the optimizer step). With AdamW at RL learning rates about 1% of bf16 elements change per step. A full checkpoint (anchor) is written in the background every `anchor_interval` steps (default 25); the first broadcast is the checkpoint the engines booted from, and the trainer reads it back (its own shard of each tensor) to encode the first delta. Each rank allocates and faults in a host buffer of the snapshot's size (2 B per parameter the rank holds) before training: the CPU snapshot, or the anchor copies of a GPU snapshot. Trainer hosts need that much free memory (e.g. next to a fully offloaded Muon optimizer). An anchor whose copy holds back the next step's diff for more than 30 s is skipped with a warning; engines then replay from the previous anchor. Engines read and decode a delta while they keep serving, patch their weights in place during the pause and verify every step against trainer checksums. For fp8 engines, changed block scales are part of the delta. An engine that is not at a delta's base step reloads the last anchor, each worker reading only the rows it holds, and replays the deltas since. Weights the engine transforms after loading other than online `fp8_per_block`, and LoRA, are not supported.
 
 ### Custom Templates
 
diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 13fa0bbf5aa..94530b05946 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -1,4 +1,5 @@
 import json
+import mmap
 import threading
 import time
 from bisect import bisect_left, bisect_right
@@ -58,11 +59,11 @@ class SparseDeltaEncoder:
     next step trains, without collectives; it only reads the weights before
     ``weights_read`` is set, and hands the delta to the writer in parts as it
     goes. The snapshot is placed at the first broadcast after a training step: on
-    the GPU when it fits next to that step's peak memory, else in one page-locked
-    host buffer (locked only between training steps: registering it can stall the
+    the GPU when it fits next to that step's peak memory, else in the host buffer,
+    page-locked (only between training steps: registering it can stall the
     trainer's CUDA calls). It starts as the checkpoint the engines booted from,
-    read back from storage. With a GPU snapshot, a pageable host buffer holds
-    anchor copies, allocated at the first periodic anchor.
+    read back from storage. With a GPU snapshot, the host buffer holds anchor
+    copies. The host buffer is allocated and faulted in at the startup broadcast.
 
     Tensors the engines hold in fp8 (``set_fp8_layout``: HF name -> scale format,
     from the engine layout) are quantized with vLLM's own per-block quantizer,
@@ -230,16 +231,24 @@ def _gather_previous(self, key: str, value: DTensor) -> torch.Tensor:
     def _bits(self, key: str) -> torch.dtype:
         return bits_dtype(str(self.wire_dtypes[key]).removeprefix("torch."))
 
-    def _allocate_host(self) -> None:
-        """One host buffer for all shards (the CPU snapshot, or anchor copies of a GPU snapshot)."""
+    def allocate_host(self) -> None:
+        """One host buffer for all shards (the CPU snapshot, or anchor copies of a GPU snapshot), faulted in
+        here, before training: on a node short of free memory, faulting it in during an anchor copy waits for
+        the kernel to reclaim page cache, for minutes, while the trainer waits for that copy."""
+        start = time.perf_counter()
         sizes = {key: -(-numel * self.wire_dtypes[key].itemsize // 8) * 8 for key, (numel, _) in self.shards.items()}
-        self.host_buffer = torch.from_numpy(np.empty(sum(sizes.values()), np.uint8))
+        flags = mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS | mmap.MAP_POPULATE
+        self.host_buffer = torch.from_numpy(np.frombuffer(mmap.mmap(-1, sum(sizes.values()), flags), np.uint8))
         offset = 0
         for key, (numel, _) in self.shards.items():
             self.host[key] = self.host_buffer[offset : offset + numel * self.wire_dtypes[key].itemsize].view(
                 self._bits(key)
             )
             offset += sizes[key]
+        get_logger().info(
+            f"Allocated the {self.host_buffer.numel() / 2**30:.1f} GiB sparse weight host buffer in "
+            f"{time.perf_counter() - start:.1f}s"
+        )
 
     def _page_lock(self) -> None:
         """Page-lock the host buffer, only while no training step runs: registering a large buffer can
@@ -269,16 +278,17 @@ def track(self, shards: list[tuple]) -> None:
         self.shards = {key: (shard.numel(), offset) for key, _, shard, offset, _ in shards}
         self.key_fqn = {key: fqn for key, fqn, *_ in shards}
 
-    def copy_snapshot_to_host(self) -> None:
-        """Copy a GPU snapshot (the last broadcast's weights) into the host buffer (allocated at the first
-        anchor), for an anchor, through the pinned staging buffer. Callers wait for earlier readers of the
-        host buffer first."""
+    def copy_snapshot_to_host(self, give_up: Callable[[], bool]) -> bool:
+        """Copy a GPU snapshot (the last broadcast's weights) into the host buffer, for an anchor, through the
+        pinned staging buffer; False if ``give_up()`` turned true before it finished. Callers wait for earlier
+        readers of the host buffer first."""
         assert self.on_gpu
-        if self.host_buffer is None:
-            self._allocate_host()
         with torch.cuda.stream(self.stream):
             for key, bits in self.snapshot.items():
+                if give_up():
+                    return False
                 self._to_host(bits, self.host[key])
+        return True
 
     def anchor_parts(self, source: dict[str, torch.Tensor]) -> list[tuple[str, list[int], int, torch.Tensor]]:
         """This rank's slices of a full HF checkpoint, from ``source`` (key -> flat wire bits of its shard)."""
@@ -294,7 +304,7 @@ def anchor_parts(self, source: dict[str, torch.Tensor]) -> list[tuple[str, list[
     @torch.no_grad()
     def initialize(self, model: nn.Module, is_master: bool, peak_reserved: int, boot_dir: Path) -> None:
         """At the first delta (v1), on the training thread: place the snapshot against the measured training
-        peak and allocate it (the host buffer only for a CPU snapshot). The snapshot starts as v0, the
+        peak and allocate it (a CPU snapshot is the host buffer). The snapshot starts as v0, the
         checkpoint in ``boot_dir`` the engines booted from, read as the encode goes (tensors gathered for fp8
         are read here, before ``local_shards`` gathers them)."""
         self.track(self.local_shards(model, is_master))
@@ -306,7 +316,6 @@ def initialize(self, model: nn.Module, is_master: bool, peak_reserved: int, boot
                 key: torch.empty(numel, dtype=self._bits(key), device=device) for key, (numel, _) in self.shards.items()
             }
         else:
-            self._allocate_host()
             self._page_lock()
             self.snapshot = dict(self.host)
         self.boot_files = checkpoint_index(boot_dir)
@@ -577,6 +586,9 @@ class SparseFileSystemWeightSender(FileSystemWeightSender):
     files are on disk."""
 
     gathers_full_weights = False
+    # Seconds an anchor copy may hold back the next encode before it gives up (that anchor is skipped): the trainer
+    # waits for that encode, well below the offload pipeline's timeout, which other ranks hit while waiting for it
+    ANCHOR_COPY_WAIT = 30.0
 
     def __init__(self, output_dir: Path, config: FileSystemWeightBroadcastConfig):
         super().__init__(output_dir, config)
@@ -594,6 +606,7 @@ def __init__(self, output_dir: Path, config: FileSystemWeightBroadcastConfig):
         self.anchor_writes: list[Future] = []
         self.errors: list[BaseException] = []
         self.weights_read: threading.Event | None = None
+        self.encode_queued_at: float | None = None
         self.part_bytes = 0
         # The checkpoint the engines hold before the first delta: the model's (None), or a resumed run's first
         # broadcast
@@ -635,6 +648,10 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
             self.encoder = SparseDeltaEncoder(model, self.config.snapshot)
         encoder = self.encoder
         base_step, self.last_step = self.last_step, step
+        if base_step is None:
+            shards = encoder.local_shards(model, self.world.is_master)
+            encoder.track(shards)
+            encoder.allocate_host()
         if base_step is None and step == 0:
             # The run's first broadcast is the checkpoint the engines booted from.
             self.anchor_step = step
@@ -646,8 +663,6 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
         if base_step is None:
             # A resumed run's startup broadcast: a full checkpoint, written from the weights before training. The
             # next broadcast's delta is against it.
-            shards = encoder.local_shards(model, self.world.is_master)
-            encoder.track(shards)
             source = {key: as_bits(shard.reshape(-1).to(encoder.wire_dtypes[key])) for key, _, shard, _, _ in shards}
             torch.cuda.synchronize()
             self._write_full(step, step_dir, encoder.anchor_parts(source))
@@ -669,6 +684,7 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
         ready = torch.cuda.Event()
         ready.record()
         self.weights_read = threading.Event()
+        self.encode_queued_at = time.monotonic()
         args = (step, base_step, step_dir, shards, ready, self.weights_read, anchor)
         self._submit(self.encodes, self._encode, *args)
 
@@ -694,8 +710,21 @@ def _write_periodic_anchor(self, step: int, step_dir: Path) -> None:
         encoder = self.encoder
         assert encoder is not None
         if encoder.on_gpu:
+            # The next encode waits for this copy (it advances the snapshot)
+            def held_too_long() -> bool:
+                queued_at = self.encode_queued_at
+                return queued_at is not None and time.monotonic() - queued_at > self.ANCHOR_COPY_WAIT
+
+            while not all(write.done() for write in self.anchor_writes) and not held_too_long():
+                time.sleep(0.1)
+            if held_too_long() or not encoder.copy_snapshot_to_host(held_too_long):
+                self.logger.warning(
+                    f"Skipped the {step_dir.name} anchor: its copy to host memory held the next encode for over "
+                    f"{self.ANCHOR_COPY_WAIT:.0f}s"
+                )
+                write_file(step_dir / f"anchor_rank{self.world.rank}.skipped", b"")
+                return
             self._wait(self.anchor_writes)
-            encoder.copy_snapshot_to_host()
             parts, writes = encoder.anchor_parts(encoder.host), self.anchor_writes
         else:
             parts, writes = encoder.anchor_parts(encoder.snapshot), self.snapshot_writes
@@ -714,6 +743,7 @@ def _encode(
         """Encode thread: diff v{step}, handing the delta to the writer in parts as it is encoded."""
         encoder = self.encoder
         assert encoder is not None
+        self.encode_queued_at = None
         self._wait(self.snapshot_writes)
         start = time.perf_counter()
         num_parts = 0
@@ -798,6 +828,9 @@ def _write_anchor(self, step: int, step_dir: Path, parts: list) -> None:
         self.logger.debug(f"Wrote the {step_dir.name} anchor in {time.perf_counter() - start:.1f}s")
         if self.world.is_master:
             while not anchor_complete(step_dir):
+                if any(step_dir.glob("anchor_rank*.skipped")):
+                    self.logger.warning(f"The {step_dir.name} anchor is incomplete: a rank skipped it")
+                    return
                 time.sleep(0.1)
             self.anchor_step = max(self.anchor_step or 0, step)
 

From 5cc0591842779a4e0cb0ada78002305df8a0651e Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sat, 17 Oct 2026 13:13:00 +0000
Subject: [PATCH 47/88] fix(weights): pin the engine's delta read buffer at
 twice the delta, so it is pinned once

---
 src/prime_rl/inference/vllm/worker/filesystem.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 6863fcd4597..e043ff7bb29 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -81,7 +81,7 @@ def _read_delta(
         self, step_dir: Path
     ) -> tuple[list[tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]], dict[str, int]]:
         """Read this worker's part of a delta into a pinned host buffer that is reused across steps
-        (allocating pinned memory every step is slow on a loaded host), in chunks of up to 2^26 changes
+        (pinning memory is slow on a loaded host), in chunks of up to 2^26 changes
         of consecutive frames."""
         global _reader_pool
         if _reader_pool is None:
@@ -103,9 +103,16 @@ def _read_delta(
             layout.append((num_bytes, num_bytes + gaps_bytes))
             num_bytes += gaps_bytes + values_bytes
         if self.host_buffer is None or self.host_buffer.numel() < num_bytes:
+            # Twice the delta, so later deltas (within ~20% of each other) fit: pinning stalls for tens of
+            # seconds on a host that must reclaim page cache first, and this read is the update's critical path
             if self.host_buffer is not None:
                 unpin(self.host_buffer)
-            self.host_buffer = pin(num_bytes * 5 // 4)
+            pin_start = time.perf_counter()
+            self.host_buffer = pin(num_bytes * 2)
+            logger.info(
+                f"Pinned a {self.host_buffer.numel() / 2**30:.2f} GiB host buffer for sparse deltas in "
+                f"{time.perf_counter() - pin_start:.2f}s"
+            )
         buffer = self.host_buffer.numpy()
         copies, pinned = [], []
         for chunk, (gaps_at, values_at) in zip(chunks, layout):

From f39c69ed6edc3af6e6bceaa022e0ee0b4fc6eee4 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sat, 17 Oct 2026 16:28:14 +0000
Subject: [PATCH 48/88] feat(weights): sparse deltas from the full-offload
 optimizer step, without a snapshot

---
 src/prime_rl/trainer/optim/offload.py     |   8 +-
 src/prime_rl/trainer/rl/train.py          |   5 +-
 src/prime_rl/transports/weights/sparse.py | 449 +++++++++++++++++-----
 src/prime_rl/utils/sparse_weights.py      | 107 ++++--
 4 files changed, 434 insertions(+), 135 deletions(-)

diff --git a/src/prime_rl/trainer/optim/offload.py b/src/prime_rl/trainer/optim/offload.py
index f0b698c9939..d783cd1c78e 100644
--- a/src/prime_rl/trainer/optim/offload.py
+++ b/src/prime_rl/trainer/optim/offload.py
@@ -719,6 +719,9 @@ def __init__(
         self.optimizer = optimizer
         self._initialized = False
         self._master_weights = master_weights
+        # Called on the H2D stream with each parameter's local shard and its new values (pinned host) instead of
+        # the copy, which it must do; the shard still holds the previous values
+        self.weight_hook: Callable[[torch.Tensor, torch.Tensor], None] | None = None
         self._chunks = self._build_chunks()
         # Reuse the transfer streams across steps: fresh streams each step land
         # their H2D/D2H staging in new per-stream allocator pools, growing
@@ -797,7 +800,10 @@ def _update_compute_weights(
                     raise TypeError(
                         f"Full-offload H2D source dtype mismatch: expected {master.compute_dtype}, got {source.dtype}"
                     )
-                local_param.copy_(source, non_blocking=True)
+                if self.weight_hook is None:
+                    local_param.copy_(source, non_blocking=True)
+                else:
+                    self.weight_hook(local_param, source)
 
     def _step_native_cpu_adamw_chunk(
         self,
diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py
index 33e9b4f0e39..4f89e0b26b1 100644
--- a/src/prime_rl/trainer/rl/train.py
+++ b/src/prime_rl/trainer/rl/train.py
@@ -9,13 +9,14 @@
 # Import environment before any other imports
 # ruff: noqa: I001
 
-from prime_rl.transports.weights import prune_broadcasts_beyond, setup_weight_sender
+from prime_rl.transports.weights import SparseFileSystemWeightSender, prune_broadcasts_beyond, setup_weight_sender
 from prime_rl.utils.act_offloading import maybe_activation_offloading
 import torch
 import torch.distributed as dist
 from torch.profiler import profile, ProfilerActivity, record_function
 from prime_rl.trainer.ckpt import Progress, setup_ckpt_manager
 from prime_rl.trainer.optim import setup_optimizer
+from prime_rl.trainer.optim.offload import FullCPUOffloadOptimizer
 from prime_rl.trainer.scheduler import setup_scheduler
 from prime_rl.configs.trainer import TrainerConfig
 from prime_rl.trainer.rl.data import DataLoader, FakeDataLoader, prepare_micro_batch
@@ -193,6 +194,8 @@ def train(config: TrainerConfig):
             config.model.lora,
         )
         logger.debug(f"Initialized weight broadcast in {format_time(time.perf_counter() - t0)}")
+        if isinstance(weight_sender, SparseFileSystemWeightSender) and isinstance(optimizer, FullCPUOffloadOptimizer):
+            weight_sender.diff_in_optimizer(optimizer)
 
     if parallel_dims.cp_enabled:
         setup_context_parallel(model, config.model, parallel_dims)
diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 94530b05946..2dd2ff0e5ce 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -1,5 +1,6 @@
 import json
 import mmap
+import queue
 import threading
 import time
 from bisect import bisect_left, bisect_right
@@ -8,6 +9,7 @@
 from concurrent.futures import Future, ThreadPoolExecutor
 from math import prod
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import numpy as np
 import torch
@@ -25,6 +27,7 @@
     DELTA_DIR,
     FP8_BLOCK,
     MANIFEST,
+    AnchorWriter,
     anchor_complete,
     as_bits,
     bits_dtype,
@@ -45,6 +48,9 @@
 )
 from prime_rl.utils.weights import convert_state_dict_to_hf, resolve_fqn, resolve_wire_dtype
 
+if TYPE_CHECKING:
+    from prime_rl.trainer.optim.offload import FullCPUOffloadOptimizer
+
 
 class SparseDeltaEncoder:
     """Rank-local deltas of the trainer's weight shards, in HF coordinates.
@@ -73,12 +79,21 @@ class SparseDeltaEncoder:
     (columns a multiple of 128). Their block scales are diffed too, as tensors
     ``"{hf_name}_scale_inv"`` (fp32, one element per 128x128 block) in the same
     delta, so engines patch scales in place.
+
+    With a fully offloaded optimizer (``bind_optimizer``), the tensors stored in its parameters are diffed in its
+    step instead (``optimizer_hook``): each parameter shard still holds the previous broadcast's weights when the
+    step writes its new values, so they need no snapshot, and anchors are written as the step produces them. The
+    encode only checks that those tensors still hold the values the step wrote.
     """
 
     # Elements per transient chunk of the diff (bf16: 128 MiB per temporary)
     CHUNK = 1 << 26
     # Changes per part handed to the writer
     PART = 1 << 26
+    # Elements per transient chunk of the diff in the optimizer pass, which overlaps the final backward
+    HOOK_CHUNK = 1 << 24
+    # GPU slots of the optimizer pass's diff thread, each two of the largest parameter shard
+    HOOK_SLOTS = 2
 
     def __init__(self, model: nn.Module, snapshot: str):
         keep_in_fp32 = getattr(model, "keep_in_fp32_for_weight_transfer", None)
@@ -115,6 +130,7 @@ def __init__(self, model: nn.Module, snapshot: str):
         self.snapshot_placement = snapshot
         # key -> (local numel, global flat offset) of this rank's shard
         self.shards: dict[str, tuple[int, int]] = {}
+        self.key_fqn: dict[str, str] = {}
         # key -> previous wire bits: GPU tensors, or views into the host buffer
         self.snapshot: dict[str, torch.Tensor] = {}
         self.on_gpu: bool | None = None  # None until placed, at the first broadcast after a training step
@@ -126,6 +142,14 @@ def __init__(self, model: nn.Module, snapshot: str):
         self.host: dict[str, torch.Tensor] = {}
         self.pinned = False
         self.initialized = False
+        # Diffed in the optimizer pass (``bind_optimizer``): parameter storage -> (key, fqn, numel, global flat
+        # offset, (shape, stride, storage offset) of the view), and that step's parts, checksums of the written wire bits and anchor
+        self.param_keys: dict[int, list[tuple[str, str, int, int, tuple]]] = defaultdict(list)
+        self.optimizer_keys: set[str] = set()
+        self.step_parts: list[tuple] = []
+        self.diffs: list[Future] = []
+        self.expected: dict[str, torch.Tensor] = {}
+        self.anchor: AnchorWriter | None = None
         self.device = torch.cuda.current_device()
         self.stream = torch.cuda.Stream()
 
@@ -276,7 +300,7 @@ def _place(self, peak_reserved: int) -> None:
     def track(self, shards: list[tuple]) -> None:
         """Record which elements of each tensor this rank's shards hold."""
         self.shards = {key: (shard.numel(), offset) for key, _, shard, offset, _ in shards}
-        self.key_fqn = {key: fqn for key, fqn, *_ in shards}
+        self.key_fqn.update({key: fqn for key, fqn, *_ in shards})
 
     def copy_snapshot_to_host(self, give_up: Callable[[], bool]) -> bool:
         """Copy a GPU snapshot (the last broadcast's weights) into the host buffer, for an anchor, through the
@@ -290,15 +314,24 @@ def copy_snapshot_to_host(self, give_up: Callable[[], bool]) -> bool:
                 self._to_host(bits, self.host[key])
         return True
 
+    def _slices(self, key: str, numel: int, offset: int) -> list[tuple[str, list[int], int, int, int]]:
+        """``(hf_name, shape, flat start in the HF tensor, begin, end in the shard)`` of the HF tensors in a
+        shard of ``numel`` elements at global flat ``offset``."""
+        slices = []
+        for start, name in self.segments[self.key_fqn[key]]:
+            shape = self.hf_tensors[name][0]
+            begin, end = max(start, offset), min(start + prod(shape), offset + numel)
+            if begin < end:
+                slices.append((name, shape, begin - start, begin - offset, end - offset))
+        return slices
+
     def anchor_parts(self, source: dict[str, torch.Tensor]) -> list[tuple[str, list[int], int, torch.Tensor]]:
         """This rank's slices of a full HF checkpoint, from ``source`` (key -> flat wire bits of its shard)."""
         parts = []
         for key, bits in source.items():
-            values, (numel, offset) = bits.view(self.wire_dtypes[key]), self.shards[key]
-            for start, name in self.segments[self.key_fqn[key]]:
-                begin, end = max(start, offset), min(start + prod(self.hf_tensors[name][0]), offset + numel)
-                if begin < end:
-                    parts.append((name, self.hf_tensors[name][0], begin - start, values[begin - offset : end - offset]))
+            values = bits.view(self.wire_dtypes[key])
+            for name, shape, start, begin, end in self._slices(key, *self.shards[key]):
+                parts.append((name, shape, start, values[begin:end]))
         return parts
 
     @torch.no_grad()
@@ -307,7 +340,7 @@ def initialize(self, model: nn.Module, is_master: bool, peak_reserved: int, boot
         peak and allocate it (a CPU snapshot is the host buffer). The snapshot starts as v0, the
         checkpoint in ``boot_dir`` the engines booted from, read as the encode goes (tensors gathered for fp8
         are read here, before ``local_shards`` gathers them)."""
-        self.track(self.local_shards(model, is_master))
+        self.track([shard for shard in self.local_shards(model, is_master) if shard[0] not in self.optimizer_keys])
         self._place(peak_reserved)
         self.staging = torch.empty(1 << 28, dtype=torch.uint8, pin_memory=True)
         if self.on_gpu:
@@ -351,17 +384,20 @@ def _read_boot(self, key: str) -> None:
     @torch.no_grad()
     def encode(
         self, shards: list[tuple], ready: torch.cuda.Event, weights_read: threading.Event, emit: Callable
-    ) -> None:
+    ) -> int:
         """Diff the shards against the snapshot and advance it, on the side stream once ``ready`` (recorded
         after the optimizer step) has passed. Hands the delta to ``emit`` in parts ``(entries, gaps, value
         bytes, new-weight checksum per HF tensor part)`` on the host as they are encoded. Sets ``weights_read``
-        once done reading the weights."""
+        once done reading the weights. Returns the number of changed fp8 block scales."""
         torch.cuda.set_device(self.device)
+        if not self.optimizer_keys:
+            self.scale_changes = 0
         try:
             with torch.cuda.stream(self.stream):
                 self.stream.wait_event(ready)
                 self._diff(shards, emit)
                 self.stream.synchronize()
+            return self.scale_changes
         finally:
             weights_read.set()
 
@@ -460,13 +496,15 @@ def _scale_part(self, fqn: str, first_block: int, previous: torch.Tensor, curren
         )
 
     def _diff(self, shards: list[tuple], emit: Callable) -> None:
-        """Diff ``shards`` against the snapshot on the current stream, writing the new bits to the
-        snapshot, in chunks of ``CHUNK`` elements (fp8: whole 128-row blocks, and their block scales).
-        Copies the delta to the host and hands it to ``emit`` every ``PART`` changes."""
-        device = torch.device(f"cuda:{self.device}")
-        pending: list[tuple] = []
-        num_pending = 0
-        self.scale_changes = 0
+        """Diff ``shards`` against the snapshot on the current stream, writing the new bits to the snapshot.
+        Shards diffed in the optimizer pass are only checked against the values that pass wrote; their parts
+        go first. Copies the delta to the host and hands it to ``emit`` every ``PART`` changes."""
+        for diff in self.diffs:
+            diff.result()
+        self.diffs.clear()
+        pending, self.step_parts = self.step_parts, []
+        num_pending = sum(len(part[2]) for part in pending)
+        checks: list[tuple[str, torch.Tensor, torch.Tensor]] = []
 
         def flush() -> None:
             names = [name for part in pending for name in part[0]]
@@ -480,90 +518,229 @@ def flush() -> None:
 
         for key, fqn, shard, offset, gathered in shards:
             shard = shard.reshape(-1)
+            if key in self.optimizer_keys:
+                if key not in self.expected:
+                    raise RuntimeError(f"{key} was not written by the optimizer step before this broadcast")
+                checks.append((key, self.expected.pop(key), self._checksum(as_bits(shard.to(self.wire_dtypes[key])))))
+                continue
             if key in self.boot_pending:
                 self._read_boot(key)
                 self.boot_pending.discard(key)
-            previous = out = self.snapshot[key]
-            wire = self.wire_dtypes[key]
-            segment_names, begins, ends, shifts = self._segments(key, self.segments[fqn], offset, shard.numel())
-            sums = torch.zeros(len(segment_names), dtype=torch.int64, device=device)
-            changed, changed_values, scales = [], [], []
+            parts = self._diff_shard(key, fqn, shard, offset, gathered, self.snapshot[key], self.snapshot[key])
+            pending += parts
+            num_pending += sum(len(part[2]) for part in parts)
+            if num_pending >= self.PART:
+                flush()
+                num_pending = 0
+        if pending:
+            flush()
+        if checks:
+            keys, expected, live = zip(*checks)
+            changed = [key for key, same in zip(keys, torch.cat(expected).eq(torch.cat(live)).tolist()) if not same]
+            if changed:
+                raise RuntimeError(f"Trainer weights changed outside the optimizer step, not sent: {changed}")
+
+    def _diff_shard(
+        self,
+        key: str,
+        fqn: str,
+        shard: torch.Tensor,
+        offset: int,
+        gathered: tuple | None,
+        previous: torch.Tensor,
+        out: torch.Tensor | None,
+        chunk: int = CHUNK,
+    ) -> list[tuple]:
+        """The part entries of a flat shard's changes against ``previous`` (its wire bits), on the current
+        stream, in chunks of ``chunk`` elements (fp8: whole 128-row blocks, and their block scales). Writes
+        the new bits to ``out`` unless None."""
+        device = torch.device(f"cuda:{self.device}")
+        wire = self.wire_dtypes[key]
+        segment_names, begins, ends, shifts = self._segments(key, self.segments[fqn], offset, shard.numel())
+        sums = torch.zeros(len(segment_names), dtype=torch.int64, device=device)
+        changed, changed_values, scales = [], [], []
+        if fqn in self.fp8:
+            ue8m0, num_cols = self.fp8[fqn]
+            if gathered is not None:
+                chunks = [(0, shard.numel())]
+            else:
+                rows = max(FP8_BLOCK[0], chunk // num_cols // FP8_BLOCK[0] * FP8_BLOCK[0])
+                chunks = [(start, start + rows * num_cols) for start in range(0, shard.numel(), rows * num_cols)]
+        else:
+            chunks = [(start, start + chunk) for start in range(0, shard.numel(), chunk)]
+        for start, end in chunks:
+            end = min(end, shard.numel())
+            current = as_bits(shard[start:end].to(device, non_blocking=True).to(wire))
+            previous_bits = previous[start:end].to(device, non_blocking=True)
             if fqn in self.fp8:
-                ue8m0, num_cols = self.fp8[fqn]
                 if gathered is not None:
-                    chunks = [(0, shard.numel())]
+                    previous_weights, current_weights = gathered
                 else:
-                    rows = max(FP8_BLOCK[0], self.CHUNK // num_cols // FP8_BLOCK[0] * FP8_BLOCK[0])
-                    chunks = [(start, start + rows * num_cols) for start in range(0, shard.numel(), rows * num_cols)]
-            else:
-                chunks = [(start, start + self.CHUNK) for start in range(0, shard.numel(), self.CHUNK)]
-            for start, end in chunks:
-                end = min(end, shard.numel())
-                current = as_bits(shard[start:end].to(device, non_blocking=True).to(wire))
-                previous_bits = previous[start:end].to(device, non_blocking=True)
-                if fqn in self.fp8:
-                    if gathered is not None:
-                        previous_weights, current_weights = gathered
-                    else:
-                        previous_weights, current_weights = (
-                            previous_bits.view(torch.bfloat16),
-                            current.view(torch.bfloat16),
-                        )
-                    (previous_fp8, previous_scales), (current_fp8, current_scales) = (
-                        quantize_fp8_blocks(weights.reshape(-1, num_cols), ue8m0, self.compiled_fp8)
-                        for weights in (previous_weights, current_weights)
+                    previous_weights, current_weights = (
+                        previous_bits.view(torch.bfloat16),
+                        current.view(torch.bfloat16),
                     )
-                    scales.append((previous_scales, current_scales))
-                    if gathered is not None:
-                        previous_fp8 = previous_fp8.view(-1)[offset : offset + shard.numel()]
-                        current_fp8 = current_fp8.view(-1)[offset : offset + shard.numel()]
-                    delta = current_fp8.view(-1) - previous_fp8.view(-1)
-                    new_bits = current_fp8.view(-1)
-                else:
-                    delta = current - previous_bits
-                    new_bits = current
+                (previous_fp8, previous_scales), (current_fp8, current_scales) = (
+                    quantize_fp8_blocks(weights.reshape(-1, num_cols), ue8m0, self.compiled_fp8)
+                    for weights in (previous_weights, current_weights)
+                )
+                scales.append((previous_scales, current_scales))
+                if gathered is not None:
+                    previous_fp8 = previous_fp8.view(-1)[offset : offset + shard.numel()]
+                    current_fp8 = current_fp8.view(-1)[offset : offset + shard.numel()]
+                delta = current_fp8.view(-1) - previous_fp8.view(-1)
+                new_bits = current_fp8.view(-1)
+            else:
+                delta = current - previous_bits
+                new_bits = current
+            if out is not None:
                 out[start:end].copy_(current, non_blocking=True)
-                index = delta.nonzero().squeeze(1)
-                changed.append(index + start)
-                changed_values.append(zigzag(delta[index]))
-                first, last = bisect_right(ends, start), bisect_left(begins, end)
-                if first < last:
-                    sums[first:last] += self._segment_sums(
-                        new_bits,
-                        [max(begin, start) - start for begin in begins[first:last]],
-                        [min(stop, end) - start for stop in ends[first:last]],
-                        [max(begin, start) + shift for begin, shift in zip(begins[first:last], shifts[first:last])],
-                    )
-            pending.append(
-                self._part(
-                    fqn,
-                    self.segments[fqn],
-                    torch.cat(changed) + offset,
-                    torch.cat(changed_values),
-                    (segment_names, sums),
+            index = delta.nonzero().squeeze(1)
+            changed.append(index + start)
+            changed_values.append(zigzag(delta[index]))
+            first, last = bisect_right(ends, start), bisect_left(begins, end)
+            if first < last:
+                sums[first:last] += self._segment_sums(
+                    new_bits,
+                    [max(begin, start) - start for begin in begins[first:last]],
+                    [min(stop, end) - start for stop in ends[first:last]],
+                    [max(begin, start) + shift for begin, shift in zip(begins[first:last], shifts[first:last])],
                 )
+        parts = [
+            self._part(
+                fqn,
+                self.segments[fqn],
+                torch.cat(changed) + offset,
+                torch.cat(changed_values),
+                (segment_names, sums),
+            )
+        ]
+        if scales:
+            # This rank sends the scales of the blocks that start in its rows
+            first_row = offset // num_cols
+            first_block, end_block = (
+                -(-row // FP8_BLOCK[0]) for row in (first_row, first_row + len(shard) // num_cols)
             )
-            num_pending += len(pending[-1][2])
-            if scales:
-                # This rank sends the scales of the blocks that start in its rows
-                first_row = offset // num_cols
-                first_block, end_block = (
-                    -(-row // FP8_BLOCK[0]) for row in (first_row, first_row + len(shard) // num_cols)
+            previous_scales, current_scales = (torch.cat(side) for side in zip(*scales))
+            if gathered is None:
+                first_block = first_row // FP8_BLOCK[0]
+            else:
+                previous_scales, current_scales = (s[first_block:end_block] for s in (previous_scales, current_scales))
+            if current_scales.numel():
+                parts.append(self._scale_part(fqn, first_block, previous_scales, current_scales))
+        return parts
+
+    def _checksum(self, bits: torch.Tensor) -> torch.Tensor:
+        """``bits_checksum`` of a flat shard, in chunks of ``HOOK_CHUNK`` elements."""
+        n = bits.numel()
+        sums = [
+            self._segment_sums(bits, [start], [min(start + self.HOOK_CHUNK, n)], [start])
+            for start in range(0, n, self.HOOK_CHUNK)
+        ]
+        return torch.cat(sums).sum(0, keepdim=True) if sums else torch.zeros(1, dtype=torch.int64, device=bits.device)
+
+    def bind_optimizer(self, model: nn.Module, is_master: bool, params: list[torch.Tensor]) -> None:
+        """Diff the trainer tensors stored in ``params`` (contiguous local shards of an offloaded optimizer's
+        parameters) in the optimizer pass (``optimizer_hook``), not against the snapshot, which only keeps the
+        other tensors and the fp8 tensors gathered across ranks. A trainer tensor may be a strided view of a
+        parameter (e.g. one projection of a packed gate/up weight)."""
+        assert all(param.is_contiguous() for param in params)
+        stored = {param.untyped_storage().data_ptr() for param in params}
+        residual = []
+        for shard in self.local_shards(model, is_master):
+            key, fqn, tensor, offset, _ = shard
+            storage = tensor.untyped_storage().data_ptr()
+            if key not in self.gather and storage in stored:
+                view = (tuple(tensor.shape), tensor.stride(), tensor.storage_offset())
+                self.param_keys[storage].append((key, fqn, tensor.numel(), offset, view))
+                self.optimizer_keys.add(key)
+            else:
+                residual.append(shard)
+        self.track(residual)
+        # GPU slots for a parameter's previous and new values, diffed on the diff thread
+        slot_bytes = 2 * -(-max(param.nbytes for param in params) // 256) * 256
+        self.hook_slots = queue.SimpleQueue()
+        for _ in range(self.HOOK_SLOTS):
+            self.hook_slots.put(torch.empty(slot_bytes, dtype=torch.uint8, device=f"cuda:{self.device}"))
+        self.differ = ThreadPoolExecutor(1)
+        self.diff_stream = torch.cuda.Stream(self.device)
+        num_bytes = sum(numel * self.wire_dtypes[key].itemsize for key, (numel, _) in self.shards.items())
+        get_logger().info(
+            f"Sparse weight deltas of {len(self.optimizer_keys)} tensors come from the optimizer step, "
+            f"{len(self.shards)} tensors ({num_bytes / 2**30:.2f} GiB) are diffed against a snapshot"
+        )
+        get_logger().debug(f"Diffed against the snapshot: {sorted(self.shards)}")
+
+    def arm(self, anchor: AnchorWriter | None) -> None:
+        """Collect the next optimizer step's delta (and anchor slices, if ``anchor``)."""
+        self.step_parts, self.expected, self.scale_changes, self.anchor = [], {}, 0, anchor
+
+    def anchor_slices(self) -> list[tuple[str, list[int], int, int, torch.dtype]]:
+        """Every ``(hf_name, shape, flat start, numel, wire dtype)`` slice of this rank's share of an anchor."""
+        shards = {key: (numel, offset) for entries in self.param_keys.values() for key, _, numel, offset, _ in entries}
+        shards.update(self.shards)
+        return [
+            (name, shape, start, end - begin, self.wire_dtypes[key])
+            for key, (numel, offset) in shards.items()
+            for name, shape, start, begin, end in self._slices(key, numel, offset)
+        ]
+
+    @torch.no_grad()
+    def optimizer_hook(self, local_param: torch.Tensor, source: torch.Tensor) -> None:
+        """``FullCPUOffloadOptimizer.weight_hook``, on its H2D stream: copy the shard, which still holds the
+        previous broadcast's weights, and its new values ``source`` (host) into a GPU slot, write the new values,
+        and diff the slot on the diff thread. At anchor steps, writes the new values from the host first."""
+        start = local_param.storage_offset()
+        entries = [
+            (key, fqn, numel, offset, (size, stride, at - start))
+            for key, fqn, numel, offset, (size, stride, at) in self.param_keys.get(
+                local_param.untyped_storage().data_ptr(), []
+            )
+            if start <= at < start + local_param.numel()
+        ]
+        if not entries:
+            local_param.copy_(source, non_blocking=True)
+            return
+        slot = self.hook_slots.get()
+        half = slot.numel() // 2
+        previous = slot[: source.nbytes].view(source.dtype)
+        new = slot[half : half + source.nbytes].view(source.dtype)
+        previous.copy_(local_param.reshape(-1))
+        new.copy_(source.reshape(-1), non_blocking=True)
+        local_param.copy_(new.view(local_param.shape))
+        ready = torch.cuda.Event()
+        ready.record()
+        if self.anchor is not None:
+            for key, _, numel, offset, (size, stride, at) in entries:
+                host = source.as_strided(size, stride, source.storage_offset() + at).reshape(-1)
+                host = host.to(self.wire_dtypes[key])
+                self.anchor.write_all(
+                    [
+                        (name, hf_start, host[begin:end].view(torch.uint8).numpy())
+                        for name, _, hf_start, begin, end in self._slices(key, numel, offset)
+                    ]
                 )
-                previous_scales, current_scales = (torch.cat(side) for side in zip(*scales))
-                if gathered is None:
-                    first_block = first_row // FP8_BLOCK[0]
-                else:
-                    previous_scales, current_scales = (
-                        s[first_block:end_block] for s in (previous_scales, current_scales)
+        self.diffs.append(self.differ.submit(self._diff_slot, entries, slot, previous, new, ready))
+
+    @torch.no_grad()
+    def _diff_slot(self, entries: list[tuple], slot: torch.Tensor, previous, new, ready: torch.cuda.Event) -> None:
+        """Diff thread: the parts of one parameter's tensors from its slot, then free the slot."""
+        torch.cuda.set_device(self.device)
+        try:
+            with torch.cuda.stream(self.diff_stream):
+                self.diff_stream.wait_event(ready)
+                for key, fqn, _, offset, (size, stride, at) in entries:
+                    wire = self.wire_dtypes[key]
+                    current, old = (
+                        values.as_strided(size, stride, values.storage_offset() + at).reshape(-1).to(wire)
+                        for values in (new, previous)
                     )
-                if current_scales.numel():
-                    pending.append(self._scale_part(fqn, first_block, previous_scales, current_scales))
-            if num_pending >= self.PART:
-                flush()
-                num_pending = 0
-        if pending:
-            flush()
+                    parts = self._diff_shard(key, fqn, current, offset, None, as_bits(old), None, self.HOOK_CHUNK)
+                    self.step_parts += parts
+                    self.expected[key] = self._checksum(as_bits(current))
+                self.diff_stream.synchronize()
+        finally:
+            self.hook_slots.put(slot)
 
 
 def _model_dir(model: nn.Module) -> Path:
@@ -611,8 +788,22 @@ def __init__(self, output_dir: Path, config: FileSystemWeightBroadcastConfig):
         # The checkpoint the engines hold before the first delta: the model's (None), or a resumed run's first
         # broadcast
         self.boot_dir: Path | None = None
+        # With an offloaded optimizer (``diff_in_optimizer``), the step its next pass collects the delta of
+        self.optimizer: "FullCPUOffloadOptimizer | None" = None
+        self.model: nn.Module | None = None
+        self.armed_step: int | None = None
+
+    def diff_in_optimizer(self, optimizer: "FullCPUOffloadOptimizer") -> None:
+        """Take the deltas of the optimizer's parameters from its pass (``SparseDeltaEncoder.optimizer_hook``),
+        which needs a broadcast after every optimizer step."""
+        self.optimizer = optimizer
 
     def wait_for_weight_reads(self) -> None:
+        self._wait_for_encode()
+        if self.optimizer is not None and self.last_step is not None and self.armed_step is None:
+            self._arm(self.last_step + 1)
+
+    def _wait_for_encode(self) -> None:
         if self.weights_read is not None:
             start = time.perf_counter()
             self.weights_read.wait()
@@ -620,6 +811,31 @@ def wait_for_weight_reads(self) -> None:
             self.logger.debug(f"Waited {time.perf_counter() - start:.2f}s for the sparse weight encode")
         self._raise_errors()
 
+    def _arm(self, step: int) -> None:
+        """Collect v{step}'s delta, and its anchor slices at anchor steps, in the coming optimizer pass."""
+        encoder = self.encoder
+        assert encoder is not None and self.optimizer is not None
+        self._set_fp8_layout()
+        if not encoder.optimizer_keys:
+            params = [master.model_param.to_local() for master in self.optimizer._master_weights.values()]
+            encoder.bind_optimizer(self.model, self.world.is_master, params)
+            encoder.allocate_host()
+        anchor = None
+        if step % self.config.anchor_interval == 0:
+            # Written next to the step dir, which the broadcast handshake recreates, and moved in once complete
+            anchor = AnchorWriter(encoder.anchor_slices(), self.step_dir(step), self.world.rank)
+        encoder.arm(anchor)
+        self.optimizer.weight_hook = encoder.optimizer_hook
+        self.armed_step = step
+
+    def _set_fp8_layout(self) -> None:
+        if self.encoder.fp8_layout is None:
+            # Every engine worker reports which tensors it holds in fp8 when it loads the first broadcast.
+            broadcast_dir = self.step_dir(0).parent
+            layout = [self._wait_for_engine_layout(broadcast_dir) if self.world.is_master else None]
+            dist.broadcast_object_list(layout, src=0)
+            self.encoder.set_fp8_layout(self.model, layout[0], self.world.is_master)
+
     def close(self) -> None:
         for pool in (self.encodes, self.writes, self.anchors):
             pool.shutdown(wait=True)
@@ -641,17 +857,24 @@ def run() -> None:
         return pool.submit(run)
 
     def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
-        self.wait_for_weight_reads()
+        self._wait_for_encode()
         if self.world.is_master:
             (step_dir / DELTA_DIR).mkdir(exist_ok=True)
         if self.encoder is None:
             self.encoder = SparseDeltaEncoder(model, self.config.snapshot)
+            self.model = model
         encoder = self.encoder
         base_step, self.last_step = self.last_step, step
+        if self.optimizer is not None and base_step is not None:
+            if self.armed_step != step:
+                raise RuntimeError(f"Sparse weight delta v{step} was not collected in the optimizer step")
+            self.optimizer.weight_hook = None
+            self.armed_step = None
         if base_step is None:
             shards = encoder.local_shards(model, self.world.is_master)
             encoder.track(shards)
-            encoder.allocate_host()
+            if self.optimizer is None:
+                encoder.allocate_host()
         if base_step is None and step == 0:
             # The run's first broadcast is the checkpoint the engines booted from.
             self.anchor_step = step
@@ -668,11 +891,7 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
             self._write_full(step, step_dir, encoder.anchor_parts(source))
             self.boot_dir = step_dir
             return
-        if encoder.fp8_layout is None:
-            # Every engine worker reports which tensors it holds in fp8 when it loads the first broadcast.
-            layout = [self._wait_for_engine_layout(step_dir.parent) if self.world.is_master else None]
-            dist.broadcast_object_list(layout, src=0)
-            encoder.set_fp8_layout(model, layout[0], self.world.is_master)
+        self._set_fp8_layout()
         if not encoder.initialized:
             encoder.initialize(
                 model, self.world.is_master, torch.cuda.max_memory_reserved(), self.boot_dir or _model_dir(model)
@@ -680,7 +899,7 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
         shards = encoder.local_shards(model, self.world.is_master)
         # Periodic anchors only serve engines that fall off the chain: written in the background, each after
         # the previous one.
-        anchor = step % self.config.anchor_interval == 0
+        anchor = encoder.anchor if self.optimizer is not None else step % self.config.anchor_interval == 0
         ready = torch.cuda.Event()
         ready.record()
         self.weights_read = threading.Event()
@@ -738,9 +957,10 @@ def _encode(
         shards,
         ready,
         weights_read: threading.Event,
-        anchor: bool,
+        anchor: bool | AnchorWriter | None,
     ) -> None:
-        """Encode thread: diff v{step}, handing the delta to the writer in parts as it is encoded."""
+        """Encode thread: diff v{step}, handing the delta to the writer in parts as it is encoded. ``anchor``: write
+        a periodic anchor, or finish the one the optimizer pass wrote."""
         encoder = self.encoder
         assert encoder is not None
         self.encode_queued_at = None
@@ -755,16 +975,24 @@ def emit(part: tuple) -> None:
             self._submit(self.writes, self._write_part, step_dir, num_parts, part)
             num_parts += 1
 
-        encoder.encode(shards, ready, weights_read, emit)
+        scale_changes = encoder.encode(shards, ready, weights_read, emit)
         # Free the gathered fp8 tensors now, not after the anchor copy below, which may overlap the next step's
         # memory peak.
         shards.clear()
         if num_parts == 0 and self.world.is_master:
             self._submit(self.writes, self._start_delta, step, base_step, step_dir)
         encode_time = time.perf_counter() - start
-        args = (step, base_step, step_dir, num_parts, encode_time, encoder.scale_changes)
+        args = (step, base_step, step_dir, num_parts, start, encode_time, scale_changes)
         self._submit(self.writes, self._finish_delta, *args)
-        if anchor:
+        if isinstance(anchor, AnchorWriter):
+            # The snapshot of the other tensors holds v{step} until the next encode, on this thread
+            for name, _, start, values in encoder.anchor_parts(encoder.snapshot):
+                written = 0
+                for chunk in encoder.stream_to_host(values):
+                    anchor.write(name, start, chunk, written)
+                    written += chunk.nbytes
+            self._submit(self.anchors, self._close_anchor, step, step_dir, anchor)
+        elif anchor:
             self._write_periodic_anchor(step, step_dir)
 
     def _start_delta(self, step: int, base_step: int, step_dir: Path) -> None:
@@ -787,7 +1015,14 @@ def _write_part(self, step_dir: Path, index: int, part: tuple) -> None:
         )
 
     def _finish_delta(
-        self, step: int, base_step: int, step_dir: Path, num_parts: int, encode_time: float, scale_changes: int
+        self,
+        step: int,
+        base_step: int,
+        step_dir: Path,
+        num_parts: int,
+        start: float,
+        encode_time: float,
+        scale_changes: int,
     ) -> None:
         """After this rank's last part: its parts marker. The master then waits for every rank's marker and
         commits the step. Runs on the writer thread, without collectives."""
@@ -813,7 +1048,8 @@ def _finish_delta(
         )
         self.logger.info(
             f"Sparse weight delta v{step} (base v{base_step}): {num_bytes / 2**20:.1f} MiB, "
-            f"encode {encode_time:.2f}s (parts written as encoded){scales}"
+            f"encode {encode_time:.2f}s (parts written as encoded), committed {time.perf_counter() - start:.2f}s after "
+            f"the encode started{scales}"
         )
 
     def _write_full(self, step: int, step_dir: Path, parts: list) -> None:
@@ -826,6 +1062,13 @@ def _write_anchor(self, step: int, step_dir: Path, parts: list) -> None:
         start = time.perf_counter()
         write_anchor_shards(parts, step_dir, self.world.rank, self.world.world_size, self.encoder.stream_to_host)
         self.logger.debug(f"Wrote the {step_dir.name} anchor in {time.perf_counter() - start:.1f}s")
+        self._wait_for_anchor(step, step_dir)
+
+    def _close_anchor(self, step: int, step_dir: Path, anchor: AnchorWriter) -> None:
+        anchor.close(self.world.world_size)
+        self._wait_for_anchor(step, step_dir)
+
+    def _wait_for_anchor(self, step: int, step_dir: Path) -> None:
         if self.world.is_master:
             while not anchor_complete(step_dir):
                 if any(step_dir.glob("anchor_rank*.skipped")):
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index ad37c5327f1..f06b6ccc025 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -12,7 +12,6 @@
 checksums of the new weights catch it if not.
 """
 
-import itertools
 import json
 import os
 import threading
@@ -22,7 +21,7 @@
 from contextlib import ExitStack
 from itertools import groupby
 from pathlib import Path
-from typing import BinaryIO
+from typing import Any, BinaryIO
 
 import numpy as np
 import pyzstd
@@ -261,6 +260,70 @@ def write_file(path: Path, data: bytes | Iterable) -> None:
     tmp.rename(path)
 
 
+class AnchorWriter:
+    """One trainer rank's share of a full checkpoint: ``(hf_name, shape, flat start, numel, dtype)`` slices,
+    stored as tensors ``"{hf_name}@{start}"`` in ~1 GiB safetensors files whose headers are written first, so
+    slices can be written in any order, from any thread (``write_all``: in parallel). The files are written next to ``step_dir`` and moved
+    into it by ``close``, followed by the ``anchor_rank{rank}.done`` marker (the world size)."""
+
+    def __init__(self, slices: list[tuple[str, list[int], int, int, torch.dtype]], step_dir: Path, rank: int):
+        self.step_dir, self.rank = step_dir, rank
+        self.tmp_dir = step_dir.with_name(f".{step_dir.name}.anchor{rank}")
+        self.tmp_dir.mkdir(parents=True, exist_ok=True)
+        self.files: list[tuple[str, int]] = []
+        self.offsets: dict[tuple[str, int], tuple[int, int]] = {}
+        # Slices go round-robin to the files: a file is written at ~0.5 GiB/s, several at once in parallel
+        num_files = max(1, -(-sum(numel * dtype.itemsize for *_, numel, dtype in slices) // 2**30))
+        groups = [slices[index::num_files] for index in range(num_files)]
+        self.pool = ThreadPoolExecutor(8)
+        for index, group in enumerate(groups):
+            header, offset = {"__metadata__": {"shapes": json.dumps({name: shape for name, shape, *_ in group})}}, 0
+            for name, _, start, numel, dtype in group:
+                safetensors_dtype = next(
+                    key for key, value in SAFETENSORS_DTYPES.items() if getattr(torch, value) == dtype
+                )
+                header[f"{name}@{start}"] = {
+                    "dtype": safetensors_dtype,
+                    "shape": [numel],
+                    "data_offsets": [offset, offset + numel * dtype.itemsize],
+                }
+                offset += numel * dtype.itemsize
+            encoded = json.dumps(header).encode()
+            encoded += b" " * (-len(encoded) % 8)
+            file_name = f"model-rank{rank:05d}-{index:05d}.safetensors"
+            fd = os.open(self.tmp_dir / file_name, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o644)
+            self._pwrite(fd, len(encoded).to_bytes(8, "little") + encoded, 0)
+            for name, _, start, _, _ in group:
+                self.offsets[(name, start)] = (fd, 8 + len(encoded) + header[f"{name}@{start}"]["data_offsets"][0])
+            self.files.append((file_name, fd))
+
+    @staticmethod
+    def _pwrite(fd: int, data, offset: int) -> None:
+        data = memoryview(data).cast("B")
+        while data:
+            written = os.pwrite(fd, data, offset)
+            data, offset = data[written:], offset + written
+
+    def write(self, name: str, start: int, data, at: int = 0) -> None:
+        """Write ``data`` (bytes) at byte ``at`` of the slice ``"{name}@{start}"``."""
+        fd, offset = self.offsets[(name, start)]
+        self._pwrite(fd, data, offset + at)
+
+    def write_all(self, items: list[tuple[str, int, Any]]) -> None:
+        """``write`` each ``(name, start, data)``, in parallel."""
+        for future in [self.pool.submit(self.write, *item) for item in items]:
+            future.result()
+
+    def close(self, world_size: int) -> None:
+        self.pool.shutdown()
+        for name, fd in self.files:
+            drop_page_cache(fd)
+            os.close(fd)
+            (self.tmp_dir / name).rename(self.step_dir / name)
+        self.tmp_dir.rmdir()
+        write_file(self.step_dir / f"anchor_rank{self.rank}.done", str(world_size).encode())
+
+
 def write_anchor_shards(
     parts: list[tuple[str, list[int], int, torch.Tensor]],
     step_dir: Path,
@@ -268,34 +331,18 @@ def write_anchor_shards(
     world_size: int,
     to_host: Callable[[torch.Tensor], Iterable] = lambda values: [values.view(torch.uint8).numpy()],
 ) -> None:
-    """Write one trainer rank's share of a full checkpoint: ``(hf_name, shape, flat start, flat values)``
-    slices, stored as tensors ``"{hf_name}@{start}"`` in ~1 GiB safetensors files, without collectives.
-    ``to_host`` yields each slice's bytes on the host, each written before the next is taken. The
-    ``anchor_rank{rank}.done`` marker (the world size) comes last."""
-    files, size = [[]], 0
-    for part in parts:
-        if size > 2**30:
-            files.append([])
-            size = 0
-        files[-1].append(part)
-        size += part[3].nbytes
-    for index, file_parts in enumerate(files):
-        header, offset = {"__metadata__": {"shapes": json.dumps({name: shape for name, shape, _, _ in file_parts})}}, 0
-        for name, _, start, values in file_parts:
-            dtype = next(key for key, dtype in SAFETENSORS_DTYPES.items() if getattr(torch, dtype) == values.dtype)
-            header[f"{name}@{start}"] = {
-                "dtype": dtype,
-                "shape": [values.numel()],
-                "data_offsets": [offset, offset + values.nbytes],
-            }
-            offset += values.nbytes
-        encoded = json.dumps(header).encode()
-        encoded += b" " * (-len(encoded) % 8)
-        buffers = itertools.chain(
-            [len(encoded).to_bytes(8, "little") + encoded], (chunk for *_, v in file_parts for chunk in to_host(v))
-        )
-        write_file(step_dir / f"model-rank{rank:05d}-{index:05d}.safetensors", buffers)
-    write_file(step_dir / f"anchor_rank{rank}.done", str(world_size).encode())
+    """Write one trainer rank's share of a full checkpoint (``AnchorWriter``) from ``(hf_name, shape, flat
+    start, flat values)`` slices. ``to_host`` yields each slice's bytes on the host, each written before the
+    next is taken."""
+    writer = AnchorWriter(
+        [(name, shape, start, values.numel(), values.dtype) for name, shape, start, values in parts], step_dir, rank
+    )
+    for name, _, start, values in parts:
+        written = 0
+        for chunk in to_host(values):
+            writer.write(name, start, chunk, written)
+            written += chunk.nbytes
+    writer.close(world_size)
 
 
 def checkpoint_index(directory: Path) -> dict[str, tuple[list[int], torch.dtype, list[tuple[int, int, Path, int]]]]:

From df71b4e41073a336db9456d976f28ed8aed8073b Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sat, 17 Oct 2026 16:34:54 +0000
Subject: [PATCH 49/88] perf(weights): write the optimizer step delta parts
 during the step

---
 src/prime_rl/transports/weights/sparse.py | 107 ++++++++++++++++------
 1 file changed, 79 insertions(+), 28 deletions(-)

diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 2dd2ff0e5ce..bcd46556c3e 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -92,6 +92,8 @@ class SparseDeltaEncoder:
     PART = 1 << 26
     # Elements per transient chunk of the diff in the optimizer pass, which overlaps the final backward
     HOOK_CHUNK = 1 << 24
+    # Changes per part handed to the writer during the optimizer step
+    HOOK_PART = 1 << 24
     # GPU slots of the optimizer pass's diff thread, each two of the largest parameter shard
     HOOK_SLOTS = 2
 
@@ -418,14 +420,17 @@ def stream_to_host(self, tensor: torch.Tensor):
                 self.write_stream.synchronize()
                 yield chunk.numpy()
 
-    def _to_host(self, tensor: torch.Tensor, host: torch.Tensor | None = None) -> torch.Tensor:
+    def _to_host(
+        self, tensor: torch.Tensor, host: torch.Tensor | None = None, staging: torch.Tensor | None = None
+    ) -> torch.Tensor:
         """Copy a GPU tensor to pageable host memory (``host``, or a new tensor) through the pinned staging
         buffer: a copy straight into pageable memory holds the CUDA driver and stalls the training step's
         own CUDA calls."""
         host = torch.empty(tensor.shape, dtype=tensor.dtype) if host is None else host
+        staging = self.staging if staging is None else staging
         data, host_bytes = tensor.view(torch.uint8), host.view(torch.uint8)
-        for begin in range(0, data.numel(), self.staging.numel()):
-            chunk = self.staging[: min(self.staging.numel(), data.numel() - begin)]
+        for begin in range(0, data.numel(), staging.numel()):
+            chunk = staging[: min(staging.numel(), data.numel() - begin)]
             chunk.copy_(data[begin : begin + chunk.numel()], non_blocking=True)
             torch.cuda.current_stream().synchronize()
             host_bytes[begin : begin + chunk.numel()].copy_(chunk)
@@ -499,21 +504,12 @@ def _diff(self, shards: list[tuple], emit: Callable) -> None:
         """Diff ``shards`` against the snapshot on the current stream, writing the new bits to the snapshot.
         Shards diffed in the optimizer pass are only checked against the values that pass wrote; their parts
         go first. Copies the delta to the host and hands it to ``emit`` every ``PART`` changes."""
-        for diff in self.diffs:
-            diff.result()
-        self.diffs.clear()
         pending, self.step_parts = self.step_parts, []
         num_pending = sum(len(part[2]) for part in pending)
         checks: list[tuple[str, torch.Tensor, torch.Tensor]] = []
 
         def flush() -> None:
-            names = [name for part in pending for name in part[0]]
-            counts = torch.cat([part[1] for part in pending]).tolist()
-            gaps = self._to_host(torch.cat([part[2] for part in pending]))
-            values = self._to_host(torch.cat([part[3] for part in pending]))
-            checksum_names = [name for part in pending for name in part[4]]
-            checksums = dict(zip(checksum_names, torch.cat([part[5] for part in pending]).tolist()))
-            emit(([(name, count) for name, count in zip(names, counts) if count], gaps, values, checksums))
+            emit(self._host_part(pending))
             pending.clear()
 
         for key, fqn, shard, offset, gathered in shards:
@@ -540,6 +536,27 @@ def flush() -> None:
             if changed:
                 raise RuntimeError(f"Trainer weights changed outside the optimizer step, not sent: {changed}")
 
+    def wait_for_diffs(self) -> None:
+        """Wait for the diff thread to finish the optimizer step's delta."""
+        for diff in self.diffs:
+            diff.result()
+        if self.diffs:
+            get_logger().debug(
+                f"Sparse weight delta in the optimizer step: waited {self.slot_wait:.2f}s for GPU slots, "
+                f"diff thread busy {self.diff_time:.2f}s"
+            )
+        self.diffs.clear()
+
+    def _host_part(self, pending: list[tuple], staging: torch.Tensor | None = None) -> tuple:
+        """One part ``(entries, gaps, value bytes, checksums)`` on the host from GPU part entries."""
+        names = [name for part in pending for name in part[0]]
+        counts = torch.cat([part[1] for part in pending]).tolist()
+        gaps = self._to_host(torch.cat([part[2] for part in pending]), staging=staging)
+        values = self._to_host(torch.cat([part[3] for part in pending]), staging=staging)
+        checksum_names = [name for part in pending for name in part[4]]
+        checksums = dict(zip(checksum_names, torch.cat([part[5] for part in pending]).tolist()))
+        return [(name, count) for name, count in zip(names, counts) if count], gaps, values, checksums
+
     def _diff_shard(
         self,
         key: str,
@@ -664,6 +681,7 @@ def bind_optimizer(self, model: nn.Module, is_master: bool, params: list[torch.T
             self.hook_slots.put(torch.empty(slot_bytes, dtype=torch.uint8, device=f"cuda:{self.device}"))
         self.differ = ThreadPoolExecutor(1)
         self.diff_stream = torch.cuda.Stream(self.device)
+        self.diff_staging = torch.empty(1 << 28, dtype=torch.uint8, pin_memory=True)
         num_bytes = sum(numel * self.wire_dtypes[key].itemsize for key, (numel, _) in self.shards.items())
         get_logger().info(
             f"Sparse weight deltas of {len(self.optimizer_keys)} tensors come from the optimizer step, "
@@ -671,9 +689,11 @@ def bind_optimizer(self, model: nn.Module, is_master: bool, params: list[torch.T
         )
         get_logger().debug(f"Diffed against the snapshot: {sorted(self.shards)}")
 
-    def arm(self, anchor: AnchorWriter | None) -> None:
-        """Collect the next optimizer step's delta (and anchor slices, if ``anchor``)."""
-        self.step_parts, self.expected, self.scale_changes, self.anchor = [], {}, 0, anchor
+    def arm(self, anchor: AnchorWriter | None, emit: Callable) -> None:
+        """Collect the next optimizer step's delta, handing it to ``emit`` in parts on the host as the step goes
+        (and write its anchor slices, if ``anchor``)."""
+        self.step_parts, self.expected, self.scale_changes, self.anchor, self.emit = [], {}, 0, anchor, emit
+        self.slot_wait = self.diff_time = 0.0
 
     def anchor_slices(self) -> list[tuple[str, list[int], int, int, torch.dtype]]:
         """Every ``(hf_name, shape, flat start, numel, wire dtype)`` slice of this rank's share of an anchor."""
@@ -701,7 +721,9 @@ def optimizer_hook(self, local_param: torch.Tensor, source: torch.Tensor) -> Non
         if not entries:
             local_param.copy_(source, non_blocking=True)
             return
+        start_wait = time.perf_counter()
         slot = self.hook_slots.get()
+        self.slot_wait += time.perf_counter() - start_wait
         half = slot.numel() // 2
         previous = slot[: source.nbytes].view(source.dtype)
         new = slot[half : half + source.nbytes].view(source.dtype)
@@ -724,7 +746,9 @@ def optimizer_hook(self, local_param: torch.Tensor, source: torch.Tensor) -> Non
 
     @torch.no_grad()
     def _diff_slot(self, entries: list[tuple], slot: torch.Tensor, previous, new, ready: torch.cuda.Event) -> None:
-        """Diff thread: the parts of one parameter's tensors from its slot, then free the slot."""
+        """Diff thread: the parts of one parameter's tensors from its slot, then free the slot. Every ``HOOK_PART``
+        changes, hands the parts so far to the host."""
+        start = time.perf_counter()
         torch.cuda.set_device(self.device)
         try:
             with torch.cuda.stream(self.diff_stream):
@@ -735,12 +759,18 @@ def _diff_slot(self, entries: list[tuple], slot: torch.Tensor, previous, new, re
                         values.as_strided(size, stride, values.storage_offset() + at).reshape(-1).to(wire)
                         for values in (new, previous)
                     )
-                    parts = self._diff_shard(key, fqn, current, offset, None, as_bits(old), None, self.HOOK_CHUNK)
-                    self.step_parts += parts
+                    self.step_parts += self._diff_shard(
+                        key, fqn, current, offset, None, as_bits(old), None, self.HOOK_CHUNK
+                    )
                     self.expected[key] = self._checksum(as_bits(current))
                 self.diff_stream.synchronize()
         finally:
             self.hook_slots.put(slot)
+        if sum(len(part[2]) for part in self.step_parts) >= self.HOOK_PART:
+            with torch.cuda.stream(self.diff_stream):
+                self.emit(self._host_part(self.step_parts, self.diff_staging))
+            self.step_parts = []
+        self.diff_time += time.perf_counter() - start
 
 
 def _model_dir(model: nn.Module) -> Path:
@@ -792,6 +822,8 @@ def __init__(self, output_dir: Path, config: FileSystemWeightBroadcastConfig):
         self.optimizer: "FullCPUOffloadOptimizer | None" = None
         self.model: nn.Module | None = None
         self.armed_step: int | None = None
+        self.streamed_dir: Path | None = None
+        self.num_streamed = 0
 
     def diff_in_optimizer(self, optimizer: "FullCPUOffloadOptimizer") -> None:
         """Take the deltas of the optimizer's parameters from its pass (``SparseDeltaEncoder.optimizer_hook``),
@@ -824,10 +856,20 @@ def _arm(self, step: int) -> None:
         if step % self.config.anchor_interval == 0:
             # Written next to the step dir, which the broadcast handshake recreates, and moved in once complete
             anchor = AnchorWriter(encoder.anchor_slices(), self.step_dir(step), self.world.rank)
-        encoder.arm(anchor)
+        # Parts written during the step go next to the step dir too, and are moved in when it is published
+        self.streamed_dir = self.step_dir(step).with_name(f".{self.step_dir(step).name}.delta{self.world.rank}")
+        self.streamed_dir.mkdir(parents=True, exist_ok=True)
+        self.num_streamed = 0
+        encoder.arm(anchor, self._stream_part)
         self.optimizer.weight_hook = encoder.optimizer_hook
         self.armed_step = step
 
+    def _stream_part(self, part: tuple) -> None:
+        """A part of the delta the optimizer step produces, written while the step runs."""
+        path = self.streamed_dir / rank_part(self.world.rank, self.num_streamed)
+        self._submit(self.writes, self._write_part, path, part)
+        self.num_streamed += 1
+
     def _set_fp8_layout(self) -> None:
         if self.encoder.fp8_layout is None:
             # Every engine worker reports which tensors it holds in fp8 when it loads the first broadcast.
@@ -966,21 +1008,24 @@ def _encode(
         self.encode_queued_at = None
         self._wait(self.snapshot_writes)
         start = time.perf_counter()
-        num_parts = 0
+        # The delta parts the optimizer step already wrote (a fully offloaded optimizer)
+        encoder.wait_for_diffs()
+        streamed_dir, num_parts = (self.streamed_dir, self.num_streamed) if self.optimizer is not None else (None, 0)
+        if self.world.is_master:
+            self._submit(self.writes, self._start_delta, step, base_step, step_dir)
+        if streamed_dir is not None:
+            self._submit(self.writes, self._move_streamed, streamed_dir, step_dir / DELTA_DIR)
 
         def emit(part: tuple) -> None:
             nonlocal num_parts
-            if num_parts == 0 and self.world.is_master:
-                self._submit(self.writes, self._start_delta, step, base_step, step_dir)
-            self._submit(self.writes, self._write_part, step_dir, num_parts, part)
+            path = step_dir / DELTA_DIR / rank_part(self.world.rank, num_parts)
+            self._submit(self.writes, self._write_part, path, part)
             num_parts += 1
 
         scale_changes = encoder.encode(shards, ready, weights_read, emit)
         # Free the gathered fp8 tensors now, not after the anchor copy below, which may overlap the next step's
         # memory peak.
         shards.clear()
-        if num_parts == 0 and self.world.is_master:
-            self._submit(self.writes, self._start_delta, step, base_step, step_dir)
         encode_time = time.perf_counter() - start
         args = (step, base_step, step_dir, num_parts, start, encode_time, scale_changes)
         self._submit(self.writes, self._finish_delta, *args)
@@ -1006,14 +1051,20 @@ def _start_delta(self, step: int, base_step: int, step_dir: Path) -> None:
         }
         _write_json(step_dir / DELTA_DIR / MANIFEST, manifest)
 
-    def _write_part(self, step_dir: Path, index: int, part: tuple) -> None:
+    def _write_part(self, path: Path, part: tuple) -> None:
         assert self.encoder is not None
         entries, gaps, values, checksums = part
-        path = step_dir / DELTA_DIR / rank_part(self.world.rank, index)
         self.part_bytes += write_rank_file(
             path, entries, gaps.numpy(), values.numpy(), self.encoder.hf_tensors, checksums
         )
 
+    @staticmethod
+    def _move_streamed(streamed_dir: Path, delta_dir: Path) -> None:
+        delta_dir.mkdir(parents=True, exist_ok=True)
+        for path in streamed_dir.iterdir():
+            path.rename(delta_dir / path.name)
+        streamed_dir.rmdir()
+
     def _finish_delta(
         self,
         step: int,

From 738b37e56350ca8009eccd64d92c8fff17d8520b Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sat, 17 Oct 2026 17:10:48 +0000
Subject: [PATCH 50/88] fix(weights): map a non-empty sparse host buffer on
 ranks without shards

---
 src/prime_rl/transports/weights/sparse.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index bcd46556c3e..90233a7c435 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -264,7 +264,8 @@ def allocate_host(self) -> None:
         start = time.perf_counter()
         sizes = {key: -(-numel * self.wire_dtypes[key].itemsize // 8) * 8 for key, (numel, _) in self.shards.items()}
         flags = mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS | mmap.MAP_POPULATE
-        self.host_buffer = torch.from_numpy(np.frombuffer(mmap.mmap(-1, sum(sizes.values()), flags), np.uint8))
+        # mmap needs a non-empty mapping: a rank may hold no shards
+        self.host_buffer = torch.from_numpy(np.frombuffer(mmap.mmap(-1, max(sum(sizes.values()), 8), flags), np.uint8))
         offset = 0
         for key, (numel, _) in self.shards.items():
             self.host[key] = self.host_buffer[offset : offset + numel * self.wire_dtypes[key].itemsize].view(

From cc5beb6918c7015c5d9b0252b98843a5ed5cf0c4 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sat, 17 Oct 2026 18:12:09 +0000
Subject: [PATCH 51/88] chore(weights): trim sparse broadcast comments, drop
 dead attributes and debug timing, dedupe small helpers

---
 .../inference/vllm/worker/filesystem.py       |  30 +-
 .../inference/vllm/worker/weight_transfer.py  |   6 +-
 src/prime_rl/trainer/optim/offload.py         |   4 +-
 src/prime_rl/transports/weights/sparse.py     | 279 +++++++-----------
 src/prime_rl/utils/sparse_weights.py          |  79 ++---
 5 files changed, 164 insertions(+), 234 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index e043ff7bb29..4ab74b37b46 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -80,9 +80,8 @@ def prefetch_weights_from_path(self, weight_path: str) -> None:
     def _read_delta(
         self, step_dir: Path
     ) -> tuple[list[tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]], dict[str, int]]:
-        """Read this worker's part of a delta into a pinned host buffer that is reused across steps
-        (pinning memory is slow on a loaded host), in chunks of up to 2^26 changes
-        of consecutive frames."""
+        """Read this worker's part of a delta into a pinned host buffer reused across steps, in chunks of up to 2^26
+        changes of consecutive frames."""
         global _reader_pool
         if _reader_pool is None:
             _reader_pool = ThreadPoolExecutor(16)
@@ -103,8 +102,8 @@ def _read_delta(
             layout.append((num_bytes, num_bytes + gaps_bytes))
             num_bytes += gaps_bytes + values_bytes
         if self.host_buffer is None or self.host_buffer.numel() < num_bytes:
-            # Twice the delta, so later deltas (within ~20% of each other) fit: pinning stalls for tens of
-            # seconds on a host that must reclaim page cache first, and this read is the update's critical path
+            # Twice the delta, so later deltas fit: pinning can stall for tens of seconds on a host that must
+            # reclaim page cache first, and this read is on the update's critical path
             if self.host_buffer is not None:
                 unpin(self.host_buffer)
             pin_start = time.perf_counter()
@@ -140,7 +139,7 @@ def _read_delta(
             copy.result()
         logger.info(
             f"Read sparse weight delta {step_dir.name}: {num_bytes / 2**30:.2f} GiB, "
-            f"read+decompress {read_time:.2f}s, pin {time.perf_counter() - start - read_time:.2f}s"
+            f"read+decompress {read_time:.2f}s, copy {time.perf_counter() - start - read_time:.2f}s"
         )
         return pinned, checksums
 
@@ -161,7 +160,7 @@ def update_weights_from_path(self, weight_path: str) -> None:
         else:
             base_step = read_manifest(step_dir)["base_step"]
             if base_step is None:
-                # The run's first broadcast: the weights the engine booted with.
+                # The run's first broadcast: the weights the engine booted with
                 if held_step is not None:
                     self._load_checkpoint(self._boot_dir())
             elif base_step == held_step:
@@ -192,17 +191,16 @@ def _apply_chunk(
             num_changes = sum(count for _, count in run)
             itemsize = bits.itemsize
             delta = values[num_bytes : num_bytes + num_changes * itemsize]
-            # A run after an odd number of narrower values is misaligned for its dtype view.
+            # A run after an odd number of narrower values is misaligned for its dtype view
             delta = delta.view(bits) if num_bytes % itemsize == 0 else delta.clone().view(bits)
             self.index_map.apply_delta(run, indices[offset : offset + num_changes], delta)
             offset += num_changes
             num_bytes += num_changes * itemsize
 
     def _build_index_map(self, step_dir: Path) -> None:
-        """At a sparse run's first update, while the engine still holds its boot weights: map checkpoint
-        tensors to the params, and tell the trainer which ones this engine holds in fp8 (and with which
-        scales), so it can send their deltas in the fp8 domain. Reads one boot checkpoint tensor per fp8
-        param."""
+        """At a sparse run's first update, while the engine still holds its boot weights: map checkpoint tensors to
+        the params, and report which ones this worker holds in fp8 (and their scale format), so the trainer sends
+        their deltas as fp8."""
         if self.index_map is not None:
             return
         if (step_dir / DELTA_DIR / MANIFEST).exists():
@@ -222,7 +220,7 @@ def read_tensor(name: str) -> torch.Tensor:
             read_checkpoint_range(boot[name], 0, prod(shape), data)
             return torch.from_numpy(data).view(dtype).view(shape)
 
-        # Each worker reports the tensors it holds; the trainer waits for all of them.
+        # The trainer waits for every worker's layout
         if self.vllm_config.parallel_config.enable_expert_parallel:
             group = get_ep_group()
         else:
@@ -252,7 +250,7 @@ def _replay(self, step_dir: Path) -> None:
     def _apply_delta(self, step_dir: Path) -> None:
         start = time.perf_counter()
         prefetched, self.prefetched = self.prefetched, None
-        # Always wait: a running prefetch writes into the same host buffer.
+        # Always wait: a running prefetch writes into the same host buffer
         delta = prefetched[1].result() if prefetched is not None else None
         if prefetched is None or prefetched[0] != step_dir.as_posix():
             delta = self._read_delta(step_dir)
@@ -295,8 +293,8 @@ def _verify(self, step_dir: Path, checksums: dict[str, int]) -> None:
             raise RuntimeError(f"Sparse weight delta {step_dir.name} diverged from the trainer: {mismatched[:8]}")
 
     def _load_checkpoint(self, step_dir: Path) -> None:
-        """Reload every weight from the HF-compatible checkpoint (or sparse-run anchor) in ``step_dir``. In a
-        sparse run, each worker reads only the rows it holds, straight into its params."""
+        """Reload every weight from the HF checkpoint (or anchor) in ``step_dir``. In a sparse run, each worker reads
+        only the rows it holds, straight into its params."""
         model = self._model()
         if self.index_map is None:
             model_loader = get_model_loader(self.load_config)
diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py
index a8672442612..fa6a5458e34 100644
--- a/src/prime_rl/inference/vllm/worker/weight_transfer.py
+++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py
@@ -63,7 +63,6 @@ def record_copy_(destination: torch.Tensor, source: torch.Tensor, non_blocking:
             source_2d = source.view(1, -1) if source.dim() == 1 else source
             destination_2d = destination.view(1, -1) if destination.dim() == 1 else destination
             if name is not None and source.dtype != destination.dtype:
-                # e.g. an fp32 router loaded from a bf16 checkpoint: fine as long as it never changes
                 unpatchable.add(name)
                 return destination
             if (
@@ -307,9 +306,7 @@ def load_rows(self, name: str, rows: torch.Tensor, begin: int) -> None:
                 self._view(index, rows_, cols, offset, strides).copy_(as_bits(block))
 
     @torch.no_grad()
-    def apply_delta(
-        self, entries: list[tuple[str, int]], indices: torch.Tensor, delta: torch.Tensor, chunk: int = 1 << 25
-    ) -> None:
+    def apply_delta(self, entries: list[tuple[str, int]], indices: torch.Tensor, delta: torch.Tensor) -> None:
         """Add ``delta`` (raw-bit differences, wrapping in the bits dtype) to the params at the checkpoint
         ``indices`` of ``entries``."""
         changed = [name for name, _ in entries if name in self.unpatchable]
@@ -332,6 +329,7 @@ def apply_delta(
             for k in range(max_records)
         ]
         entry_ends = torch.tensor([count for _, count in entries], device=device).cumsum(0)
+        chunk = 1 << 25
         for start in range(0, indices.numel(), chunk):
             index = indices[start : start + chunk]
             source = delta[start : start + chunk]
diff --git a/src/prime_rl/trainer/optim/offload.py b/src/prime_rl/trainer/optim/offload.py
index d783cd1c78e..2d88109c0f5 100644
--- a/src/prime_rl/trainer/optim/offload.py
+++ b/src/prime_rl/trainer/optim/offload.py
@@ -719,8 +719,8 @@ def __init__(
         self.optimizer = optimizer
         self._initialized = False
         self._master_weights = master_weights
-        # Called on the H2D stream with each parameter's local shard and its new values (pinned host) instead of
-        # the copy, which it must do; the shard still holds the previous values
+        # Replaces the H2D copy of each parameter's new values (pinned host) into its local shard, which still
+        # holds the previous values; it must do the copy itself
         self.weight_hook: Callable[[torch.Tensor, torch.Tensor], None] | None = None
         self._chunks = self._build_chunks()
         # Reuse the transfer streams across steps: fresh streams each step land
diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 90233a7c435..c2de543869b 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -55,46 +55,30 @@
 class SparseDeltaEncoder:
     """Rank-local deltas of the trainer's weight shards, in HF coordinates.
 
-    Every rank keeps the wire-dtype bits of its own shards from the previous
-    broadcast (the snapshot) and diffs against them, so no weights are gathered.
-    The prime->HF conversion is replayed once on meta tensors: every HF tensor
-    must come out as a contiguous view of one trainer tensor (renames, expert
-    unstacking), which maps a changed trainer element to its HF name and index.
-
-    ``encode`` runs on a background thread and a side CUDA stream while the
-    next step trains, without collectives; it only reads the weights before
-    ``weights_read`` is set, and hands the delta to the writer in parts as it
-    goes. The snapshot is placed at the first broadcast after a training step: on
-    the GPU when it fits next to that step's peak memory, else in the host buffer,
-    page-locked (only between training steps: registering it can stall the
-    trainer's CUDA calls). It starts as the checkpoint the engines booted from,
-    read back from storage. With a GPU snapshot, the host buffer holds anchor
-    copies. The host buffer is allocated and faulted in at the startup broadcast.
-
-    Tensors the engines hold in fp8 (``set_fp8_layout``: HF name -> scale format,
-    from the engine layout) are quantized with vLLM's own per-block quantizer,
-    previous and current weights alike, and diffed as fp8 bytes. Their 128-row
-    blocks must not straddle trainer ranks, else the trainer tensor is gathered.
-    The quantizer is compiled ahead of time, so it can run on the encode thread
-    (columns a multiple of 128). Their block scales are diffed too, as tensors
-    ``"{hf_name}_scale_inv"`` (fp32, one element per 128x128 block) in the same
-    delta, so engines patch scales in place.
+    Each rank diffs its own shards against the wire-dtype bits of the previous broadcast (the snapshot), so no
+    weights are gathered. The prime->HF conversion is replayed once on meta tensors: every HF tensor must be a
+    contiguous view of one trainer tensor, which maps a changed trainer element to its HF name and flat index.
+
+    ``encode`` runs on a background thread and a side CUDA stream while the next step trains, without collectives.
+    The snapshot starts as the boot checkpoint, read back from storage, and lives on the GPU if it fits next to the
+    first step's peak memory, else in the host buffer.
+
+    Tensors the engines hold in fp8 are quantized with vLLM's own quantizer (old and new weights) and diffed as fp8
+    bytes; their changed block scales travel as ``"{hf_name}_scale_inv"`` tensors. fp8 tensors whose 128-row blocks
+    straddle ranks are gathered.
 
     With a fully offloaded optimizer (``bind_optimizer``), the tensors stored in its parameters are diffed in its
-    step instead (``optimizer_hook``): each parameter shard still holds the previous broadcast's weights when the
-    step writes its new values, so they need no snapshot, and anchors are written as the step produces them. The
-    encode only checks that those tensors still hold the values the step wrote.
+    step (``optimizer_hook``): each shard still holds the previous broadcast's weights when the step writes the new
+    ones, so they need no snapshot. ``encode`` then only checks that they still hold what the step wrote.
     """
 
-    # Elements per transient chunk of the diff (bf16: 128 MiB per temporary)
+    # Elements per diff chunk and changes per written part. Smaller in the optimizer step, which overlaps the final
+    # backward.
     CHUNK = 1 << 26
-    # Changes per part handed to the writer
     PART = 1 << 26
-    # Elements per transient chunk of the diff in the optimizer pass, which overlaps the final backward
     HOOK_CHUNK = 1 << 24
-    # Changes per part handed to the writer during the optimizer step
     HOOK_PART = 1 << 24
-    # GPU slots of the optimizer pass's diff thread, each two of the largest parameter shard
+    # GPU slots of the optimizer step's diff thread, each twice the largest parameter shard
     HOOK_SLOTS = 2
 
     def __init__(self, model: nn.Module, snapshot: str):
@@ -119,7 +103,7 @@ def __init__(self, model: nn.Module, snapshot: str):
             self.hf_tensors[name] = (list(tensor.shape), str(tensor.dtype).removeprefix("torch."))
         self.segments = {fqn: sorted(segs) for fqn, segs in segments.items()}
         self.numel = {fqn: tensor.numel() for fqn, tensor in meta.items()}
-        # fqn -> (ue8m0 scales, columns) for trainer tensors held in fp8; quantized as one [rows, columns] matrix
+        # fqn -> (ue8m0 scales, columns) of trainer tensors held in fp8, quantized as one [rows, columns] matrix
         self.fp8: dict[str, tuple[bool, int]] = {}
         self.fp8_layout: dict[str, str] | None = None
         self.compiled_fp8 = None
@@ -135,17 +119,16 @@ def __init__(self, model: nn.Module, snapshot: str):
         self.key_fqn: dict[str, str] = {}
         # key -> previous wire bits: GPU tensors, or views into the host buffer
         self.snapshot: dict[str, torch.Tensor] = {}
-        self.on_gpu: bool | None = None  # None until placed, at the first broadcast after a training step
+        self.on_gpu: bool | None = None
         self.host_buffer: torch.Tensor | None = None
         self.write_staging: torch.Tensor | None = None
         # The boot checkpoint's ``checkpoint_index``, and keys whose v0 is still to read from it
         self.boot_files: dict[str, tuple] = {}
         self.boot_pending: set[str] = set()
         self.host: dict[str, torch.Tensor] = {}
-        self.pinned = False
         self.initialized = False
-        # Diffed in the optimizer pass (``bind_optimizer``): parameter storage -> (key, fqn, numel, global flat
-        # offset, (shape, stride, storage offset) of the view), and that step's parts, checksums of the written wire bits and anchor
+        # Optimizer-step deltas: parameter storage -> (key, fqn, numel, global flat offset, (shape, stride, storage
+        # offset) of the view); the step's parts, diff futures and checksums of the written wire bits
         self.param_keys: dict[int, list[tuple[str, str, int, int, tuple]]] = defaultdict(list)
         self.optimizer_keys: set[str] = set()
         self.step_parts: list[tuple] = []
@@ -190,7 +173,7 @@ def set_fp8_layout(self, model: nn.Module, fp8: dict[str, str], is_master: bool)
 
     def _local_shard(self, key: str, value: torch.Tensor, is_master: bool) -> tuple[torch.Tensor, int] | None:
         """This rank's dim-0 slice of ``value`` and its global flat offset, or None when another rank owns the
-        elements. Flattened only where it is read (``reshape(-1)``), as that copies a non-contiguous slice."""
+        elements. Not flattened here: ``reshape(-1)`` copies a non-contiguous slice."""
         if not isinstance(value, DTensor):
             return (value.detach(), 0) if is_master else None
         if all(placement.is_replicate() for placement in value.placements):
@@ -225,10 +208,10 @@ def _find_gathered(self, model: nn.Module, is_master: bool) -> set[str]:
 
     @torch.no_grad()
     def local_shards(self, model: nn.Module, is_master: bool) -> list[tuple]:
-        """This rank's ``(key, fqn, shard, global flat offset, gathered)`` for ``encode``, taken on the
-        training thread: FSDP swaps module parameters during forward, but the sharded storage stays put.
-        ``gathered`` is ``(previous, current)`` bf16 full tensors of fp8 tensors whose blocks straddle ranks
-        (the all-gathers run here, as ``encode`` runs no collectives)."""
+        """This rank's ``(key, fqn, shard, global flat offset, gathered)``, taken on the training thread (FSDP swaps
+        module parameters during forward, but the sharded storage stays put). ``gathered`` is ``(previous,
+        current)`` full bf16 tensors of fp8 tensors whose blocks straddle ranks, all-gathered here because
+        ``encode`` runs no collectives."""
         shards = []
         for key, value in model.state_dict().items():
             fqn = resolve_fqn(model, key)
@@ -258,9 +241,9 @@ def _bits(self, key: str) -> torch.dtype:
         return bits_dtype(str(self.wire_dtypes[key]).removeprefix("torch."))
 
     def allocate_host(self) -> None:
-        """One host buffer for all shards (the CPU snapshot, or anchor copies of a GPU snapshot), faulted in
-        here, before training: on a node short of free memory, faulting it in during an anchor copy waits for
-        the kernel to reclaim page cache, for minutes, while the trainer waits for that copy."""
+        """One host buffer for all shards (the CPU snapshot, or anchor copies of a GPU snapshot). Faulted in before
+        training: on a node short of free memory, faulting it in during an anchor copy waits minutes for the kernel
+        to reclaim page cache while the trainer waits for that copy."""
         start = time.perf_counter()
         sizes = {key: -(-numel * self.wire_dtypes[key].itemsize // 8) * 8 for key, (numel, _) in self.shards.items()}
         flags = mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS | mmap.MAP_POPULATE
@@ -278,19 +261,20 @@ def allocate_host(self) -> None:
         )
 
     def _page_lock(self) -> None:
-        """Page-lock the host buffer, only while no training step runs: registering a large buffer can
-        stall the trainer's CUDA calls."""
+        """Only while no training step runs: registering a large buffer stalls the trainer's CUDA calls."""
         start = time.perf_counter()
         page_lock(self.host_buffer)
-        self.pinned = True
         get_logger().info(
             f"Page-locked the {self.host_buffer.numel() / 2**30:.1f} GiB sparse weight host buffer in "
             f"{time.perf_counter() - start:.1f}s"
         )
 
+    def _snapshot_bytes(self) -> int:
+        return sum(numel * self.wire_dtypes[key].itemsize for key, (numel, _) in self.shards.items())
+
     def _place(self, peak_reserved: int) -> None:
         """Keep the snapshot on the GPU if it fits next to the training step's peak, else on the host."""
-        num_bytes = sum(numel * self.wire_dtypes[key].itemsize for key, (numel, _) in self.shards.items())
+        num_bytes = self._snapshot_bytes()
         total = torch.cuda.get_device_properties(self.device).total_memory
         headroom = total - peak_reserved - num_bytes
         self.on_gpu = self.snapshot_placement == "auto" and headroom >= max(4 * 2**30, total // 20)
@@ -306,10 +290,8 @@ def track(self, shards: list[tuple]) -> None:
         self.key_fqn.update({key: fqn for key, fqn, *_ in shards})
 
     def copy_snapshot_to_host(self, give_up: Callable[[], bool]) -> bool:
-        """Copy a GPU snapshot (the last broadcast's weights) into the host buffer, for an anchor, through the
-        pinned staging buffer; False if ``give_up()`` turned true before it finished. Callers wait for earlier
-        readers of the host buffer first."""
-        assert self.on_gpu
+        """Copy a GPU snapshot into the host buffer for an anchor; False if ``give_up()`` turned true first. Callers
+        wait for earlier readers of the host buffer."""
         with torch.cuda.stream(self.stream):
             for key, bits in self.snapshot.items():
                 if give_up():
@@ -339,10 +321,9 @@ def anchor_parts(self, source: dict[str, torch.Tensor]) -> list[tuple[str, list[
 
     @torch.no_grad()
     def initialize(self, model: nn.Module, is_master: bool, peak_reserved: int, boot_dir: Path) -> None:
-        """At the first delta (v1), on the training thread: place the snapshot against the measured training
-        peak and allocate it (a CPU snapshot is the host buffer). The snapshot starts as v0, the
-        checkpoint in ``boot_dir`` the engines booted from, read as the encode goes (tensors gathered for fp8
-        are read here, before ``local_shards`` gathers them)."""
+        """At the first delta (v1), on the training thread: place and allocate the snapshot. It starts as v0, the
+        checkpoint in ``boot_dir``, read as the encode goes; tensors gathered for fp8 are read here, before
+        ``local_shards`` gathers them."""
         self.track([shard for shard in self.local_shards(model, is_master) if shard[0] not in self.optimizer_keys])
         self._place(peak_reserved)
         self.staging = torch.empty(1 << 28, dtype=torch.uint8, pin_memory=True)
@@ -365,33 +346,26 @@ def initialize(self, model: nn.Module, is_master: bool, peak_reserved: int, boot
         self.initialized = True
 
     def _read_boot(self, key: str) -> None:
-        """Fill the snapshot of ``key`` with this rank's slice of the boot checkpoint, through the pinned
-        staging buffer, on the current stream."""
+        """Fill the snapshot of ``key`` with this rank's slice of the boot checkpoint, on the current stream."""
         out = self.snapshot[key].view(self.wire_dtypes[key])
-        numel, offset = self.shards[key]
         staging = self.staging.numpy()
-        for start, name in self.segments[self.key_fqn[key]]:
-            begin, end = max(start, offset), min(start + prod(self.hf_tensors[name][0]), offset + numel)
-            if begin >= end:
-                continue
+        for name, _, hf_start, begin, end in self._slices(key, *self.shards[key]):
             dtype = self.boot_files[name][1]
             step = self.staging.numel() // dtype.itemsize
             for chunk in range(begin, end, step):
                 count = min(step, end - chunk)
-                read_checkpoint_range(self.boot_files[name], chunk - start, chunk - start + count, staging)
-                out[chunk - offset : chunk - offset + count].copy_(
-                    self.staging[: count * dtype.itemsize].view(dtype), non_blocking=True
-                )
+                first = hf_start + chunk - begin
+                read_checkpoint_range(self.boot_files[name], first, first + count, staging)
+                out[chunk : chunk + count].copy_(self.staging[: count * dtype.itemsize].view(dtype), non_blocking=True)
                 torch.cuda.current_stream().synchronize()
 
     @torch.no_grad()
     def encode(
         self, shards: list[tuple], ready: torch.cuda.Event, weights_read: threading.Event, emit: Callable
     ) -> int:
-        """Diff the shards against the snapshot and advance it, on the side stream once ``ready`` (recorded
-        after the optimizer step) has passed. Hands the delta to ``emit`` in parts ``(entries, gaps, value
-        bytes, new-weight checksum per HF tensor part)`` on the host as they are encoded. Sets ``weights_read``
-        once done reading the weights. Returns the number of changed fp8 block scales."""
+        """Diff the shards against the snapshot and advance it, on the side stream after ``ready``. Hands host parts
+        ``(entries, gaps, value bytes, checksums)`` to ``emit`` as they are encoded and sets ``weights_read`` once
+        done reading the weights. Returns the number of changed fp8 block scales."""
         torch.cuda.set_device(self.device)
         if not self.optimizer_keys:
             self.scale_changes = 0
@@ -405,8 +379,8 @@ def encode(
             weights_read.set()
 
     def stream_to_host(self, tensor: torch.Tensor):
-        """A tensor's bytes on the host; for a GPU tensor, in chunks through a pinned staging buffer, each to be
-        written before the next is taken."""
+        """A tensor's bytes on the host; for a GPU tensor, in pinned staging chunks, each to be consumed before the
+        next is taken."""
         data = tensor.reshape(-1).view(torch.uint8)
         if not data.is_cuda:
             yield data.numpy()
@@ -424,9 +398,8 @@ def stream_to_host(self, tensor: torch.Tensor):
     def _to_host(
         self, tensor: torch.Tensor, host: torch.Tensor | None = None, staging: torch.Tensor | None = None
     ) -> torch.Tensor:
-        """Copy a GPU tensor to pageable host memory (``host``, or a new tensor) through the pinned staging
-        buffer: a copy straight into pageable memory holds the CUDA driver and stalls the training step's
-        own CUDA calls."""
+        """Copy a GPU tensor to pageable host memory (``host``, or a new tensor) through pinned staging: a copy
+        straight into pageable memory holds the CUDA driver and stalls the training step's CUDA calls."""
         host = torch.empty(tensor.shape, dtype=tensor.dtype) if host is None else host
         staging = self.staging if staging is None else staging
         data, host_bytes = tensor.view(torch.uint8), host.view(torch.uint8)
@@ -456,9 +429,9 @@ def _segments(
 
     @staticmethod
     def _segment_sums(bits: torch.Tensor, begins: list[int], ends: list[int], firsts: list[int]) -> torch.Tensor:
-        """``bits_checksum`` of each ``bits[begin:end]`` whose first element is at HF flat index ``first``:
-        Σ bits[k]·(first + k - begin + 1), wrapping int64. One reduction per run of adjacent equal-length
-        segments (e.g. unstacked experts): few kernels, as the encode thread shares the GIL with training."""
+        """``bits_checksum`` of each ``bits[begin:end]`` whose first element is at HF flat index ``first``. One
+        reduction per run of adjacent equal-length segments (e.g. unstacked experts): few kernel launches, as the
+        encode thread shares the GIL with training."""
         sums, i = [], 0
         while i < len(begins):
             length, j = ends[i] - begins[i], i + 1
@@ -502,9 +475,9 @@ def _scale_part(self, fqn: str, first_block: int, previous: torch.Tensor, curren
         )
 
     def _diff(self, shards: list[tuple], emit: Callable) -> None:
-        """Diff ``shards`` against the snapshot on the current stream, writing the new bits to the snapshot.
-        Shards diffed in the optimizer pass are only checked against the values that pass wrote; their parts
-        go first. Copies the delta to the host and hands it to ``emit`` every ``PART`` changes."""
+        """Diff ``shards`` against the snapshot on the current stream and advance it, handing a host part to ``emit``
+        every ``PART`` changes. Shards diffed in the optimizer step are only checked against what it wrote; its
+        remaining parts go first."""
         pending, self.step_parts = self.step_parts, []
         num_pending = sum(len(part[2]) for part in pending)
         checks: list[tuple[str, torch.Tensor, torch.Tensor]] = []
@@ -538,14 +511,8 @@ def flush() -> None:
                 raise RuntimeError(f"Trainer weights changed outside the optimizer step, not sent: {changed}")
 
     def wait_for_diffs(self) -> None:
-        """Wait for the diff thread to finish the optimizer step's delta."""
         for diff in self.diffs:
             diff.result()
-        if self.diffs:
-            get_logger().debug(
-                f"Sparse weight delta in the optimizer step: waited {self.slot_wait:.2f}s for GPU slots, "
-                f"diff thread busy {self.diff_time:.2f}s"
-            )
         self.diffs.clear()
 
     def _host_part(self, pending: list[tuple], staging: torch.Tensor | None = None) -> tuple:
@@ -569,9 +536,9 @@ def _diff_shard(
         out: torch.Tensor | None,
         chunk: int = CHUNK,
     ) -> list[tuple]:
-        """The part entries of a flat shard's changes against ``previous`` (its wire bits), on the current
-        stream, in chunks of ``chunk`` elements (fp8: whole 128-row blocks, and their block scales). Writes
-        the new bits to ``out`` unless None."""
+        """The part entries of a flat shard's changes against ``previous`` (its wire bits), on the current stream, in
+        chunks of ``chunk`` elements (fp8: whole 128-row blocks, plus their block scales). Writes the new bits to
+        ``out`` unless None."""
         device = torch.device(f"cuda:{self.device}")
         wire = self.wire_dtypes[key]
         segment_names, begins, ends, shifts = self._segments(key, self.segments[fqn], offset, shard.numel())
@@ -658,10 +625,9 @@ def _checksum(self, bits: torch.Tensor) -> torch.Tensor:
         return torch.cat(sums).sum(0, keepdim=True) if sums else torch.zeros(1, dtype=torch.int64, device=bits.device)
 
     def bind_optimizer(self, model: nn.Module, is_master: bool, params: list[torch.Tensor]) -> None:
-        """Diff the trainer tensors stored in ``params`` (contiguous local shards of an offloaded optimizer's
-        parameters) in the optimizer pass (``optimizer_hook``), not against the snapshot, which only keeps the
-        other tensors and the fp8 tensors gathered across ranks. A trainer tensor may be a strided view of a
-        parameter (e.g. one projection of a packed gate/up weight)."""
+        """Diff the trainer tensors stored in ``params`` (the offloaded optimizer's contiguous local shards) in the
+        optimizer step (``optimizer_hook``); the snapshot keeps only the other tensors and gathered fp8 tensors. A
+        trainer tensor may be a strided view of a parameter (e.g. one half of a packed gate/up weight)."""
         assert all(param.is_contiguous() for param in params)
         stored = {param.untyped_storage().data_ptr() for param in params}
         residual = []
@@ -675,7 +641,6 @@ def bind_optimizer(self, model: nn.Module, is_master: bool, params: list[torch.T
             else:
                 residual.append(shard)
         self.track(residual)
-        # GPU slots for a parameter's previous and new values, diffed on the diff thread
         slot_bytes = 2 * -(-max(param.nbytes for param in params) // 256) * 256
         self.hook_slots = queue.SimpleQueue()
         for _ in range(self.HOOK_SLOTS):
@@ -683,18 +648,16 @@ def bind_optimizer(self, model: nn.Module, is_master: bool, params: list[torch.T
         self.differ = ThreadPoolExecutor(1)
         self.diff_stream = torch.cuda.Stream(self.device)
         self.diff_staging = torch.empty(1 << 28, dtype=torch.uint8, pin_memory=True)
-        num_bytes = sum(numel * self.wire_dtypes[key].itemsize for key, (numel, _) in self.shards.items())
+        num_bytes = self._snapshot_bytes()
         get_logger().info(
             f"Sparse weight deltas of {len(self.optimizer_keys)} tensors come from the optimizer step, "
             f"{len(self.shards)} tensors ({num_bytes / 2**30:.2f} GiB) are diffed against a snapshot"
         )
-        get_logger().debug(f"Diffed against the snapshot: {sorted(self.shards)}")
 
     def arm(self, anchor: AnchorWriter | None, emit: Callable) -> None:
-        """Collect the next optimizer step's delta, handing it to ``emit`` in parts on the host as the step goes
-        (and write its anchor slices, if ``anchor``)."""
+        """Collect the next optimizer step's delta, handing host parts to ``emit`` as the step goes, and write its
+        anchor slices if ``anchor``."""
         self.step_parts, self.expected, self.scale_changes, self.anchor, self.emit = [], {}, 0, anchor, emit
-        self.slot_wait = self.diff_time = 0.0
 
     def anchor_slices(self) -> list[tuple[str, list[int], int, int, torch.dtype]]:
         """Every ``(hf_name, shape, flat start, numel, wire dtype)`` slice of this rank's share of an anchor."""
@@ -708,9 +671,10 @@ def anchor_slices(self) -> list[tuple[str, list[int], int, int, torch.dtype]]:
 
     @torch.no_grad()
     def optimizer_hook(self, local_param: torch.Tensor, source: torch.Tensor) -> None:
-        """``FullCPUOffloadOptimizer.weight_hook``, on its H2D stream: copy the shard, which still holds the
-        previous broadcast's weights, and its new values ``source`` (host) into a GPU slot, write the new values,
-        and diff the slot on the diff thread. At anchor steps, writes the new values from the host first."""
+        """``FullCPUOffloadOptimizer.weight_hook``, on its H2D stream: copy the shard (still the previous
+        broadcast's weights) and its new values ``source`` (pinned host) into a GPU slot, write the new values and
+        diff the slot on the diff thread. At anchor steps, also writes the new values to the anchor from the
+        host."""
         start = local_param.storage_offset()
         entries = [
             (key, fqn, numel, offset, (size, stride, at - start))
@@ -722,9 +686,7 @@ def optimizer_hook(self, local_param: torch.Tensor, source: torch.Tensor) -> Non
         if not entries:
             local_param.copy_(source, non_blocking=True)
             return
-        start_wait = time.perf_counter()
         slot = self.hook_slots.get()
-        self.slot_wait += time.perf_counter() - start_wait
         half = slot.numel() // 2
         previous = slot[: source.nbytes].view(source.dtype)
         new = slot[half : half + source.nbytes].view(source.dtype)
@@ -747,9 +709,8 @@ def optimizer_hook(self, local_param: torch.Tensor, source: torch.Tensor) -> Non
 
     @torch.no_grad()
     def _diff_slot(self, entries: list[tuple], slot: torch.Tensor, previous, new, ready: torch.cuda.Event) -> None:
-        """Diff thread: the parts of one parameter's tensors from its slot, then free the slot. Every ``HOOK_PART``
-        changes, hands the parts so far to the host."""
-        start = time.perf_counter()
+        """Diff thread: the parts of one parameter's tensors from its slot, then free the slot. Hands the parts so far
+        to ``emit`` every ``HOOK_PART`` changes."""
         torch.cuda.set_device(self.device)
         try:
             with torch.cuda.stream(self.diff_stream):
@@ -771,7 +732,6 @@ def _diff_slot(self, entries: list[tuple], slot: torch.Tensor, previous, new, re
             with torch.cuda.stream(self.diff_stream):
                 self.emit(self._host_part(self.step_parts, self.diff_staging))
             self.step_parts = []
-        self.diff_time += time.perf_counter() - start
 
 
 def _model_dir(model: nn.Module) -> Path:
@@ -785,17 +745,15 @@ def _write_json(path: Path, value) -> None:
 
 
 class SparseFileSystemWeightSender(FileSystemWeightSender):
-    """Filesystem broadcast where steps only carry the delta against the
-    previous broadcast (see ``prime_rl.utils.sparse_weights``). The delta is
-    encoded and written in the background while the next step trains: the
-    trainer only waits, before its next step's final backward (which may run
-    the optimizer step), until the encode has read the weights. Full checkpoints (anchors) serve engines that fall off
-    the delta chain. Steps are committed by the writer thread once every rank's
-    files are on disk."""
+    """Filesystem broadcast where steps carry only the delta against the previous broadcast (format in
+    ``prime_rl.utils.sparse_weights``). The delta is encoded and written in the background while the next step
+    trains; the trainer waits only before its next final backward, until the encode has read the weights. Anchors
+    (full checkpoints) serve engines that fall off the delta chain. The writer thread commits a step once every
+    rank's files are on disk."""
 
     gathers_full_weights = False
-    # Seconds an anchor copy may hold back the next encode before it gives up (that anchor is skipped): the trainer
-    # waits for that encode, well below the offload pipeline's timeout, which other ranks hit while waiting for it
+    # Seconds an anchor copy may hold back the next encode before that anchor is skipped: the trainer waits for the
+    # encode, and other ranks would hit the offload pipeline's timeout
     ANCHOR_COPY_WAIT = 30.0
 
     def __init__(self, output_dir: Path, config: FileSystemWeightBroadcastConfig):
@@ -804,12 +762,12 @@ def __init__(self, output_dir: Path, config: FileSystemWeightBroadcastConfig):
         self.encoder: SparseDeltaEncoder | None = None
         self.last_step: int | None = None
         self.anchor_step: int | None = None
-        # Background work, each in order: encodes, delta writes and commits, periodic anchor writes
+        # Each in order: encodes, delta writes and commits, anchor writes
         self.encodes = ThreadPoolExecutor(1)
         self.writes = ThreadPoolExecutor(1)
         self.anchors = ThreadPoolExecutor(1)
-        # Background writes that read the snapshot (the next encode waits for them), and anchor writes from the
-        # host buffer of a GPU snapshot (the next anchor copy waits for them)
+        # Anchor writes that read the snapshot (the next encode waits for them), and anchor writes from the host
+        # buffer of a GPU snapshot (the next anchor copy waits for them)
         self.snapshot_writes: list[Future] = []
         self.anchor_writes: list[Future] = []
         self.errors: list[BaseException] = []
@@ -819,16 +777,16 @@ def __init__(self, output_dir: Path, config: FileSystemWeightBroadcastConfig):
         # The checkpoint the engines hold before the first delta: the model's (None), or a resumed run's first
         # broadcast
         self.boot_dir: Path | None = None
-        # With an offloaded optimizer (``diff_in_optimizer``), the step its next pass collects the delta of
         self.optimizer: "FullCPUOffloadOptimizer | None" = None
         self.model: nn.Module | None = None
+        # The step whose delta the next optimizer step collects
         self.armed_step: int | None = None
         self.streamed_dir: Path | None = None
         self.num_streamed = 0
 
     def diff_in_optimizer(self, optimizer: "FullCPUOffloadOptimizer") -> None:
-        """Take the deltas of the optimizer's parameters from its pass (``SparseDeltaEncoder.optimizer_hook``),
-        which needs a broadcast after every optimizer step."""
+        """Take the deltas of the optimizer's parameters from its step (``SparseDeltaEncoder.optimizer_hook``); needs
+        a broadcast after every optimizer step."""
         self.optimizer = optimizer
 
     def wait_for_weight_reads(self) -> None:
@@ -845,9 +803,8 @@ def _wait_for_encode(self) -> None:
         self._raise_errors()
 
     def _arm(self, step: int) -> None:
-        """Collect v{step}'s delta, and its anchor slices at anchor steps, in the coming optimizer pass."""
+        """Collect v{step}'s delta, and its anchor slices at anchor steps, in the coming optimizer step."""
         encoder = self.encoder
-        assert encoder is not None and self.optimizer is not None
         self._set_fp8_layout()
         if not encoder.optimizer_keys:
             params = [master.model_param.to_local() for master in self.optimizer._master_weights.values()]
@@ -866,14 +823,13 @@ def _arm(self, step: int) -> None:
         self.armed_step = step
 
     def _stream_part(self, part: tuple) -> None:
-        """A part of the delta the optimizer step produces, written while the step runs."""
         path = self.streamed_dir / rank_part(self.world.rank, self.num_streamed)
         self._submit(self.writes, self._write_part, path, part)
         self.num_streamed += 1
 
     def _set_fp8_layout(self) -> None:
         if self.encoder.fp8_layout is None:
-            # Every engine worker reports which tensors it holds in fp8 when it loads the first broadcast.
+            # Every engine worker reports which tensors it holds in fp8 at its first update
             broadcast_dir = self.step_dir(0).parent
             layout = [self._wait_for_engine_layout(broadcast_dir) if self.world.is_master else None]
             dist.broadcast_object_list(layout, src=0)
@@ -919,16 +875,14 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
             if self.optimizer is None:
                 encoder.allocate_host()
         if base_step is None and step == 0:
-            # The run's first broadcast is the checkpoint the engines booted from.
+            # The run's first broadcast is the checkpoint the engines booted from
             self.anchor_step = step
             if self.world.is_master:
-                manifest = {"step": step, "base_step": None, "ranks": 0, "tensors": encoder.hf_tensors}
-                _write_json(step_dir / DELTA_DIR / MANIFEST, manifest)
+                self._write_manifest(step_dir, step, None, 0)
                 WeightSender._finish(self, step, step_dir)
             return
         if base_step is None:
-            # A resumed run's startup broadcast: a full checkpoint, written from the weights before training. The
-            # next broadcast's delta is against it.
+            # A resumed run's first broadcast: a full checkpoint, which the next delta is against
             source = {key: as_bits(shard.reshape(-1).to(encoder.wire_dtypes[key])) for key, _, shard, _, _ in shards}
             torch.cuda.synchronize()
             self._write_full(step, step_dir, encoder.anchor_parts(source))
@@ -940,8 +894,6 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
                 model, self.world.is_master, torch.cuda.max_memory_reserved(), self.boot_dir or _model_dir(model)
             )
         shards = encoder.local_shards(model, self.world.is_master)
-        # Periodic anchors only serve engines that fall off the chain: written in the background, each after
-        # the previous one.
         anchor = encoder.anchor if self.optimizer is not None else step % self.config.anchor_interval == 0
         ready = torch.cuda.Event()
         ready.record()
@@ -956,7 +908,6 @@ def _wait_for_engine_layout(self, broadcast_dir: Path) -> dict[str, str]:
             if time.monotonic() - start > self.timeout:
                 raise TimeoutError(f"Not every engine worker reported its weight layout within {self.timeout}s")
             time.sleep(0.1)
-        self.logger.debug(f"Waited {time.monotonic() - start:.2f}s for the engine weight layout")
         return layout
 
     @staticmethod
@@ -966,13 +917,12 @@ def _wait(writes: list[Future]) -> None:
         writes.clear()
 
     def _write_periodic_anchor(self, step: int, step_dir: Path) -> None:
-        """Write this rank's slices of the last broadcast's full checkpoint in the background. A GPU snapshot
-        goes through the host buffer, so the next encode does not wait for the write; a CPU snapshot is
-        written from itself, and the next encode waits for that."""
+        """Write this rank's slices of the last broadcast in the background. A GPU snapshot is first copied to the
+        host buffer, so the next encode waits only for that copy; a CPU snapshot is written from itself, and the
+        next encode waits for the write."""
         encoder = self.encoder
-        assert encoder is not None
         if encoder.on_gpu:
-            # The next encode waits for this copy (it advances the snapshot)
+
             def held_too_long() -> bool:
                 queued_at = self.encode_queued_at
                 return queued_at is not None and time.monotonic() - queued_at > self.ANCHOR_COPY_WAIT
@@ -1002,18 +952,16 @@ def _encode(
         weights_read: threading.Event,
         anchor: bool | AnchorWriter | None,
     ) -> None:
-        """Encode thread: diff v{step}, handing the delta to the writer in parts as it is encoded. ``anchor``: write
-        a periodic anchor, or finish the one the optimizer pass wrote."""
+        """Encode thread: diff v{step}, handing the delta to the writer in parts. ``anchor``: write a periodic
+        anchor (True), or finish the one the optimizer step wrote."""
         encoder = self.encoder
-        assert encoder is not None
         self.encode_queued_at = None
         self._wait(self.snapshot_writes)
         start = time.perf_counter()
-        # The delta parts the optimizer step already wrote (a fully offloaded optimizer)
         encoder.wait_for_diffs()
         streamed_dir, num_parts = (self.streamed_dir, self.num_streamed) if self.optimizer is not None else (None, 0)
         if self.world.is_master:
-            self._submit(self.writes, self._start_delta, step, base_step, step_dir)
+            self._submit(self.writes, self._write_manifest, step_dir, step, base_step, self.world.world_size)
         if streamed_dir is not None:
             self._submit(self.writes, self._move_streamed, streamed_dir, step_dir / DELTA_DIR)
 
@@ -1024,36 +972,24 @@ def emit(part: tuple) -> None:
             num_parts += 1
 
         scale_changes = encoder.encode(shards, ready, weights_read, emit)
-        # Free the gathered fp8 tensors now, not after the anchor copy below, which may overlap the next step's
-        # memory peak.
+        # Free the gathered fp8 tensors before the anchor copy below, which may overlap the next step's memory peak
         shards.clear()
         encode_time = time.perf_counter() - start
         args = (step, base_step, step_dir, num_parts, start, encode_time, scale_changes)
         self._submit(self.writes, self._finish_delta, *args)
         if isinstance(anchor, AnchorWriter):
-            # The snapshot of the other tensors holds v{step} until the next encode, on this thread
-            for name, _, start, values in encoder.anchor_parts(encoder.snapshot):
-                written = 0
-                for chunk in encoder.stream_to_host(values):
-                    anchor.write(name, start, chunk, written)
-                    written += chunk.nbytes
+            # The snapshot tensors hold v{step} until the next encode, which runs on this thread
+            anchor.write_parts(encoder.anchor_parts(encoder.snapshot), encoder.stream_to_host)
             self._submit(self.anchors, self._close_anchor, step, step_dir, anchor)
         elif anchor:
             self._write_periodic_anchor(step, step_dir)
 
-    def _start_delta(self, step: int, base_step: int, step_dir: Path) -> None:
-        """The manifest comes first, so engines can read part files as they appear."""
-        assert self.encoder is not None
-        manifest = {
-            "step": step,
-            "base_step": base_step,
-            "ranks": self.world.world_size,
-            "tensors": self.encoder.hf_tensors,
-        }
+    def _write_manifest(self, step_dir: Path, step: int, base_step: int | None, ranks: int) -> None:
+        """Written before the parts, so engines can read part files as they appear."""
+        manifest = {"step": step, "base_step": base_step, "ranks": ranks, "tensors": self.encoder.hf_tensors}
         _write_json(step_dir / DELTA_DIR / MANIFEST, manifest)
 
     def _write_part(self, path: Path, part: tuple) -> None:
-        assert self.encoder is not None
         entries, gaps, values, checksums = part
         self.part_bytes += write_rank_file(
             path, entries, gaps.numpy(), values.numpy(), self.encoder.hf_tensors, checksums
@@ -1076,8 +1012,8 @@ def _finish_delta(
         encode_time: float,
         scale_changes: int,
     ) -> None:
-        """After this rank's last part: its parts marker. The master then waits for every rank's marker and
-        commits the step. Runs on the writer thread, without collectives."""
+        """After this rank's last part, its parts marker. The master then waits for every rank's marker and commits
+        the step (on the writer thread, without collectives)."""
         delta_dir = step_dir / DELTA_DIR
         _write_json(
             delta_dir / rank_parts_marker(self.world.rank),
@@ -1100,8 +1036,7 @@ def _finish_delta(
         )
         self.logger.info(
             f"Sparse weight delta v{step} (base v{base_step}): {num_bytes / 2**20:.1f} MiB, "
-            f"encode {encode_time:.2f}s (parts written as encoded), committed {time.perf_counter() - start:.2f}s after "
-            f"the encode started{scales}"
+            f"encode {encode_time:.2f}s, committed {time.perf_counter() - start:.2f}s after the encode started{scales}"
         )
 
     def _write_full(self, step: int, step_dir: Path, parts: list) -> None:
@@ -1111,9 +1046,7 @@ def _write_full(self, step: int, step_dir: Path, parts: list) -> None:
             WeightSender._finish(self, step, step_dir)
 
     def _write_anchor(self, step: int, step_dir: Path, parts: list) -> None:
-        start = time.perf_counter()
         write_anchor_shards(parts, step_dir, self.world.rank, self.world.world_size, self.encoder.stream_to_host)
-        self.logger.debug(f"Wrote the {step_dir.name} anchor in {time.perf_counter() - start:.1f}s")
         self._wait_for_anchor(step, step_dir)
 
     def _close_anchor(self, step: int, step_dir: Path, anchor: AnchorWriter) -> None:
@@ -1133,5 +1066,5 @@ def _finish(self, step: int, step_dir: Path) -> None:
         """Committed in ``_broadcast`` at v0, else in the background (``_finish_delta``, ``_write_full``)."""
 
     def _clean(self, step: int) -> None:
-        # Sparse chains replay from their anchor, so keep everything since it.
+        # Engines replay from the last anchor, so keep everything since it
         super()._clean(step if self.anchor_step is None else min(step, self.anchor_step + 1))
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index f06b6ccc025..56e994bcc27 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -49,11 +49,10 @@ def bits_dtype(dtype: str) -> torch.dtype:
 
 
 def quantize_fp8_blocks(weight: torch.Tensor, ue8m0: bool, compiled=None) -> tuple[torch.Tensor, torch.Tensor]:
-    """vLLM's online ``fp8_per_block`` quantization of a 2-D weight: fp8 bits (uint8) and fp32 block
-    scales. With ``ue8m0``, also DeepGEMM's requantization to power-of-two scales. Uses vLLM's own
-    compiled kernels: a reimplementation rounds differently on a few elements. ``compiled``: the
-    ahead-of-time compiled kernels from ``compile_fp8_quantizer`` (columns a multiple of 128), which,
-    unlike vLLM's torch.compile'd function, may run on a thread next to a compiled training step."""
+    """vLLM's online ``fp8_per_block`` quantization of a 2-D weight: fp8 bits (uint8) and fp32 block scales; with
+    ``ue8m0``, also DeepGEMM's requantization to power-of-two scales. Uses vLLM's own kernels, as a reimplementation
+    rounds differently on a few elements. ``compiled`` (from ``compile_fp8_quantizer``) may, unlike vLLM's
+    torch.compile'd function, run on a thread next to a compiled training step."""
     from vllm.model_executor.layers.quantization.utils.fp8_utils import requant_weight_ue8m0_inplace
     from vllm.utils.deep_gemm import per_block_cast_to_fp8
 
@@ -73,10 +72,9 @@ def quantize_fp8_blocks(weight: torch.Tensor, ue8m0: bool, compiled=None) -> tup
 
 
 def compile_fp8_quantizer() -> tuple:
-    """vLLM's ``per_block_cast_to_fp8`` compiled ahead of time with AOTInductor for any shape (bf16 input;
-    and fp32 input with ue8m0 scales, for the requantization). Matches vLLM's compiled function bit for
-    bit for columns a multiple of 128. Compile on the thread that trains: compiling next to a compiled
-    training step breaks activation checkpointing there."""
+    """vLLM's ``per_block_cast_to_fp8`` compiled with AOTInductor for any shape (bf16 input, and fp32 input with
+    ue8m0 scales for the requantization). Bit-identical to vLLM's for columns a multiple of 128. Compile on the
+    training thread: compiling next to a compiled training step breaks its activation checkpointing."""
     import tempfile
 
     from torch.export import Dim
@@ -148,11 +146,10 @@ def write_rank_file(
     checksums: dict[str, int],
     frame_size: int = 1 << 20,
 ) -> int:
-    """Write a rank's delta as zstd frames of about ``frame_size`` changes each (whole entries), so
-    readers can skip frames without entries they hold. ``values`` holds each entry's ``zigzag``
-    differences in its bits dtype. Gaps and values are stored as uint8; those of 255 or more go to
-    the int32 escape stream (gaps first). ``checksums`` (HF name -> ``bits_checksum`` of the new weights
-    this rank holds) let engines verify every step. Returns the file size."""
+    """Write a rank's delta as zstd frames of about ``frame_size`` changes (whole entries), so readers can skip
+    frames without entries they hold. ``values`` holds each entry's ``zigzag`` differences in its bits dtype;
+    ``checksums`` maps HF names to the ``bits_checksum`` of the new weights this rank holds. Returns the file
+    size."""
     options = {pyzstd.CParameter.compressionLevel: 1, pyzstd.CParameter.checksumFlag: 1}
     frames, size = [[]], 0
     for index, (_, count) in enumerate(entries):
@@ -202,7 +199,7 @@ def decode_gaps(gaps: torch.Tensor, counts: torch.Tensor) -> torch.Tensor:
 
 
 def _rank_file_frames(
-    file: BinaryIO, tensors: dict, names: set[str] | None
+    file: BinaryIO, names: set[str] | None
 ) -> tuple[list[tuple[list[tuple[str, int]], list[tuple[int, int]]]], dict[str, int]]:
     """The frames of an open rank file with entries in ``names``: ``(entries, [(offset, size)] of the
     gaps, escapes and values streams)``, and the file's checksums."""
@@ -261,10 +258,10 @@ def write_file(path: Path, data: bytes | Iterable) -> None:
 
 
 class AnchorWriter:
-    """One trainer rank's share of a full checkpoint: ``(hf_name, shape, flat start, numel, dtype)`` slices,
-    stored as tensors ``"{hf_name}@{start}"`` in ~1 GiB safetensors files whose headers are written first, so
-    slices can be written in any order, from any thread (``write_all``: in parallel). The files are written next to ``step_dir`` and moved
-    into it by ``close``, followed by the ``anchor_rank{rank}.done`` marker (the world size)."""
+    """One trainer rank's share of a full checkpoint: ``(hf_name, shape, flat start, numel, dtype)`` slices, stored
+    as tensors ``"{hf_name}@{start}"`` in ~1 GiB safetensors files whose headers are written first, so slices can
+    be written in any order and from any thread. The files are written next to ``step_dir``; ``close`` moves them
+    in and writes the ``anchor_rank{rank}.done`` marker (holding the world size)."""
 
     def __init__(self, slices: list[tuple[str, list[int], int, int, torch.dtype]], step_dir: Path, rank: int):
         self.step_dir, self.rank = step_dir, rank
@@ -272,7 +269,7 @@ def __init__(self, slices: list[tuple[str, list[int], int, int, torch.dtype]], s
         self.tmp_dir.mkdir(parents=True, exist_ok=True)
         self.files: list[tuple[str, int]] = []
         self.offsets: dict[tuple[str, int], tuple[int, int]] = {}
-        # Slices go round-robin to the files: a file is written at ~0.5 GiB/s, several at once in parallel
+        # Round-robin over the files, which are written in parallel
         num_files = max(1, -(-sum(numel * dtype.itemsize for *_, numel, dtype in slices) // 2**30))
         groups = [slices[index::num_files] for index in range(num_files)]
         self.pool = ThreadPoolExecutor(8)
@@ -309,6 +306,17 @@ def write(self, name: str, start: int, data, at: int = 0) -> None:
         fd, offset = self.offsets[(name, start)]
         self._pwrite(fd, data, offset + at)
 
+    def write_parts(
+        self, parts: list[tuple[str, list[int], int, torch.Tensor]], to_host: Callable[[torch.Tensor], Iterable]
+    ) -> None:
+        """Write ``(hf_name, shape, flat start, flat values)`` slices; ``to_host`` yields each slice's bytes, each
+        written before the next is taken."""
+        for name, _, start, values in parts:
+            written = 0
+            for chunk in to_host(values):
+                self.write(name, start, chunk, written)
+                written += chunk.nbytes
+
     def write_all(self, items: list[tuple[str, int, Any]]) -> None:
         """``write`` each ``(name, start, data)``, in parallel."""
         for future in [self.pool.submit(self.write, *item) for item in items]:
@@ -329,19 +337,13 @@ def write_anchor_shards(
     step_dir: Path,
     rank: int,
     world_size: int,
-    to_host: Callable[[torch.Tensor], Iterable] = lambda values: [values.view(torch.uint8).numpy()],
+    to_host: Callable[[torch.Tensor], Iterable],
 ) -> None:
-    """Write one trainer rank's share of a full checkpoint (``AnchorWriter``) from ``(hf_name, shape, flat
-    start, flat values)`` slices. ``to_host`` yields each slice's bytes on the host, each written before the
-    next is taken."""
+    """Write one trainer rank's share of a full checkpoint from ``AnchorWriter.write_parts`` slices."""
     writer = AnchorWriter(
         [(name, shape, start, values.numel(), values.dtype) for name, shape, start, values in parts], step_dir, rank
     )
-    for name, _, start, values in parts:
-        written = 0
-        for chunk in to_host(values):
-            writer.write(name, start, chunk, written)
-            written += chunk.nbytes
+    writer.write_parts(parts, to_host)
     writer.close(world_size)
 
 
@@ -384,7 +386,7 @@ def read_checkpoint_range(entry: tuple, begin: int, end: int, out: np.ndarray) -
 
 
 def anchor_complete(step_dir: Path) -> bool:
-    """Whether ``step_dir`` holds a complete full checkpoint written by ``write_anchor_shards``."""
+    """Whether every trainer rank has written its share of an anchor in ``step_dir``."""
     revalidate(step_dir)
     markers = list(step_dir.glob("anchor_rank*.done"))
     return bool(markers) and len(markers) == int(markers[0].read_text())
@@ -421,12 +423,11 @@ def rank_parts_marker(rank: int) -> str:
 def read_delta(
     step_dir: Path, pool: ThreadPoolExecutor, names: set[str] | None = None
 ) -> tuple[list[tuple[list[tuple[str, int]], np.ndarray, np.ndarray]], dict[str, int]]:
-    """Read and decompress a delta step dir into ``(entries, gaps, value bytes)`` per frame in host
-    memory (values: new minus previous raw bits, in each entry's bits dtype) on ``pool``; ``decode_gaps``
-    turns the gaps into flat indices on the device that applies them. With ``names``, keep only those
-    tensors (frames without any are not read). Part files are read as the trainer writes them, until
-    every rank has written its parts marker. Also returns the new weights' ``bits_checksum`` per HF
-    tensor, summed over ranks."""
+    """Read and decompress a delta into host ``(entries, gaps, value bytes)`` per frame (values: new minus previous
+    raw bits, in each entry's bits dtype); ``decode_gaps`` turns the gaps into flat indices. With ``names``, keep
+    only those tensors and skip frames without any. Part files are read as the trainer writes them, until every
+    rank's parts marker exists. Also returns the new weights' ``bits_checksum`` per HF tensor, summed over
+    ranks."""
     manifest = read_manifest(step_dir)
     tensors = manifest["tensors"]
     delta_dir = step_dir / DELTA_DIR
@@ -446,7 +447,7 @@ def read_delta(
                 if (num_parts[rank] is None or next_part[rank] < num_parts[rank]) and path.exists():
                     file = stack.enter_context(open(path, "rb"))
                     files.append(file)
-                    frames, partial = _rank_file_frames(file, tensors, names)
+                    frames, partial = _rank_file_frames(file, names)
                     jobs.extend(pool.submit(_read_frame, file.fileno(), frame, tensors, names) for frame in frames)
                     for tensor_name, value in partial.items():
                         checksums[tensor_name] = add_checksums(checksums.get(tensor_name, 0), value)
@@ -461,8 +462,8 @@ def read_delta(
 
 
 def pin(num_bytes: int) -> torch.Tensor:
-    """Page-locked host memory of exactly ``num_bytes``, outside torch's caching host allocator
-    (which rounds sizes up and keeps freed blocks). Release it with ``unpin``."""
+    """Page-locked host memory of exactly ``num_bytes``, outside torch's caching host allocator (which rounds sizes
+    up and keeps freed blocks). Release it with ``unpin``."""
     return page_lock(torch.from_numpy(np.empty(num_bytes, np.uint8)))
 
 

From 8e87fe986a5b120ead2f313ba00dd731ab39da29 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sat, 17 Oct 2026 18:12:09 +0000
Subject: [PATCH 52/88] docs(weights): concise sparse broadcast docs and config
 descriptions

---
 docs/scaling.md                                   | 15 ++++++++++++++-
 .../src/prime_rl/configs/shared.py                | 13 ++++++-------
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index c3a8b83a496..b65353b4c00 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -276,7 +276,20 @@ By default, the trainer and inference worker each allocate one transfer arena. S
 
 ### Sparse filesystem weight broadcast
 
-With `[weight_broadcast] type = "filesystem"` and `sparse = true`, the trainer publishes only the elements of the weights that changed since the previous broadcast, in the engine's weight format (bf16, or fp8 for engines that quantize online to `fp8_per_block`). Each trainer rank diffs its own shards against a copy of the previous broadcast (`snapshot = "auto"` keeps it on the GPU when it fits next to the first training step's peak memory, else in pinned host memory), so no weights are gathered. The diff runs on a side CUDA stream while the next step trains; the trainer only waits for it before the next step's final backward (which may already run the optimizer step). With AdamW at RL learning rates about 1% of bf16 elements change per step. A full checkpoint (anchor) is written in the background every `anchor_interval` steps (default 25); the first broadcast is the checkpoint the engines booted from, and the trainer reads it back (its own shard of each tensor) to encode the first delta. Each rank allocates and faults in a host buffer of the snapshot's size (2 B per parameter the rank holds) before training: the CPU snapshot, or the anchor copies of a GPU snapshot. Trainer hosts need that much free memory (e.g. next to a fully offloaded Muon optimizer). An anchor whose copy holds back the next step's diff for more than 30 s is skipped with a warning; engines then replay from the previous anchor. Engines read and decode a delta while they keep serving, patch their weights in place during the pause and verify every step against trainer checksums. For fp8 engines, changed block scales are part of the delta. An engine that is not at a delta's base step reloads the last anchor, each worker reading only the rows it holds, and replays the deltas since. Weights the engine transforms after loading other than online `fp8_per_block`, and LoRA, are not supported.
+With `[weight_broadcast] type = "filesystem"` and `sparse = true`, the trainer publishes only the weight elements that changed since the previous broadcast, in the engine's weight format (bf16, or fp8 for engines that quantize online to `fp8_per_block`). Engines read and decode each delta while they keep serving, patch their weights in place during the pause, and verify every update against trainer checksums.
+
+```toml
+[weight_broadcast]
+type = "filesystem"
+sparse = true
+anchor_interval = 25  # full checkpoint every N steps
+snapshot = "auto"     # previous weights on the GPU if they fit, else in pinned host memory
+```
+
+- With `model.offload = "full"`, deltas come from the optimizer step: no snapshot of the weights is kept, except a small one for tensors that are not optimizer parameters.
+- Otherwise each trainer rank diffs its own shards against a snapshot of the previous broadcast (2 B per parameter) while the next step trains; the trainer only waits before its next final backward. Each rank also allocates a host buffer of the same size before training, so trainer hosts need that much free memory.
+- Anchors (full checkpoints) are written every `anchor_interval` steps. An engine that joins, restarts or fails an update loads the latest complete anchor, reading only the rows it holds, and replays the deltas since. An anchor whose copy holds back the next step's diff for more than 30 s is skipped. Anchors written in the optimizer step slow that step down, so prefer a larger `anchor_interval` there.
+- Not supported: LoRA, fp8 tensors whose columns are not a multiple of 128, and engines that transform weights after loading other than online `fp8_per_block` (e.g. FlashInfer TRTLLM/CUTLASS MoE, Marlin). fp8 deltas need the same vLLM and torch versions on the trainer and the engines.
 
 ### Custom Templates
 
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
index 4ac8a6f2675..38c0866944f 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
@@ -49,18 +49,17 @@ class FileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig):
     type: Literal["filesystem"] = "filesystem"
 
     sparse: bool = False
-    """Between anchors, publish only the delta of the wire-dtype weights against the previous
-    broadcast. Each trainer rank diffs its own shards; engines decode the delta while serving and
-    patch their weights in place during the pause. Not supported with LoRA."""
+    """Between anchors, publish only the weight elements that changed since the previous broadcast.
+    Engines patch their weights in place and verify every update. Not supported with LoRA."""
 
     anchor_interval: int = Field(25, ge=1)
-    """With ``sparse``, publish a full checkpoint every this many steps (and at the first broadcast).
-    Engines that fall off the delta chain reload the last anchor and replay the deltas since."""
+    """With ``sparse``, publish a full checkpoint every this many steps. Engines that join or fall
+    off the delta chain load the last one and replay the deltas since."""
 
     snapshot: Literal["auto", "cpu"] = "auto"
     """With ``sparse``, where each trainer rank keeps the previous broadcast's weights to diff against.
-    ``auto`` keeps them on the GPU if they fit next to the first training step's peak memory, else in
-    page-locked host memory; ``cpu`` always in host memory."""
+    ``auto``: on the GPU if they fit next to the first training step's peak memory, else in pinned
+    host memory; ``cpu``: always in host memory."""
 
 
 class InMemoryWeightBroadcastConfig(BaseWeightBroadcastConfig):

From 129b22552e5c387f86c9270c65eaa0dd7ea4444b Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sat, 17 Oct 2026 18:46:29 +0000
Subject: [PATCH 53/88] fix(weights): fail the trainer when the optimizer-step
 diff fails; take the engine layout at the startup broadcast

---
 src/prime_rl/transports/weights/sparse.py | 30 ++++++++++++++++-------
 1 file changed, 21 insertions(+), 9 deletions(-)

diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index c2de543869b..1e8a24247c9 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -880,13 +880,17 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
             if self.world.is_master:
                 self._write_manifest(step_dir, step, None, 0)
                 WeightSender._finish(self, step, step_dir)
-            return
-        if base_step is None:
+        elif base_step is None:
             # A resumed run's first broadcast: a full checkpoint, which the next delta is against
             source = {key: as_bits(shard.reshape(-1).to(encoder.wire_dtypes[key])) for key, _, shard, _, _ in shards}
             torch.cuda.synchronize()
             self._write_full(step, step_dir, encoder.anchor_parts(source))
             self.boot_dir = step_dir
+        if base_step is None:
+            if self.optimizer is not None:
+                # Engines report their layout as they load this broadcast. Take it (and compile the fp8 quantizer)
+                # now, not in the first optimizer step, next to its activations.
+                self._set_fp8_layout()
             return
         self._set_fp8_layout()
         if not encoder.initialized:
@@ -958,12 +962,7 @@ def _encode(
         self.encode_queued_at = None
         self._wait(self.snapshot_writes)
         start = time.perf_counter()
-        encoder.wait_for_diffs()
-        streamed_dir, num_parts = (self.streamed_dir, self.num_streamed) if self.optimizer is not None else (None, 0)
-        if self.world.is_master:
-            self._submit(self.writes, self._write_manifest, step_dir, step, base_step, self.world.world_size)
-        if streamed_dir is not None:
-            self._submit(self.writes, self._move_streamed, streamed_dir, step_dir / DELTA_DIR)
+        num_parts = 0
 
         def emit(part: tuple) -> None:
             nonlocal num_parts
@@ -971,7 +970,20 @@ def emit(part: tuple) -> None:
             self._submit(self.writes, self._write_part, path, part)
             num_parts += 1
 
-        scale_changes = encoder.encode(shards, ready, weights_read, emit)
+        try:
+            encoder.wait_for_diffs()
+            if self.optimizer is not None:
+                num_parts = self.num_streamed
+            if self.world.is_master:
+                self._submit(self.writes, self._write_manifest, step_dir, step, base_step, self.world.world_size)
+            if self.optimizer is not None:
+                self._submit(self.writes, self._move_streamed, self.streamed_dir, step_dir / DELTA_DIR)
+            scale_changes = encoder.encode(shards, ready, weights_read, emit)
+        except BaseException as error:
+            # Recorded before the trainer wakes up, so it fails at its next wait instead of training on
+            self.errors.append(error)
+            weights_read.set()
+            raise
         # Free the gathered fp8 tensors before the anchor copy below, which may overlap the next step's memory peak
         shards.clear()
         encode_time = time.perf_counter() - start

From d445d80f05c0fb8c5353fea435c106a1fbce5590 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sun, 18 Oct 2026 09:15:14 +0000
Subject: [PATCH 54/88] feat(weights): sparse deltas from the optimizer step
 for every optimizer, drop the snapshot path

- Muon (update_hook around each chunk's update), AdamW and SignSGD on the GPU (ChunkedUpdates: the step a chunk
  at a time) hand each chunk's old wire bits to the encoder right before the update; a diff thread diffs them
  against the new values, the bf16 tensors of a chunk as one flat tensor, and writes the parts and anchor slices
  as the step goes. Full offload keeps FullCPUOffloadOptimizer.weight_hook.
- Removed: the snapshot (placement, `snapshot` option, host buffer and prefault, v1 read-back of the boot
  checkpoint, anchor-copy skip). Tensors that are not optimizer parameters keep a small copy, taken from the live
  weights at each full checkpoint.
- A broadcast not collected in the optimizer step (first broadcast, SFT online evals) publishes a full checkpoint.
- The diff thread waits for the GPU without holding the GIL (no syncing ops).
---
 docs/scaling.md                               |   7 +-
 .../src/prime_rl/configs/shared.py            |   5 -
 src/prime_rl/trainer/optim/__init__.py        |  15 +-
 src/prime_rl/trainer/optim/base.py            |  42 +
 src/prime_rl/trainer/rl/train.py              |   3 +-
 src/prime_rl/trainer/sft/train.py             |   7 +-
 src/prime_rl/transports/weights/sparse.py     | 843 +++++++++---------
 src/prime_rl/utils/sparse_weights.py          |  25 +-
 8 files changed, 485 insertions(+), 462 deletions(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index b65353b4c00..f84657a0d60 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -283,12 +283,11 @@ With `[weight_broadcast] type = "filesystem"` and `sparse = true`, the trainer p
 type = "filesystem"
 sparse = true
 anchor_interval = 25  # full checkpoint every N steps
-snapshot = "auto"     # previous weights on the GPU if they fit, else in pinned host memory
 ```
 
-- With `model.offload = "full"`, deltas come from the optimizer step: no snapshot of the weights is kept, except a small one for tensors that are not optimizer parameters.
-- Otherwise each trainer rank diffs its own shards against a snapshot of the previous broadcast (2 B per parameter) while the next step trains; the trainer only waits before its next final backward. Each rank also allocates a host buffer of the same size before training, so trainer hosts need that much free memory.
-- Anchors (full checkpoints) are written every `anchor_interval` steps. An engine that joins, restarts or fails an update loads the latest complete anchor, reading only the rows it holds, and replays the deltas since. An anchor whose copy holds back the next step's diff for more than 30 s is skipped. Anchors written in the optimizer step slow that step down, so prefer a larger `anchor_interval` there.
+- RL: deltas come from the optimizer step, with every optimizer and offload mode. Right before the step updates a chunk of parameters, their shards still hold the previous broadcast's weights; the trainer diffs the old and new values of each chunk as the step goes and writes the delta out. No copy of the weights is kept, except of the few tensors that are not optimizer parameters (buffers, frozen parameters, fp8 tensors whose 128-row blocks straddle trainer ranks). The trainer checks that no weight changed outside the optimizer step.
+- A broadcast that does not follow every optimizer step (the startup broadcast of a resumed run, SFT online evals) publishes a full checkpoint.
+- Anchors (full checkpoints) are written every `anchor_interval` steps, from the optimizer step. They slow that step down, so prefer a larger `anchor_interval`. An engine that joins, restarts or fails an update loads the latest complete anchor, reading only the rows it holds, and replays the deltas since.
 - Not supported: LoRA, fp8 tensors whose columns are not a multiple of 128, and engines that transform weights after loading other than online `fp8_per_block` (e.g. FlashInfer TRTLLM/CUTLASS MoE, Marlin). fp8 deltas need the same vLLM and torch versions on the trainer and the engines.
 
 ### Custom Templates
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
index 38c0866944f..c054e9c236b 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
@@ -56,11 +56,6 @@ class FileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig):
     """With ``sparse``, publish a full checkpoint every this many steps. Engines that join or fall
     off the delta chain load the last one and replay the deltas since."""
 
-    snapshot: Literal["auto", "cpu"] = "auto"
-    """With ``sparse``, where each trainer rank keeps the previous broadcast's weights to diff against.
-    ``auto``: on the GPU if they fit next to the first training step's peak memory, else in pinned
-    host memory; ``cpu``: always in host memory."""
-
 
 class InMemoryWeightBroadcastConfig(BaseWeightBroadcastConfig):
     host: str = "localhost"
diff --git a/src/prime_rl/trainer/optim/__init__.py b/src/prime_rl/trainer/optim/__init__.py
index 5988996e969..281f6580021 100644
--- a/src/prime_rl/trainer/optim/__init__.py
+++ b/src/prime_rl/trainer/optim/__init__.py
@@ -3,23 +3,32 @@
 from dion import Muon
 from torch import nn
 from torch.distributed.device_mesh import DeviceMesh
-from torch.optim import AdamW, Optimizer
+from torch.optim import Optimizer
 
 from prime_rl.configs.trainer import OffloadMode, OptimizerConfig
 from prime_rl.trainer.models.fusions import get_model_packed_parameters
+from prime_rl.trainer.optim.base import ChunkedUpdates, OptimizerLike
 from prime_rl.trainer.optim.base import OffloadOptimizer as OffloadOptimizer
-from prime_rl.trainer.optim.base import OptimizerLike
 from prime_rl.trainer.optim.offload import (
     FullCPUOffloadOptimizer,
     GradientOffloadManager,
     _create_cpu_master_weights,
 )
-from prime_rl.trainer.optim.sign_sgd import SignSGD
+from prime_rl.trainer.optim.sign_sgd import SignSGD as _SignSGD
 from prime_rl.trainer.optim.state_offload import CPUOffloadOptimizer
 from prime_rl.trainer.parallel_dims import ParallelDims
 from prime_rl.utils.logger import get_logger
 
 
+# The optimizers on the GPU take an ``update_hook``
+class AdamW(ChunkedUpdates, torch.optim.AdamW):
+    pass
+
+
+class SignSGD(ChunkedUpdates, _SignSGD):
+    pass
+
+
 def _warmup_muon_mesh(mesh: DeviceMesh) -> None:
     """Establish NCCL peer connections before Muon's first optimizer step.
 
diff --git a/src/prime_rl/trainer/optim/base.py b/src/prime_rl/trainer/optim/base.py
index bbdf364f089..a85a1e3a2dc 100644
--- a/src/prime_rl/trainer/optim/base.py
+++ b/src/prime_rl/trainer/optim/base.py
@@ -1,8 +1,50 @@
 from abc import ABC, abstractmethod
+from collections.abc import Callable
 from typing import TypeAlias
 
+import torch
+from torch.distributed.tensor import DTensor
 from torch.optim import Optimizer
 
+# Called right before the optimizer updates some parameters, on the stream of the update; the callable it returns is
+# called right after the update
+UpdateHook: TypeAlias = Callable[[list[torch.Tensor]], Callable[[], None]]
+
+
+class ChunkedUpdates:
+    """Optimizer mixin: with ``update_hook`` set, ``step`` updates the parameters a chunk of up to ``CHUNK_NUMEL``
+    local elements at a time and calls the hook around each chunk's update."""
+
+    CHUNK_NUMEL = 1 << 28
+    update_hook: UpdateHook | None = None
+
+    def step(self, closure=None):
+        if self.update_hook is None:
+            return super().step(closure)
+        assert closure is None
+        groups = self.param_groups
+        try:
+            for group in groups:
+                chunk, numel = [], 0
+                for param in group["params"]:
+                    if param.grad is None:
+                        continue
+                    chunk.append(param)
+                    numel += (param.to_local() if isinstance(param, DTensor) else param).numel()
+                    if numel >= self.CHUNK_NUMEL:
+                        self._step_params(group, chunk)
+                        chunk, numel = [], 0
+                if chunk:
+                    self._step_params(group, chunk)
+        finally:
+            self.param_groups = groups
+
+    def _step_params(self, group: dict, params: list[torch.Tensor]) -> None:
+        self.param_groups = [{**group, "params": params}]
+        finish = self.update_hook(params)
+        super().step()
+        finish()
+
 
 class OffloadOptimizer(ABC):
     @property
diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py
index 4f89e0b26b1..694ff4d837d 100644
--- a/src/prime_rl/trainer/rl/train.py
+++ b/src/prime_rl/trainer/rl/train.py
@@ -16,7 +16,6 @@
 from torch.profiler import profile, ProfilerActivity, record_function
 from prime_rl.trainer.ckpt import Progress, setup_ckpt_manager
 from prime_rl.trainer.optim import setup_optimizer
-from prime_rl.trainer.optim.offload import FullCPUOffloadOptimizer
 from prime_rl.trainer.scheduler import setup_scheduler
 from prime_rl.configs.trainer import TrainerConfig
 from prime_rl.trainer.rl.data import DataLoader, FakeDataLoader, prepare_micro_batch
@@ -194,7 +193,7 @@ def train(config: TrainerConfig):
             config.model.lora,
         )
         logger.debug(f"Initialized weight broadcast in {format_time(time.perf_counter() - t0)}")
-        if isinstance(weight_sender, SparseFileSystemWeightSender) and isinstance(optimizer, FullCPUOffloadOptimizer):
+        if isinstance(weight_sender, SparseFileSystemWeightSender):
             weight_sender.diff_in_optimizer(optimizer)
 
     if parallel_dims.cp_enabled:
diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py
index 178ab0d9f68..4008e381854 100644
--- a/src/prime_rl/trainer/sft/train.py
+++ b/src/prime_rl/trainer/sft/train.py
@@ -473,11 +473,7 @@ def is_online_eval_step(step: int) -> bool:
                 scaled_loss = local_loss_sum / grad_accum_steps
 
             with maybe_record_function("backward"):
-                final_backward = micro_step == grad_accum_steps - 1
-                if final_backward and weight_sender is not None:
-                    # The offloaded optimizer step overlaps the final backward and updates the weights
-                    weight_sender.wait_for_weight_reads()
-                begin_backward(gradient_manager, final_backward=final_backward)
+                begin_backward(gradient_manager, final_backward=micro_step == grad_accum_steps - 1)
                 scaled_loss.backward()
                 finish_backward(gradient_manager)
 
@@ -687,7 +683,6 @@ def is_online_eval_step(step: int) -> bool:
     if weight_sender is not None and not paused:
         logger.info("Broadcasting final weights")
         weight_sender.broadcast(model, step=progress.step)
-        weight_sender.close()
 
     if gradient_manager is not None:
         gradient_manager.close()
diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 1e8a24247c9..07f9811a6fc 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -1,5 +1,4 @@
 import json
-import mmap
 import queue
 import threading
 import time
@@ -7,15 +6,14 @@
 from collections import defaultdict
 from collections.abc import Callable
 from concurrent.futures import Future, ThreadPoolExecutor
+from functools import partial
 from math import prod
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-import numpy as np
 import torch
 import torch.distributed as dist
 import torch.nn as nn
-from huggingface_hub import snapshot_download
 from torch.distributed.tensor import DTensor
 from torch.distributed.tensor._utils import compute_local_shape_and_global_offset
 
@@ -30,18 +28,13 @@
     AnchorWriter,
     anchor_complete,
     as_bits,
-    bits_dtype,
-    checkpoint_index,
     compile_fp8_quantizer,
     encode_gaps,
-    page_lock,
     quantize_fp8_blocks,
     rank_part,
     rank_parts_marker,
-    read_checkpoint_range,
     read_engine_layout,
     revalidate,
-    write_anchor_shards,
     write_file,
     write_rank_file,
     zigzag,
@@ -49,46 +42,47 @@
 from prime_rl.utils.weights import convert_state_dict_to_hf, resolve_fqn, resolve_wire_dtype
 
 if TYPE_CHECKING:
-    from prime_rl.trainer.optim.offload import FullCPUOffloadOptimizer
+    from prime_rl.trainer.optim import OptimizerLike
 
 
 class SparseDeltaEncoder:
     """Rank-local deltas of the trainer's weight shards, in HF coordinates.
 
-    Each rank diffs its own shards against the wire-dtype bits of the previous broadcast (the snapshot), so no
-    weights are gathered. The prime->HF conversion is replayed once on meta tensors: every HF tensor must be a
-    contiguous view of one trainer tensor, which maps a changed trainer element to its HF name and flat index.
-
-    ``encode`` runs on a background thread and a side CUDA stream while the next step trains, without collectives.
-    The snapshot starts as the boot checkpoint, read back from storage, and lives on the GPU if it fits next to the
-    first step's peak memory, else in the host buffer.
-
-    Tensors the engines hold in fp8 are quantized with vLLM's own quantizer (old and new weights) and diffed as fp8
-    bytes; their changed block scales travel as ``"{hf_name}_scale_inv"`` tensors. fp8 tensors whose 128-row blocks
-    straddle ranks are gathered.
-
-    With a fully offloaded optimizer (``bind_optimizer``), the tensors stored in its parameters are diffed in its
-    step (``optimizer_hook``): each shard still holds the previous broadcast's weights when the step writes the new
-    ones, so they need no snapshot. ``encode`` then only checks that they still hold what the step wrote.
+    The deltas of the optimizer's parameters come from its step (``bind``): right before the step updates a chunk of
+    parameters, their shards still hold the previous broadcast's weights. ``update_hook`` (optimizers on the GPU)
+    copies the chunk's old wire bits, ``optimizer_hook`` (``FullCPUOffloadOptimizer``) receives the new values next
+    to the old ones; a diff thread compares the old and new wire bits and streams the changes out as the step goes.
+    No copy of the weights is kept. The other tensors (the residual: buffers, frozen parameters, state-dict copies,
+    fp8 tensors whose 128-row blocks straddle ranks) are diffed in ``encode`` against a copy of their previous wire
+    bits, which also checks that the optimizer's tensors still hold what its step wrote.
+
+    The prime->HF conversion is replayed once on meta tensors: every HF tensor must be a contiguous view of one trainer
+    tensor, which maps a changed trainer element to its HF name and flat index. Tensors the engines hold in fp8 are
+    quantized with vLLM's own quantizer (old and new weights) and diffed as fp8 bytes; their changed block scales
+    travel as ``"{hf_name}_scale_inv"`` tensors.
     """
 
-    # Elements per diff chunk and changes per written part. Smaller in the optimizer step, which overlaps the final
-    # backward.
+    # Elements per diff chunk and changes per written part: in ``encode``, and in the optimizer step (smaller, as it
+    # overlaps training)
     CHUNK = 1 << 26
     PART = 1 << 26
     HOOK_CHUNK = 1 << 24
     HOOK_PART = 1 << 24
-    # GPU slots of the optimizer step's diff thread, each twice the largest parameter shard
+    # GPU slots of ``optimizer_hook``, each twice the largest parameter shard; bytes of old values that
+    # ``update_hook`` holds for the diff thread (at least one update's)
     HOOK_SLOTS = 2
+    UPDATE_BYTES = 1 << 32
 
-    def __init__(self, model: nn.Module, snapshot: str):
+    def __init__(self, model: nn.Module):
         keep_in_fp32 = getattr(model, "keep_in_fp32_for_weight_transfer", None)
         self.wire_dtypes: dict[str, torch.dtype] = {}
+        self.key_fqn: dict[str, str] = {}
         meta: dict[str, torch.Tensor] = {}
         for key, value in model.state_dict().items():
             dtype = resolve_wire_dtype(keep_in_fp32, key, torch.bfloat16) if isinstance(value, DTensor) else value.dtype
             self.wire_dtypes[key] = dtype
-            meta[resolve_fqn(model, key)] = torch.empty(value.shape, dtype=dtype, device="meta")
+            self.key_fqn[key] = resolve_fqn(model, key)
+            meta[self.key_fqn[key]] = torch.empty(value.shape, dtype=dtype, device="meta")
         owner = {id(tensor): fqn for fqn, tensor in meta.items()}
 
         segments = defaultdict(list)
@@ -113,28 +107,21 @@ def __init__(self, model: nn.Module, snapshot: str):
         self.gather: set[str] = set()
         self.segment_starts: dict[str, torch.Tensor] = {}
         self.shard_segments: dict[str, tuple[list[str], list[int], list[int], list[int]]] = {}
-        self.snapshot_placement = snapshot
-        # key -> (local numel, global flat offset) of this rank's shard
+        # Residual keys -> (local numel, global flat offset) of this rank's shard, and their previous wire bits
         self.shards: dict[str, tuple[int, int]] = {}
-        self.key_fqn: dict[str, str] = {}
-        # key -> previous wire bits: GPU tensors, or views into the host buffer
-        self.snapshot: dict[str, torch.Tensor] = {}
-        self.on_gpu: bool | None = None
-        self.host_buffer: torch.Tensor | None = None
-        self.write_staging: torch.Tensor | None = None
-        # The boot checkpoint's ``checkpoint_index``, and keys whose v0 is still to read from it
-        self.boot_files: dict[str, tuple] = {}
-        self.boot_pending: set[str] = set()
-        self.host: dict[str, torch.Tensor] = {}
-        self.initialized = False
-        # Optimizer-step deltas: parameter storage -> (key, fqn, numel, global flat offset, (shape, stride, storage
-        # offset) of the view); the step's parts, diff futures and checksums of the written wire bits
+        self.previous: dict[str, torch.Tensor] = {}
+        # Optimizer parameter storage -> (key, fqn, numel, global flat offset, (shape, stride, storage offset) of the
+        # view) of the trainer tensors stored in it, and the checksum of each one's last written wire bits
         self.param_keys: dict[int, list[tuple[str, str, int, int, tuple]]] = defaultdict(list)
-        self.optimizer_keys: set[str] = set()
+        self.checksums: dict[str, torch.Tensor] = {}
+        self.residual: set[str] = set()
+        self.differ: ThreadPoolExecutor | None = None
+        # The armed step's parts and diffs, and its anchor
+        self.armed = False
         self.step_parts: list[tuple] = []
         self.diffs: list[Future] = []
-        self.expected: dict[str, torch.Tensor] = {}
         self.anchor: AnchorWriter | None = None
+        self.write_staging: torch.Tensor | None = None
         self.device = torch.cuda.current_device()
         self.stream = torch.cuda.Stream()
 
@@ -207,20 +194,20 @@ def _find_gathered(self, model: nn.Module, is_master: bool) -> set[str]:
         return {key for key, flag in zip(keys, flags.tolist()) if flag}
 
     @torch.no_grad()
-    def local_shards(self, model: nn.Module, is_master: bool) -> list[tuple]:
+    def local_shards(self, model: nn.Module, is_master: bool, gather: bool = False) -> list[tuple]:
         """This rank's ``(key, fqn, shard, global flat offset, gathered)``, taken on the training thread (FSDP swaps
-        module parameters during forward, but the sharded storage stays put). ``gathered`` is ``(previous,
-        current)`` full bf16 tensors of fp8 tensors whose blocks straddle ranks, all-gathered here because
-        ``encode`` runs no collectives."""
+        module parameters during forward, but the sharded storage stays put). With ``gather``, ``gathered`` is
+        ``(previous, current)`` full bf16 tensors of fp8 tensors whose blocks straddle ranks, all-gathered here
+        because ``encode`` runs no collectives."""
         shards = []
         for key, value in model.state_dict().items():
-            fqn = resolve_fqn(model, key)
+            fqn = self.key_fqn[key]
             if fqn not in self.segments:
                 continue
             shard = self._local_shard(key, value, is_master)
             empty = shard is None or shard[0].numel() == 0
             gathered = None
-            if key in self.gather and self.initialized:
+            if gather and key in self.gather:
                 if empty:
                     value.full_tensor()  # every rank joins the all-gathers
                     value.full_tensor()
@@ -231,79 +218,21 @@ def local_shards(self, model: nn.Module, is_master: bool) -> list[tuple]:
         return shards
 
     def _gather_previous(self, key: str, value: DTensor) -> torch.Tensor:
-        previous = self.snapshot[key].to(f"cuda:{self.device}").view(torch.bfloat16)
+        previous = self.previous[key].view(torch.bfloat16)
         local_shape = (previous.numel() // prod(value.shape[1:]), *value.shape[1:])
         return DTensor.from_local(
             previous.view(local_shape), value.device_mesh, value.placements, shape=value.shape, stride=value.stride()
         ).full_tensor()
 
-    def _bits(self, key: str) -> torch.dtype:
-        return bits_dtype(str(self.wire_dtypes[key]).removeprefix("torch."))
-
-    def allocate_host(self) -> None:
-        """One host buffer for all shards (the CPU snapshot, or anchor copies of a GPU snapshot). Faulted in before
-        training: on a node short of free memory, faulting it in during an anchor copy waits minutes for the kernel
-        to reclaim page cache while the trainer waits for that copy."""
-        start = time.perf_counter()
-        sizes = {key: -(-numel * self.wire_dtypes[key].itemsize // 8) * 8 for key, (numel, _) in self.shards.items()}
-        flags = mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS | mmap.MAP_POPULATE
-        # mmap needs a non-empty mapping: a rank may hold no shards
-        self.host_buffer = torch.from_numpy(np.frombuffer(mmap.mmap(-1, max(sum(sizes.values()), 8), flags), np.uint8))
-        offset = 0
-        for key, (numel, _) in self.shards.items():
-            self.host[key] = self.host_buffer[offset : offset + numel * self.wire_dtypes[key].itemsize].view(
-                self._bits(key)
-            )
-            offset += sizes[key]
-        get_logger().info(
-            f"Allocated the {self.host_buffer.numel() / 2**30:.1f} GiB sparse weight host buffer in "
-            f"{time.perf_counter() - start:.1f}s"
-        )
-
-    def _page_lock(self) -> None:
-        """Only while no training step runs: registering a large buffer stalls the trainer's CUDA calls."""
-        start = time.perf_counter()
-        page_lock(self.host_buffer)
-        get_logger().info(
-            f"Page-locked the {self.host_buffer.numel() / 2**30:.1f} GiB sparse weight host buffer in "
-            f"{time.perf_counter() - start:.1f}s"
-        )
-
-    def _snapshot_bytes(self) -> int:
-        return sum(numel * self.wire_dtypes[key].itemsize for key, (numel, _) in self.shards.items())
-
-    def _place(self, peak_reserved: int) -> None:
-        """Keep the snapshot on the GPU if it fits next to the training step's peak, else on the host."""
-        num_bytes = self._snapshot_bytes()
-        total = torch.cuda.get_device_properties(self.device).total_memory
-        headroom = total - peak_reserved - num_bytes
-        self.on_gpu = self.snapshot_placement == "auto" and headroom >= max(4 * 2**30, total // 20)
-        get_logger().info(
-            f"Sparse weight snapshot ({num_bytes / 2**30:.1f} GiB) on the {'GPU' if self.on_gpu else 'host'} "
-            f"(training peak {peak_reserved / 2**30:.1f} GiB of {total / 2**30:.1f} GiB, "
-            f"snapshot={self.snapshot_placement})"
-        )
+    def _bits(self, key: str, shard: torch.Tensor) -> torch.Tensor:
+        """A shard's flat wire bits (a copy)."""
+        return as_bits(shard.reshape(-1).to(self.wire_dtypes[key], copy=True))
 
-    def track(self, shards: list[tuple]) -> None:
-        """Record which elements of each tensor this rank's shards hold."""
-        self.shards = {key: (shard.numel(), offset) for key, _, shard, offset, _ in shards}
-        self.key_fqn.update({key: fqn for key, fqn, *_ in shards})
-
-    def copy_snapshot_to_host(self, give_up: Callable[[], bool]) -> bool:
-        """Copy a GPU snapshot into the host buffer for an anchor; False if ``give_up()`` turned true first. Callers
-        wait for earlier readers of the host buffer."""
-        with torch.cuda.stream(self.stream):
-            for key, bits in self.snapshot.items():
-                if give_up():
-                    return False
-                self._to_host(bits, self.host[key])
-        return True
-
-    def _slices(self, key: str, numel: int, offset: int) -> list[tuple[str, list[int], int, int, int]]:
+    def _slices(self, fqn: str, numel: int, offset: int) -> list[tuple[str, list[int], int, int, int]]:
         """``(hf_name, shape, flat start in the HF tensor, begin, end in the shard)`` of the HF tensors in a
         shard of ``numel`` elements at global flat ``offset``."""
         slices = []
-        for start, name in self.segments[self.key_fqn[key]]:
+        for start, name in self.segments[fqn]:
             shape = self.hf_tensors[name][0]
             begin, end = max(start, offset), min(start + prod(shape), offset + numel)
             if begin < end:
@@ -311,64 +240,46 @@ def _slices(self, key: str, numel: int, offset: int) -> list[tuple[str, list[int
         return slices
 
     def anchor_parts(self, source: dict[str, torch.Tensor]) -> list[tuple[str, list[int], int, torch.Tensor]]:
-        """This rank's slices of a full HF checkpoint, from ``source`` (key -> flat wire bits of its shard)."""
+        """This rank's ``(hf_name, shape, flat start, flat values)`` slices of a full HF checkpoint, from ``source``
+        (key -> (global flat offset, flat wire bits of its shard))."""
         parts = []
-        for key, bits in source.items():
+        for key, (offset, bits) in source.items():
             values = bits.view(self.wire_dtypes[key])
-            for name, shape, start, begin, end in self._slices(key, *self.shards[key]):
+            for name, shape, start, begin, end in self._slices(self.key_fqn[key], len(bits), offset):
                 parts.append((name, shape, start, values[begin:end]))
         return parts
 
     @torch.no_grad()
-    def initialize(self, model: nn.Module, is_master: bool, peak_reserved: int, boot_dir: Path) -> None:
-        """At the first delta (v1), on the training thread: place and allocate the snapshot. It starts as v0, the
-        checkpoint in ``boot_dir``, read as the encode goes; tensors gathered for fp8 are read here, before
-        ``local_shards`` gathers them."""
-        self.track([shard for shard in self.local_shards(model, is_master) if shard[0] not in self.optimizer_keys])
-        self._place(peak_reserved)
-        self.staging = torch.empty(1 << 28, dtype=torch.uint8, pin_memory=True)
-        if self.on_gpu:
-            device = f"cuda:{self.device}"
-            self.snapshot = {
-                key: torch.empty(numel, dtype=self._bits(key), device=device) for key, (numel, _) in self.shards.items()
-            }
-        else:
-            self._page_lock()
-            self.snapshot = dict(self.host)
-        self.boot_files = checkpoint_index(boot_dir)
-        self.boot_pending = set(self.shards) - self.gather
-        # The snapshot was just allocated on this stream, whose earlier work may still use that memory
-        self.stream.wait_stream(torch.cuda.current_stream())
-        with torch.cuda.stream(self.stream):
-            for key in self.gather & set(self.shards):
-                self._read_boot(key)
-            self.stream.synchronize()
-        self.initialized = True
-
-    def _read_boot(self, key: str) -> None:
-        """Fill the snapshot of ``key`` with this rank's slice of the boot checkpoint, on the current stream."""
-        out = self.snapshot[key].view(self.wire_dtypes[key])
-        staging = self.staging.numpy()
-        for name, _, hf_start, begin, end in self._slices(key, *self.shards[key]):
-            dtype = self.boot_files[name][1]
-            step = self.staging.numel() // dtype.itemsize
-            for chunk in range(begin, end, step):
-                count = min(step, end - chunk)
-                first = hf_start + chunk - begin
-                read_checkpoint_range(self.boot_files[name], first, first + count, staging)
-                out[chunk : chunk + count].copy_(self.staging[: count * dtype.itemsize].view(dtype), non_blocking=True)
-                torch.cuda.current_stream().synchronize()
+    def write_full_checkpoint(self, model: nn.Module, is_master: bool, step_dir: Path, rank: int, world_size: int):
+        """Write this rank's slices of a full HF checkpoint of the live weights, one tensor at a time."""
+        shards = [
+            (key, shard, self._slices(fqn, shard.numel(), offset))
+            for key, fqn, shard, offset, _ in self.local_shards(model, is_master)
+        ]
+        writer = AnchorWriter(
+            [
+                (name, shape, start, end - begin, self.wire_dtypes[key])
+                for key, _, slices in shards
+                for name, shape, start, begin, end in slices
+            ],
+            step_dir,
+            rank,
+        )
+        for key, shard, slices in shards:
+            values = shard.reshape(-1).to(self.wire_dtypes[key])
+            parts = [(name, shape, start, values[begin:end]) for name, shape, start, begin, end in slices]
+            writer.write_parts(parts, self.stream_to_host)
+        writer.close(world_size)
 
     @torch.no_grad()
     def encode(
         self, shards: list[tuple], ready: torch.cuda.Event, weights_read: threading.Event, emit: Callable
     ) -> int:
-        """Diff the shards against the snapshot and advance it, on the side stream after ``ready``. Hands host parts
-        ``(entries, gaps, value bytes, checksums)`` to ``emit`` as they are encoded and sets ``weights_read`` once
-        done reading the weights. Returns the number of changed fp8 block scales."""
+        """After ``wait_for_diffs``, on the side stream after ``ready``: diff the residual shards against their
+        previous bits and advance them, and check the optimizer's tensors. Hands host parts ``(entries, gaps, value
+        bytes, checksums)`` to ``emit`` as they are encoded and sets ``weights_read`` once done reading the weights.
+        Returns the number of changed fp8 block scales."""
         torch.cuda.set_device(self.device)
-        if not self.optimizer_keys:
-            self.scale_changes = 0
         try:
             with torch.cuda.stream(self.stream):
                 self.stream.wait_event(ready)
@@ -378,6 +289,15 @@ def encode(
         finally:
             weights_read.set()
 
+    def wait_for_diffs(self) -> None:
+        for diff in self.diffs:
+            diff.result()
+        self.diffs.clear()
+        get_logger().debug(
+            f"Sparse weight delta in the optimizer step: waited {self.hook_wait:.2f}s for earlier diffs, diff thread "
+            f"busy {self.diff_busy:.2f}s"
+        )
+
     def stream_to_host(self, tensor: torch.Tensor):
         """A tensor's bytes on the host; for a GPU tensor, in pinned staging chunks, each to be consumed before the
         next is taken."""
@@ -388,6 +308,7 @@ def stream_to_host(self, tensor: torch.Tensor):
         if self.write_staging is None:
             self.write_staging = torch.empty(1 << 28, dtype=torch.uint8, pin_memory=True)
             self.write_stream = torch.cuda.Stream(self.device)
+        self.write_stream.wait_stream(torch.cuda.current_stream())
         with torch.cuda.stream(self.write_stream):
             for begin in range(0, data.numel(), self.write_staging.numel()):
                 chunk = self.write_staging[: min(self.write_staging.numel(), data.numel() - begin)]
@@ -439,30 +360,33 @@ def _segment_sums(bits: torch.Tensor, begins: list[int], ends: list[int], firsts
                 j += 1
             rows = bits[begins[i] : begins[i] + (j - i) * length].view(j - i, length)
             positions = torch.arange(1, length + 1, device=bits.device)
-            sums.append(
-                (rows * positions).sum(1)
-                + rows.sum(1, dtype=torch.int64) * torch.tensor(firsts[i:j], device=bits.device)
-            )
+            # Through pinned memory: a copy from pageable memory waits for the stream, holding the GIL
+            first = torch.tensor(firsts[i:j]).pin_memory().to(bits.device, non_blocking=True)
+            sums.append((rows * positions).sum(1) + rows.sum(1, dtype=torch.int64) * first)
             i = j
         return torch.cat(sums)
 
-    def _part(self, cache_key, segments: list[tuple[int, str]], position, values, checksums) -> tuple:
-        """A part entry for the changes at global flat ``position`` of a trainer tensor made of ``segments``:
-        ``(names, counts, gaps, value bytes, checksum names, checksums)``."""
+    def _part(self, cache_key, names, begins, shifts, position, values, checksums) -> tuple:
+        """A part entry for changes at flat ``position``s in segments that start at ``begins`` (sorted), where the HF
+        flat index is the position plus the segment's shift: ``(names, counts, gaps, value bytes, checksum names,
+        checksums)``."""
         if cache_key not in self.segment_starts:
-            self.segment_starts[cache_key] = torch.tensor([start for start, _ in segments], device=position.device)
-        starts = self.segment_starts[cache_key]
+            device = position.device
+            self.segment_starts[cache_key] = (torch.tensor(begins, device=device), torch.tensor(shifts, device=device))
+        starts, shift = self.segment_starts[cache_key]
         segment = torch.searchsorted(starts, position, right=True) - 1
         entry_starts = torch.ones_like(segment, dtype=torch.bool)
         entry_starts[1:] = segment[1:] != segment[:-1]
-        gaps = encode_gaps(position - starts[segment], entry_starts)
-        counts = torch.bincount(segment, minlength=len(starts))
-        return [name for _, name in segments], counts, gaps, values.view(torch.uint8), *checksums
+        gaps = encode_gaps(position + shift[segment], entry_starts)
+        # Not bincount, which syncs to find the largest value
+        counts = torch.zeros_like(starts).index_add_(0, segment, torch.ones_like(segment))
+        return names, counts, gaps, values.view(torch.uint8), *checksums
 
     def _scale_part(self, fqn: str, first_block: int, previous: torch.Tensor, current: torch.Tensor) -> tuple:
         """The part entry of an fp8 tensor's changed block scales: block rows from ``first_block`` on."""
         scale_cols = current.shape[1]
         delta = as_bits(current).reshape(-1) - as_bits(previous).reshape(-1)
+        torch.cuda.current_stream().synchronize()
         index = delta.nonzero().squeeze(1)
         self.scale_changes += len(index)
         offset = first_block * scale_cols
@@ -470,17 +394,24 @@ def _scale_part(self, fqn: str, first_block: int, previous: torch.Tensor, curren
         sums = self._segment_sums(
             as_bits(current).reshape(-1), begins, ends, [begin + shift for begin, shift in zip(begins, shifts)]
         )
+        segments = self.scale_segments[fqn]
         return self._part(
-            (fqn, "scales"), self.scale_segments[fqn], index + offset, zigzag(delta[index]), (names, sums)
+            (fqn, "scales"),
+            [name for _, name in segments],
+            [start for start, _ in segments],
+            [-start for start, _ in segments],
+            index + offset,
+            zigzag(delta[index]),
+            (names, sums),
         )
 
     def _diff(self, shards: list[tuple], emit: Callable) -> None:
-        """Diff ``shards`` against the snapshot on the current stream and advance it, handing a host part to ``emit``
-        every ``PART`` changes. Shards diffed in the optimizer step are only checked against what it wrote; its
-        remaining parts go first."""
+        """Hand the optimizer step's remaining parts to ``emit``, then diff the residual ``shards`` on the current
+        stream, handing a host part to ``emit`` every ``PART`` changes, and check that the optimizer's tensors hold
+        what its step wrote."""
         pending, self.step_parts = self.step_parts, []
         num_pending = sum(len(part[2]) for part in pending)
-        checks: list[tuple[str, torch.Tensor, torch.Tensor]] = []
+        checks: list[tuple[str, torch.Tensor]] = []
 
         def flush() -> None:
             emit(self._host_part(pending))
@@ -488,15 +419,10 @@ def flush() -> None:
 
         for key, fqn, shard, offset, gathered in shards:
             shard = shard.reshape(-1)
-            if key in self.optimizer_keys:
-                if key not in self.expected:
-                    raise RuntimeError(f"{key} was not written by the optimizer step before this broadcast")
-                checks.append((key, self.expected.pop(key), self._checksum(as_bits(shard.to(self.wire_dtypes[key])))))
+            if key in self.checksums:
+                checks.append((key, self._checksum(key, offset, as_bits(shard.to(self.wire_dtypes[key])))))
                 continue
-            if key in self.boot_pending:
-                self._read_boot(key)
-                self.boot_pending.discard(key)
-            parts = self._diff_shard(key, fqn, shard, offset, gathered, self.snapshot[key], self.snapshot[key])
+            parts = self._diff_shard(key, fqn, shard, offset, gathered, self.previous[key], self.previous[key])
             pending += parts
             num_pending += sum(len(part[2]) for part in parts)
             if num_pending >= self.PART:
@@ -505,22 +431,18 @@ def flush() -> None:
         if pending:
             flush()
         if checks:
-            keys, expected, live = zip(*checks)
-            changed = [key for key, same in zip(keys, torch.cat(expected).eq(torch.cat(live)).tolist()) if not same]
+            keys, live = zip(*checks)
+            expected = torch.cat([self.checksums[key] for key in keys])
+            changed = [key for key, same in zip(keys, expected.eq(torch.cat(live)).tolist()) if not same]
             if changed:
                 raise RuntimeError(f"Trainer weights changed outside the optimizer step, not sent: {changed}")
 
-    def wait_for_diffs(self) -> None:
-        for diff in self.diffs:
-            diff.result()
-        self.diffs.clear()
-
     def _host_part(self, pending: list[tuple], staging: torch.Tensor | None = None) -> tuple:
         """One part ``(entries, gaps, value bytes, checksums)`` on the host from GPU part entries."""
         names = [name for part in pending for name in part[0]]
-        counts = torch.cat([part[1] for part in pending]).tolist()
         gaps = self._to_host(torch.cat([part[2] for part in pending]), staging=staging)
         values = self._to_host(torch.cat([part[3] for part in pending]), staging=staging)
+        counts = torch.cat([part[1] for part in pending]).tolist()
         checksum_names = [name for part in pending for name in part[4]]
         checksums = dict(zip(checksum_names, torch.cat([part[5] for part in pending]).tolist()))
         return [(name, count) for name, count in zip(names, counts) if count], gaps, values, checksums
@@ -539,65 +461,50 @@ def _diff_shard(
         """The part entries of a flat shard's changes against ``previous`` (its wire bits), on the current stream, in
         chunks of ``chunk`` elements (fp8: whole 128-row blocks, plus their block scales). Writes the new bits to
         ``out`` unless None."""
-        device = torch.device(f"cuda:{self.device}")
         wire = self.wire_dtypes[key]
         segment_names, begins, ends, shifts = self._segments(key, self.segments[fqn], offset, shard.numel())
-        sums = torch.zeros(len(segment_names), dtype=torch.int64, device=device)
+        if fqn not in self.fp8:
+            current = as_bits(shard.to(wire))
+            part = self._diff_flat(key, segment_names, begins, ends, shifts, current, previous, chunk)
+            if out is not None:
+                out.copy_(current)
+            return [part]
+        sums = torch.zeros(len(segment_names), dtype=torch.int64, device=shard.device)
         changed, changed_values, scales = [], [], []
-        if fqn in self.fp8:
-            ue8m0, num_cols = self.fp8[fqn]
-            if gathered is not None:
-                chunks = [(0, shard.numel())]
-            else:
-                rows = max(FP8_BLOCK[0], chunk // num_cols // FP8_BLOCK[0] * FP8_BLOCK[0])
-                chunks = [(start, start + rows * num_cols) for start in range(0, shard.numel(), rows * num_cols)]
+        ue8m0, num_cols = self.fp8[fqn]
+        if gathered is not None:
+            chunks = [(0, shard.numel())]
         else:
-            chunks = [(start, start + chunk) for start in range(0, shard.numel(), chunk)]
+            rows = max(FP8_BLOCK[0], chunk // num_cols // FP8_BLOCK[0] * FP8_BLOCK[0])
+            chunks = [(start, start + rows * num_cols) for start in range(0, shard.numel(), rows * num_cols)]
         for start, end in chunks:
             end = min(end, shard.numel())
-            current = as_bits(shard[start:end].to(device, non_blocking=True).to(wire))
-            previous_bits = previous[start:end].to(device, non_blocking=True)
-            if fqn in self.fp8:
-                if gathered is not None:
-                    previous_weights, current_weights = gathered
-                else:
-                    previous_weights, current_weights = (
-                        previous_bits.view(torch.bfloat16),
-                        current.view(torch.bfloat16),
-                    )
-                (previous_fp8, previous_scales), (current_fp8, current_scales) = (
-                    quantize_fp8_blocks(weights.reshape(-1, num_cols), ue8m0, self.compiled_fp8)
-                    for weights in (previous_weights, current_weights)
-                )
-                scales.append((previous_scales, current_scales))
-                if gathered is not None:
-                    previous_fp8 = previous_fp8.view(-1)[offset : offset + shard.numel()]
-                    current_fp8 = current_fp8.view(-1)[offset : offset + shard.numel()]
-                delta = current_fp8.view(-1) - previous_fp8.view(-1)
-                new_bits = current_fp8.view(-1)
+            current = as_bits(shard[start:end].to(wire))
+            previous_bits = previous[start:end]
+            if gathered is not None:
+                previous_weights, current_weights = gathered
             else:
-                delta = current - previous_bits
-                new_bits = current
+                previous_weights, current_weights = previous_bits.view(torch.bfloat16), current.view(torch.bfloat16)
+            (previous_fp8, previous_scales), (current_fp8, current_scales) = (
+                quantize_fp8_blocks(weights.reshape(-1, num_cols), ue8m0, self.compiled_fp8)
+                for weights in (previous_weights, current_weights)
+            )
+            scales.append((previous_scales, current_scales))
+            if gathered is not None:
+                previous_fp8 = previous_fp8.view(-1)[offset : offset + shard.numel()]
+                current_fp8 = current_fp8.view(-1)[offset : offset + shard.numel()]
+            delta = current_fp8.view(-1) - previous_fp8.view(-1)
             if out is not None:
                 out[start:end].copy_(current, non_blocking=True)
+            # nonzero syncs to size its output, holding the GIL: wait for the stream first, without it
+            torch.cuda.current_stream().synchronize()
             index = delta.nonzero().squeeze(1)
             changed.append(index + start)
             changed_values.append(zigzag(delta[index]))
-            first, last = bisect_right(ends, start), bisect_left(begins, end)
-            if first < last:
-                sums[first:last] += self._segment_sums(
-                    new_bits,
-                    [max(begin, start) - start for begin in begins[first:last]],
-                    [min(stop, end) - start for stop in ends[first:last]],
-                    [max(begin, start) + shift for begin, shift in zip(begins[first:last], shifts[first:last])],
-                )
+            sums += self._window_sums(current_fp8.view(-1), start, end, begins, ends, shifts)
         parts = [
             self._part(
-                fqn,
-                self.segments[fqn],
-                torch.cat(changed) + offset,
-                torch.cat(changed_values),
-                (segment_names, sums),
+                key, segment_names, begins, shifts, torch.cat(changed), torch.cat(changed_values), (segment_names, sums)
             )
         ]
         if scales:
@@ -615,74 +522,166 @@ def _diff_shard(
                 parts.append(self._scale_part(fqn, first_block, previous_scales, current_scales))
         return parts
 
-    def _checksum(self, bits: torch.Tensor) -> torch.Tensor:
-        """``bits_checksum`` of a flat shard, in chunks of ``HOOK_CHUNK`` elements."""
-        n = bits.numel()
-        sums = [
-            self._segment_sums(bits, [start], [min(start + self.HOOK_CHUNK, n)], [start])
-            for start in range(0, n, self.HOOK_CHUNK)
-        ]
-        return torch.cat(sums).sum(0, keepdim=True) if sums else torch.zeros(1, dtype=torch.int64, device=bits.device)
+    def _diff_flat(self, cache_key, names, begins, ends, shifts, current, previous, chunk: int) -> tuple:
+        """The part entry of the changes between the flat bits ``previous`` and ``current``, made of segments
+        ``[begin, end)`` whose HF flat index is the position plus the shift, in chunks of ``chunk`` elements."""
+        sums = torch.zeros(len(names), dtype=torch.int64, device=current.device)
+        changed, changed_values = [], []
+        for start in range(0, current.numel(), chunk):
+            end = min(start + chunk, current.numel())
+            delta = current[start:end] - previous[start:end]
+            # nonzero syncs to size its output, holding the GIL: wait for the stream first, without it
+            torch.cuda.current_stream().synchronize()
+            index = delta.nonzero().squeeze(1)
+            changed.append(index + start)
+            changed_values.append(zigzag(delta[index]))
+            sums += self._window_sums(current[start:end], start, end, begins, ends, shifts)
+        return self._part(
+            cache_key, names, begins, shifts, torch.cat(changed), torch.cat(changed_values), (names, sums)
+        )
 
-    def bind_optimizer(self, model: nn.Module, is_master: bool, params: list[torch.Tensor]) -> None:
-        """Diff the trainer tensors stored in ``params`` (the offloaded optimizer's contiguous local shards) in the
-        optimizer step (``optimizer_hook``); the snapshot keeps only the other tensors and gathered fp8 tensors. A
-        trainer tensor may be a strided view of a parameter (e.g. one half of a packed gate/up weight)."""
+    def _window_sums(self, bits, start: int, end: int, begins, ends, shifts) -> torch.Tensor:
+        """Every segment's ``bits_checksum`` contribution from the window ``[start, end)``, whose flat bits are
+        ``bits``."""
+        sums = torch.zeros(len(begins), dtype=torch.int64, device=bits.device)
+        first, last = bisect_right(ends, start), bisect_left(begins, end)
+        if first < last:
+            sums[first:last] = self._segment_sums(
+                bits,
+                [max(begin, start) - start for begin in begins[first:last]],
+                [min(stop, end) - start for stop in ends[first:last]],
+                [max(begin, start) + shift for begin, shift in zip(begins[first:last], shifts[first:last])],
+            )
+        return sums
+
+    def _checksum(self, key: str, offset: int, bits: torch.Tensor) -> torch.Tensor:
+        """The sum of the ``bits_checksum``s of the HF tensors in a key's flat shard bits, in chunks of
+        ``HOOK_CHUNK`` elements: what the optimizer step records for the tensors it writes."""
+        _, begins, ends, shifts = self._segments(key, self.segments[self.key_fqn[key]], offset, bits.numel())
+        total = torch.zeros(1, dtype=torch.int64, device=bits.device)
+        for start in range(0, bits.numel(), self.HOOK_CHUNK):
+            end = min(start + self.HOOK_CHUNK, bits.numel())
+            total += self._window_sums(bits[start:end], start, end, begins, ends, shifts).sum()
+        return total
+
+    def bind(self, model: nn.Module, is_master: bool, params: list[torch.Tensor], full_offload: bool) -> None:
+        """Diff the trainer tensors stored in ``params`` (the optimizer's local parameter shards, contiguous) in its
+        step; the others are the residual. A trainer tensor may be a strided view of a parameter (e.g. one half of
+        a packed gate/up weight). fp8 tensors whose blocks straddle ranks need collectives and stay residual."""
         assert all(param.is_contiguous() for param in params)
         stored = {param.untyped_storage().data_ptr() for param in params}
         residual = []
-        for shard in self.local_shards(model, is_master):
-            key, fqn, tensor, offset, _ = shard
+        for key, fqn, tensor, offset, _ in self.local_shards(model, is_master):
             storage = tensor.untyped_storage().data_ptr()
             if key not in self.gather and storage in stored:
                 view = (tuple(tensor.shape), tensor.stride(), tensor.storage_offset())
                 self.param_keys[storage].append((key, fqn, tensor.numel(), offset, view))
-                self.optimizer_keys.add(key)
             else:
-                residual.append(shard)
-        self.track(residual)
-        slot_bytes = 2 * -(-max(param.nbytes for param in params) // 256) * 256
-        self.hook_slots = queue.SimpleQueue()
-        for _ in range(self.HOOK_SLOTS):
-            self.hook_slots.put(torch.empty(slot_bytes, dtype=torch.uint8, device=f"cuda:{self.device}"))
+                residual.append(key)
+        self.residual = set(residual)
+        device = f"cuda:{self.device}"
+        if full_offload:
+            slot_bytes = 2 * -(-max(param.nbytes for param in params) // 256) * 256
+            self.hook_slots = queue.SimpleQueue()
+            for _ in range(self.HOOK_SLOTS):
+                self.hook_slots.put(torch.empty(slot_bytes, dtype=torch.uint8, device=device))
+        self.update_bytes = threading.Condition()
+        self.held_bytes = 0
         self.differ = ThreadPoolExecutor(1)
         self.diff_stream = torch.cuda.Stream(self.device)
         self.diff_staging = torch.empty(1 << 28, dtype=torch.uint8, pin_memory=True)
-        num_bytes = self._snapshot_bytes()
+        self.staging = torch.empty(1 << 28, dtype=torch.uint8, pin_memory=True)
+
+    @torch.no_grad()
+    def reset(self, model: nn.Module, is_master: bool) -> None:
+        """After a full checkpoint: take the residual's previous bits and the optimizer tensors' checksums from the
+        live weights, on the training thread."""
+        self.previous, self.checksums, self.shards = {}, {}, {}
+        for key, _, shard, offset, _ in self.local_shards(model, is_master):
+            if key in self.residual:
+                self.previous[key] = self._bits(key, shard)
+                self.shards[key] = (shard.numel(), offset)
+            else:
+                self.checksums[key] = self._checksum(key, offset, as_bits(shard.reshape(-1).to(self.wire_dtypes[key])))
+        num_bytes = sum(bits.nbytes for bits in self.previous.values())
         get_logger().info(
-            f"Sparse weight deltas of {len(self.optimizer_keys)} tensors come from the optimizer step, "
-            f"{len(self.shards)} tensors ({num_bytes / 2**30:.2f} GiB) are diffed against a snapshot"
+            f"Sparse weight deltas of {len(self.checksums)} tensors come from the optimizer step, "
+            f"{len(self.previous)} tensors ({num_bytes / 2**30:.2f} GiB) are diffed against a copy"
         )
 
     def arm(self, anchor: AnchorWriter | None, emit: Callable) -> None:
         """Collect the next optimizer step's delta, handing host parts to ``emit`` as the step goes, and write its
         anchor slices if ``anchor``."""
-        self.step_parts, self.expected, self.scale_changes, self.anchor, self.emit = [], {}, 0, anchor, emit
+        self.step_parts, self.scale_changes, self.anchor, self.emit = [], 0, anchor, emit
+        self.hook_wait = self.diff_busy = 0.0
+        self.armed = True
+
+    def disarm(self) -> None:
+        self.armed = False
 
     def anchor_slices(self) -> list[tuple[str, list[int], int, int, torch.dtype]]:
         """Every ``(hf_name, shape, flat start, numel, wire dtype)`` slice of this rank's share of an anchor."""
-        shards = {key: (numel, offset) for entries in self.param_keys.values() for key, _, numel, offset, _ in entries}
-        shards.update(self.shards)
+        shards = {
+            key: (fqn, numel, offset) for entries in self.param_keys.values() for key, fqn, numel, offset, _ in entries
+        }
+        shards.update({key: (self.key_fqn[key], numel, offset) for key, (numel, offset) in self.shards.items()})
         return [
             (name, shape, start, end - begin, self.wire_dtypes[key])
-            for key, (numel, offset) in shards.items()
-            for name, shape, start, begin, end in self._slices(key, numel, offset)
+            for key, (fqn, numel, offset) in shards.items()
+            for name, shape, start, begin, end in self._slices(fqn, numel, offset)
         ]
 
-    @torch.no_grad()
-    def optimizer_hook(self, local_param: torch.Tensor, source: torch.Tensor) -> None:
-        """``FullCPUOffloadOptimizer.weight_hook``, on its H2D stream: copy the shard (still the previous
-        broadcast's weights) and its new values ``source`` (pinned host) into a GPU slot, write the new values and
-        diff the slot on the diff thread. At anchor steps, also writes the new values to the anchor from the
-        host."""
+    def _entries(self, local_param: torch.Tensor) -> list[tuple]:
+        """The ``param_keys`` entries of the trainer tensors in ``local_param``, views relative to it."""
         start = local_param.storage_offset()
-        entries = [
+        return [
             (key, fqn, numel, offset, (size, stride, at - start))
             for key, fqn, numel, offset, (size, stride, at) in self.param_keys.get(
                 local_param.untyped_storage().data_ptr(), []
             )
             if start <= at < start + local_param.numel()
         ]
+
+    @torch.no_grad()
+    def update_hook(self, params: list[torch.Tensor]) -> Callable[[], None]:
+        """``update_hook`` of the optimizers on the GPU, right before they update ``params`` on the current stream:
+        copy the old wire bits of their shards, and diff them against the updated shards on the diff thread."""
+        if not self.armed:
+            return lambda: None
+        locals_ = [param.to_local() if isinstance(param, DTensor) else param for param in params]
+        num_bytes = sum(local.numel() for local in locals_) * 2
+        start = time.perf_counter()
+        with self.update_bytes:
+            self.update_bytes.wait_for(lambda: self.held_bytes == 0 or self.held_bytes + num_bytes <= self.UPDATE_BYTES)
+            self.held_bytes += num_bytes
+        self.hook_wait += time.perf_counter() - start
+        items = []
+        for local in locals_:
+            entries = self._entries(local)
+            if entries:
+                dtypes = {self.wire_dtypes[key] for key, *_ in entries}
+                dtype = dtypes.pop() if len(dtypes) == 1 else local.dtype
+                items.append((entries, local.to(dtype, copy=True), local))
+
+        def release() -> None:
+            with self.update_bytes:
+                self.held_bytes -= num_bytes
+                self.update_bytes.notify_all()
+
+        def finish() -> None:
+            ready = torch.cuda.Event()
+            ready.record()
+            self.diffs.append(self.differ.submit(self._diff_slot, items, ready, release, True))
+
+        return finish
+
+    @torch.no_grad()
+    def optimizer_hook(self, local_param: torch.Tensor, source: torch.Tensor) -> None:
+        """``FullCPUOffloadOptimizer.weight_hook``, on its H2D stream: copy the shard (still the previous
+        broadcast's weights) and its new values ``source`` (pinned host) into a GPU slot, write the new values and
+        diff the slot on the diff thread. At anchor steps, also writes the new values to the anchor from the
+        host."""
+        entries = self._entries(local_param) if self.armed else []
         if not entries:
             local_param.copy_(source, non_blocking=True)
             return
@@ -696,48 +695,106 @@ def optimizer_hook(self, local_param: torch.Tensor, source: torch.Tensor) -> Non
         ready = torch.cuda.Event()
         ready.record()
         if self.anchor is not None:
-            for key, _, numel, offset, (size, stride, at) in entries:
+            for key, fqn, numel, offset, (size, stride, at) in entries:
                 host = source.as_strided(size, stride, source.storage_offset() + at).reshape(-1)
                 host = host.to(self.wire_dtypes[key])
                 self.anchor.write_all(
                     [
                         (name, hf_start, host[begin:end].view(torch.uint8).numpy())
-                        for name, _, hf_start, begin, end in self._slices(key, numel, offset)
+                        for name, _, hf_start, begin, end in self._slices(fqn, numel, offset)
                     ]
                 )
-        self.diffs.append(self.differ.submit(self._diff_slot, entries, slot, previous, new, ready))
+        release = partial(self.hook_slots.put, slot)
+        self.diffs.append(self.differ.submit(self._diff_slot, [(entries, previous, new)], ready, release, False))
 
     @torch.no_grad()
-    def _diff_slot(self, entries: list[tuple], slot: torch.Tensor, previous, new, ready: torch.cuda.Event) -> None:
-        """Diff thread: the parts of one parameter's tensors from its slot, then free the slot. Hands the parts so far
-        to ``emit`` every ``HOOK_PART`` changes."""
+    def _diff_slot(self, items: list[tuple], ready: torch.cuda.Event, release: Callable, on_gpu: bool) -> None:
+        """Diff thread: the parts of the trainer tensors in ``items`` (``(entries, old values, new values)`` of
+        parameter shards), and their anchor slices if the optimizer is ``on_gpu``; then ``release`` the old values.
+        The bf16 tensors are diffed together, as one flat tensor. Hands the parts so far to ``emit`` every
+        ``HOOK_PART`` changes."""
         torch.cuda.set_device(self.device)
+        ready.synchronize()
+        start = time.perf_counter()
+        chunk = self.CHUNK if on_gpu else self.HOOK_CHUNK
         try:
             with torch.cuda.stream(self.diff_stream):
-                self.diff_stream.wait_event(ready)
-                for key, fqn, _, offset, (size, stride, at) in entries:
-                    wire = self.wire_dtypes[key]
-                    current, old = (
-                        values.as_strided(size, stride, values.storage_offset() + at).reshape(-1).to(wire)
-                        for values in (new, previous)
+                batch, written = [], []
+                for entries, previous, new in items:
+                    for key, fqn, numel, offset, (size, stride, at) in entries:
+                        current, old = (
+                            values.as_strided(size, stride, values.storage_offset() + at) for values in (new, previous)
+                        )
+                        if fqn in self.fp8 or self.wire_dtypes[key] != torch.bfloat16:
+                            current = current.reshape(-1).to(self.wire_dtypes[key])
+                            old = as_bits(old.reshape(-1).to(self.wire_dtypes[key]))
+                            self.step_parts += self._diff_shard(key, fqn, current, offset, None, old, None, chunk)
+                            self.checksums[key] = self._checksum(key, offset, as_bits(current))
+                            written.append((fqn, numel, offset, current))
+                        else:
+                            batch.append((key, fqn, numel, offset, current, old))
+                if batch:
+                    total = sum(numel for _, _, numel, *_ in batch)
+                    current_bits, previous_bits = (
+                        torch.empty(total, dtype=torch.bfloat16, device=f"cuda:{self.device}") for _ in range(2)
                     )
-                    self.step_parts += self._diff_shard(
-                        key, fqn, current, offset, None, as_bits(old), None, self.HOOK_CHUNK
+                    names, begins, ends, shifts, ranges, position = [], [], [], [], [], 0
+                    for key, fqn, numel, offset, current, old in batch:
+                        current_bits[position : position + numel].view(current.shape).copy_(current)
+                        previous_bits[position : position + numel].view(old.shape).copy_(old)
+                        segment_names, *segments = self._segments(key, self.segments[fqn], offset, numel)
+                        ranges.append((key, len(names), len(names) + len(segment_names)))
+                        names += segment_names
+                        begins += [begin + position for begin in segments[0]]
+                        ends += [end + position for end in segments[1]]
+                        shifts += [shift - position for shift in segments[2]]
+                        written.append((fqn, numel, offset, current_bits[position : position + numel]))
+                        position += numel
+                    part = self._diff_flat(
+                        tuple(key for key, *_ in batch),
+                        names,
+                        begins,
+                        ends,
+                        shifts,
+                        as_bits(current_bits),
+                        as_bits(previous_bits),
+                        chunk,
                     )
-                    self.expected[key] = self._checksum(as_bits(current))
+                    self.step_parts.append(part)
+                    for key, first, last in ranges:
+                        self.checksums[key] = part[5][first:last].sum(0, keepdim=True)
+                if on_gpu and self.anchor is not None:
+                    for fqn, numel, offset, values in written:
+                        self._write_anchor_slices(fqn, numel, offset, values)
                 self.diff_stream.synchronize()
         finally:
-            self.hook_slots.put(slot)
+            release()
+            self.diff_busy += time.perf_counter() - start
         if sum(len(part[2]) for part in self.step_parts) >= self.HOOK_PART:
             with torch.cuda.stream(self.diff_stream):
                 self.emit(self._host_part(self.step_parts, self.diff_staging))
             self.step_parts = []
 
-
-def _model_dir(model: nn.Module) -> Path:
-    """The local dir of the checkpoint the model was loaded from, as the trainer resolves it."""
-    name = model.config.name_or_path
-    return Path(name) if Path(name).exists() else Path(snapshot_download(repo_id=name, repo_type="model"))
+    def _write_anchor_slices(self, fqn: str, numel: int, offset: int, values: torch.Tensor) -> None:
+        """Write a GPU shard's anchor slices through the pinned staging buffer: the slices in each staged window in
+        parallel, in pieces of up to 32 MiB."""
+        data, itemsize = values.view(torch.uint8), values.element_size()
+        slices = [
+            (name, start, begin * itemsize, end * itemsize)
+            for name, _, start, begin, end in self._slices(fqn, numel, offset)
+        ]
+        window, piece = self.diff_staging.numel(), 1 << 25
+        for at in range(0, data.numel(), window):
+            count = min(window, data.numel() - at)
+            self.diff_staging[:count].copy_(data[at : at + count], non_blocking=True)
+            self.diff_stream.synchronize()
+            staged = self.diff_staging[:count].numpy()
+            items = []
+            for name, start, begin, end in slices:
+                for first in range(max(begin, at), min(end, at + count), piece):
+                    last = min(first + piece, end, at + count)
+                    items.append((name, start, staged[first - at : last - at], first - begin))
+            self.anchor.write_all(items)
 
 
 def _write_json(path: Path, value) -> None:
@@ -746,15 +803,14 @@ def _write_json(path: Path, value) -> None:
 
 class SparseFileSystemWeightSender(FileSystemWeightSender):
     """Filesystem broadcast where steps carry only the delta against the previous broadcast (format in
-    ``prime_rl.utils.sparse_weights``). The delta is encoded and written in the background while the next step
-    trains; the trainer waits only before its next final backward, until the encode has read the weights. Anchors
-    (full checkpoints) serve engines that fall off the delta chain. The writer thread commits a step once every
-    rank's files are on disk."""
+    ``prime_rl.utils.sparse_weights``). The delta comes from the optimizer step (``diff_in_optimizer``) and is
+    written as the step goes; the rest is finished in the background while the next step trains, and the trainer
+    waits only before its next final backward, until that work has read the weights. A broadcast that does not
+    follow an optimizer step collected this way (the first one, or any broadcast without ``diff_in_optimizer``)
+    publishes a full checkpoint. Anchors (full checkpoints every ``anchor_interval`` steps) serve engines that fall
+    off the delta chain. The writer thread commits a step once every rank's files are on disk."""
 
     gathers_full_weights = False
-    # Seconds an anchor copy may hold back the next encode before that anchor is skipped: the trainer waits for the
-    # encode, and other ranks would hit the offload pipeline's timeout
-    ANCHOR_COPY_WAIT = 30.0
 
     def __init__(self, output_dir: Path, config: FileSystemWeightBroadcastConfig):
         super().__init__(output_dir, config)
@@ -766,29 +822,37 @@ def __init__(self, output_dir: Path, config: FileSystemWeightBroadcastConfig):
         self.encodes = ThreadPoolExecutor(1)
         self.writes = ThreadPoolExecutor(1)
         self.anchors = ThreadPoolExecutor(1)
-        # Anchor writes that read the snapshot (the next encode waits for them), and anchor writes from the host
-        # buffer of a GPU snapshot (the next anchor copy waits for them)
-        self.snapshot_writes: list[Future] = []
-        self.anchor_writes: list[Future] = []
         self.errors: list[BaseException] = []
         self.weights_read: threading.Event | None = None
-        self.encode_queued_at: float | None = None
         self.part_bytes = 0
-        # The checkpoint the engines hold before the first delta: the model's (None), or a resumed run's first
-        # broadcast
-        self.boot_dir: Path | None = None
-        self.optimizer: "FullCPUOffloadOptimizer | None" = None
+        self.optimizer: "OptimizerLike | None" = None
         self.model: nn.Module | None = None
         # The step whose delta the next optimizer step collects
         self.armed_step: int | None = None
         self.streamed_dir: Path | None = None
         self.num_streamed = 0
 
-    def diff_in_optimizer(self, optimizer: "FullCPUOffloadOptimizer") -> None:
-        """Take the deltas of the optimizer's parameters from its step (``SparseDeltaEncoder.optimizer_hook``); needs
-        a broadcast after every optimizer step."""
+    def diff_in_optimizer(self, optimizer: "OptimizerLike") -> None:
+        """Take the deltas of the optimizer's parameters from its steps, which must each be followed by a broadcast."""
         self.optimizer = optimizer
 
+    def _bind(self) -> None:
+        """Hook the encoder into the optimizer's parameter updates."""
+        from prime_rl.trainer.optim import FullCPUOffloadOptimizer, OffloadOptimizer
+
+        optimizer, encoder = self.optimizer, self.encoder
+        if isinstance(optimizer, FullCPUOffloadOptimizer):
+            params = [master.model_param.to_local() for master in optimizer._master_weights.values()]
+            encoder.bind(self.model, self.world.is_master, params, full_offload=True)
+            optimizer.weight_hook = encoder.optimizer_hook
+            return
+        base = optimizer.base_optimizer if isinstance(optimizer, OffloadOptimizer) else optimizer
+        if not hasattr(base, "update_hook"):
+            raise NotImplementedError(f"Sparse weight broadcast cannot hook into {type(base).__name__} updates")
+        params = [p.to_local() if isinstance(p, DTensor) else p for group in base.param_groups for p in group["params"]]
+        encoder.bind(self.model, self.world.is_master, params, full_offload=False)
+        base.update_hook = encoder.update_hook
+
     def wait_for_weight_reads(self) -> None:
         self._wait_for_encode()
         if self.optimizer is not None and self.last_step is not None and self.armed_step is None:
@@ -805,11 +869,6 @@ def _wait_for_encode(self) -> None:
     def _arm(self, step: int) -> None:
         """Collect v{step}'s delta, and its anchor slices at anchor steps, in the coming optimizer step."""
         encoder = self.encoder
-        self._set_fp8_layout()
-        if not encoder.optimizer_keys:
-            params = [master.model_param.to_local() for master in self.optimizer._master_weights.values()]
-            encoder.bind_optimizer(self.model, self.world.is_master, params)
-            encoder.allocate_host()
         anchor = None
         if step % self.config.anchor_interval == 0:
             # Written next to the step dir, which the broadcast handshake recreates, and moved in once complete
@@ -819,7 +878,6 @@ def _arm(self, step: int) -> None:
         self.streamed_dir.mkdir(parents=True, exist_ok=True)
         self.num_streamed = 0
         encoder.arm(anchor, self._stream_part)
-        self.optimizer.weight_hook = encoder.optimizer_hook
         self.armed_step = step
 
     def _stream_part(self, part: tuple) -> None:
@@ -860,51 +918,47 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
         if self.world.is_master:
             (step_dir / DELTA_DIR).mkdir(exist_ok=True)
         if self.encoder is None:
-            self.encoder = SparseDeltaEncoder(model, self.config.snapshot)
+            self.encoder = SparseDeltaEncoder(model)
             self.model = model
+        if self.armed_step is not None and self.armed_step != step:
+            raise RuntimeError(f"Sparse weight delta v{self.armed_step} was collected, but v{step} is broadcast")
+        if self.armed_step is None:
+            self._broadcast_full(model, step, step_dir)
+            return
         encoder = self.encoder
+        encoder.disarm()
+        self.armed_step = None
         base_step, self.last_step = self.last_step, step
-        if self.optimizer is not None and base_step is not None:
-            if self.armed_step != step:
-                raise RuntimeError(f"Sparse weight delta v{step} was not collected in the optimizer step")
-            self.optimizer.weight_hook = None
-            self.armed_step = None
-        if base_step is None:
-            shards = encoder.local_shards(model, self.world.is_master)
-            encoder.track(shards)
-            if self.optimizer is None:
-                encoder.allocate_host()
-        if base_step is None and step == 0:
-            # The run's first broadcast is the checkpoint the engines booted from
+        shards = encoder.local_shards(model, self.world.is_master, gather=True)
+        ready = torch.cuda.Event()
+        ready.record()
+        self.weights_read = threading.Event()
+        args = (step, base_step, step_dir, shards, ready, self.weights_read, encoder.anchor)
+        self._submit(self.encodes, self._encode, *args)
+
+    def _broadcast_full(self, model: nn.Module, step: int, step_dir: Path) -> None:
+        """Publish v{step} as a full checkpoint of the live weights; the run's first broadcast at step 0 is the
+        checkpoint the engines booted from. The next delta is against it."""
+        if self.last_step is None and step == 0:
             self.anchor_step = step
             if self.world.is_master:
                 self._write_manifest(step_dir, step, None, 0)
                 WeightSender._finish(self, step, step_dir)
-        elif base_step is None:
-            # A resumed run's first broadcast: a full checkpoint, which the next delta is against
-            source = {key: as_bits(shard.reshape(-1).to(encoder.wire_dtypes[key])) for key, _, shard, _, _ in shards}
-            torch.cuda.synchronize()
-            self._write_full(step, step_dir, encoder.anchor_parts(source))
-            self.boot_dir = step_dir
-        if base_step is None:
-            if self.optimizer is not None:
-                # Engines report their layout as they load this broadcast. Take it (and compile the fp8 quantizer)
-                # now, not in the first optimizer step, next to its activations.
-                self._set_fp8_layout()
-            return
-        self._set_fp8_layout()
-        if not encoder.initialized:
-            encoder.initialize(
-                model, self.world.is_master, torch.cuda.max_memory_reserved(), self.boot_dir or _model_dir(model)
+        else:
+            self.encoder.write_full_checkpoint(
+                model, self.world.is_master, step_dir, self.world.rank, self.world.world_size
             )
-        shards = encoder.local_shards(model, self.world.is_master)
-        anchor = encoder.anchor if self.optimizer is not None else step % self.config.anchor_interval == 0
-        ready = torch.cuda.Event()
-        ready.record()
-        self.weights_read = threading.Event()
-        self.encode_queued_at = time.monotonic()
-        args = (step, base_step, step_dir, shards, ready, self.weights_read, anchor)
-        self._submit(self.encodes, self._encode, *args)
+            self._wait_for_anchor(step, step_dir)
+            if self.world.is_master:
+                WeightSender._finish(self, step, step_dir)
+        self.last_step = step
+        if self.optimizer is not None:
+            # Engines report their layout as they load this broadcast. Take it (and compile the fp8 quantizer) now,
+            # not in the first optimizer step, next to its activations.
+            self._set_fp8_layout()
+            if self.encoder.differ is None:
+                self._bind()
+            self.encoder.reset(model, self.world.is_master)
 
     def _wait_for_engine_layout(self, broadcast_dir: Path) -> dict[str, str]:
         start = time.monotonic()
@@ -914,38 +968,6 @@ def _wait_for_engine_layout(self, broadcast_dir: Path) -> dict[str, str]:
             time.sleep(0.1)
         return layout
 
-    @staticmethod
-    def _wait(writes: list[Future]) -> None:
-        for write in writes:
-            write.exception()
-        writes.clear()
-
-    def _write_periodic_anchor(self, step: int, step_dir: Path) -> None:
-        """Write this rank's slices of the last broadcast in the background. A GPU snapshot is first copied to the
-        host buffer, so the next encode waits only for that copy; a CPU snapshot is written from itself, and the
-        next encode waits for the write."""
-        encoder = self.encoder
-        if encoder.on_gpu:
-
-            def held_too_long() -> bool:
-                queued_at = self.encode_queued_at
-                return queued_at is not None and time.monotonic() - queued_at > self.ANCHOR_COPY_WAIT
-
-            while not all(write.done() for write in self.anchor_writes) and not held_too_long():
-                time.sleep(0.1)
-            if held_too_long() or not encoder.copy_snapshot_to_host(held_too_long):
-                self.logger.warning(
-                    f"Skipped the {step_dir.name} anchor: its copy to host memory held the next encode for over "
-                    f"{self.ANCHOR_COPY_WAIT:.0f}s"
-                )
-                write_file(step_dir / f"anchor_rank{self.world.rank}.skipped", b"")
-                return
-            self._wait(self.anchor_writes)
-            parts, writes = encoder.anchor_parts(encoder.host), self.anchor_writes
-        else:
-            parts, writes = encoder.anchor_parts(encoder.snapshot), self.snapshot_writes
-        writes.append(self._submit(self.anchors, self._write_anchor, step, step_dir, parts))
-
     def _encode(
         self,
         step: int,
@@ -954,13 +976,10 @@ def _encode(
         shards,
         ready,
         weights_read: threading.Event,
-        anchor: bool | AnchorWriter | None,
+        anchor: AnchorWriter | None,
     ) -> None:
-        """Encode thread: diff v{step}, handing the delta to the writer in parts. ``anchor``: write a periodic
-        anchor (True), or finish the one the optimizer step wrote."""
+        """Encode thread: finish v{step}'s delta, and the anchor the optimizer step wrote."""
         encoder = self.encoder
-        self.encode_queued_at = None
-        self._wait(self.snapshot_writes)
         start = time.perf_counter()
         num_parts = 0
 
@@ -972,29 +991,26 @@ def emit(part: tuple) -> None:
 
         try:
             encoder.wait_for_diffs()
-            if self.optimizer is not None:
-                num_parts = self.num_streamed
+            num_parts = self.num_streamed
             if self.world.is_master:
                 self._submit(self.writes, self._write_manifest, step_dir, step, base_step, self.world.world_size)
-            if self.optimizer is not None:
-                self._submit(self.writes, self._move_streamed, self.streamed_dir, step_dir / DELTA_DIR)
+            self._submit(self.writes, self._move_streamed, self.streamed_dir, step_dir / DELTA_DIR)
             scale_changes = encoder.encode(shards, ready, weights_read, emit)
         except BaseException as error:
             # Recorded before the trainer wakes up, so it fails at its next wait instead of training on
             self.errors.append(error)
             weights_read.set()
             raise
-        # Free the gathered fp8 tensors before the anchor copy below, which may overlap the next step's memory peak
+        # Free the gathered fp8 tensors
         shards.clear()
         encode_time = time.perf_counter() - start
         args = (step, base_step, step_dir, num_parts, start, encode_time, scale_changes)
         self._submit(self.writes, self._finish_delta, *args)
-        if isinstance(anchor, AnchorWriter):
-            # The snapshot tensors hold v{step} until the next encode, which runs on this thread
-            anchor.write_parts(encoder.anchor_parts(encoder.snapshot), encoder.stream_to_host)
+        if anchor is not None:
+            # The residual's previous bits hold v{step} until the next encode, which runs on this thread
+            source = {key: (self.encoder.shards[key][1], bits) for key, bits in encoder.previous.items()}
+            anchor.write_parts(encoder.anchor_parts(source), encoder.stream_to_host)
             self._submit(self.anchors, self._close_anchor, step, step_dir, anchor)
-        elif anchor:
-            self._write_periodic_anchor(step, step_dir)
 
     def _write_manifest(self, step_dir: Path, step: int, base_step: int | None, ranks: int) -> None:
         """Written before the parts, so engines can read part files as they appear."""
@@ -1051,16 +1067,6 @@ def _finish_delta(
             f"encode {encode_time:.2f}s, committed {time.perf_counter() - start:.2f}s after the encode started{scales}"
         )
 
-    def _write_full(self, step: int, step_dir: Path, parts: list) -> None:
-        """Publish v{step} as a full checkpoint: every rank writes its slices, the master commits."""
-        self._write_anchor(step, step_dir, parts)
-        if self.world.is_master:
-            WeightSender._finish(self, step, step_dir)
-
-    def _write_anchor(self, step: int, step_dir: Path, parts: list) -> None:
-        write_anchor_shards(parts, step_dir, self.world.rank, self.world.world_size, self.encoder.stream_to_host)
-        self._wait_for_anchor(step, step_dir)
-
     def _close_anchor(self, step: int, step_dir: Path, anchor: AnchorWriter) -> None:
         anchor.close(self.world.world_size)
         self._wait_for_anchor(step, step_dir)
@@ -1068,14 +1074,11 @@ def _close_anchor(self, step: int, step_dir: Path, anchor: AnchorWriter) -> None
     def _wait_for_anchor(self, step: int, step_dir: Path) -> None:
         if self.world.is_master:
             while not anchor_complete(step_dir):
-                if any(step_dir.glob("anchor_rank*.skipped")):
-                    self.logger.warning(f"The {step_dir.name} anchor is incomplete: a rank skipped it")
-                    return
                 time.sleep(0.1)
             self.anchor_step = max(self.anchor_step or 0, step)
 
     def _finish(self, step: int, step_dir: Path) -> None:
-        """Committed in ``_broadcast`` at v0, else in the background (``_finish_delta``, ``_write_full``)."""
+        """Committed in ``_broadcast_full`` or in the background (``_finish_delta``)."""
 
     def _clean(self, step: int) -> None:
         # Engines replay from the last anchor, so keep everything since it
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index 56e994bcc27..6ac0b066263 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -133,7 +133,7 @@ def encode_gaps(indices: torch.Tensor, entry_starts: torch.Tensor) -> torch.Tens
     ``entry_starts`` position (where the gap is the index itself)."""
     previous = torch.empty_like(indices)
     previous[1:] = indices[:-1]
-    previous[entry_starts] = -1
+    previous.masked_fill_(entry_starts, -1)
     return (indices - previous - 1).to(torch.int32)
 
 
@@ -332,21 +332,6 @@ def close(self, world_size: int) -> None:
         write_file(self.step_dir / f"anchor_rank{self.rank}.done", str(world_size).encode())
 
 
-def write_anchor_shards(
-    parts: list[tuple[str, list[int], int, torch.Tensor]],
-    step_dir: Path,
-    rank: int,
-    world_size: int,
-    to_host: Callable[[torch.Tensor], Iterable],
-) -> None:
-    """Write one trainer rank's share of a full checkpoint from ``AnchorWriter.write_parts`` slices."""
-    writer = AnchorWriter(
-        [(name, shape, start, values.numel(), values.dtype) for name, shape, start, values in parts], step_dir, rank
-    )
-    writer.write_parts(parts, to_host)
-    writer.close(world_size)
-
-
 def checkpoint_index(directory: Path) -> dict[str, tuple[list[int], torch.dtype, list[tuple[int, int, Path, int]]]]:
     """Tensor name -> ``(shape, dtype, slices)`` of the safetensors files in ``directory``, with slices
     ``(flat start, numel, file, data offset)``: a HF checkpoint (one slice at 0), or an anchor
@@ -463,12 +448,8 @@ def read_delta(
 
 def pin(num_bytes: int) -> torch.Tensor:
     """Page-locked host memory of exactly ``num_bytes``, outside torch's caching host allocator (which rounds sizes
-    up and keeps freed blocks). Release it with ``unpin``."""
-    return page_lock(torch.from_numpy(np.empty(num_bytes, np.uint8)))
-
-
-def page_lock(tensor: torch.Tensor) -> torch.Tensor:
-    """Page-lock a host tensor in place with one ``cudaHostRegister`` (a copy may not span two)."""
+    up and keeps freed blocks), with one ``cudaHostRegister`` (a copy may not span two). Release it with ``unpin``."""
+    tensor = torch.from_numpy(np.empty(num_bytes, np.uint8))
     if tensor.numel():
         error = torch.cuda.cudart().cudaHostRegister(tensor.data_ptr(), tensor.nbytes, 0)
         assert error == 0, f"cudaHostRegister failed: {error}"

From aa3c6957cbb4c2ef2e5fcc690d590c3a11d8b353 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sun, 18 Oct 2026 10:27:29 +0000
Subject: [PATCH 55/88] fix(weights): write the anchor slices of optimizer
 tensors the step did not update; check every anchor slice is written

---
 src/prime_rl/transports/weights/sparse.py | 18 +++++++++++++-----
 src/prime_rl/utils/sparse_weights.py      | 11 +++++++++++
 tests/unit/utils/test_sparse_weights.py   | 14 ++++++++++++++
 3 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 07f9811a6fc..174b40df543 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -420,7 +420,11 @@ def flush() -> None:
         for key, fqn, shard, offset, gathered in shards:
             shard = shard.reshape(-1)
             if key in self.checksums:
-                checks.append((key, self._checksum(key, offset, as_bits(shard.to(self.wire_dtypes[key])))))
+                bits = as_bits(shard.to(self.wire_dtypes[key]))
+                checks.append((key, self._checksum(key, offset, bits)))
+                if self.anchor is not None and key not in self.anchored:
+                    # Not updated in this step (e.g. no gradient): its anchor slices come from the live weights
+                    self._write_anchor_slices(fqn, len(bits), offset, bits)
                 continue
             parts = self._diff_shard(key, fqn, shard, offset, gathered, self.previous[key], self.previous[key])
             pending += parts
@@ -613,6 +617,8 @@ def arm(self, anchor: AnchorWriter | None, emit: Callable) -> None:
         """Collect the next optimizer step's delta, handing host parts to ``emit`` as the step goes, and write its
         anchor slices if ``anchor``."""
         self.step_parts, self.scale_changes, self.anchor, self.emit = [], 0, anchor, emit
+        # Optimizer tensors whose anchor slices the step wrote
+        self.anchored: set[str] = set()
         self.hook_wait = self.diff_busy = 0.0
         self.armed = True
 
@@ -698,6 +704,7 @@ def optimizer_hook(self, local_param: torch.Tensor, source: torch.Tensor) -> Non
             for key, fqn, numel, offset, (size, stride, at) in entries:
                 host = source.as_strided(size, stride, source.storage_offset() + at).reshape(-1)
                 host = host.to(self.wire_dtypes[key])
+                self.anchored.add(key)
                 self.anchor.write_all(
                     [
                         (name, hf_start, host[begin:end].view(torch.uint8).numpy())
@@ -730,7 +737,7 @@ def _diff_slot(self, items: list[tuple], ready: torch.cuda.Event, release: Calla
                             old = as_bits(old.reshape(-1).to(self.wire_dtypes[key]))
                             self.step_parts += self._diff_shard(key, fqn, current, offset, None, old, None, chunk)
                             self.checksums[key] = self._checksum(key, offset, as_bits(current))
-                            written.append((fqn, numel, offset, current))
+                            written.append((key, fqn, numel, offset, current))
                         else:
                             batch.append((key, fqn, numel, offset, current, old))
                 if batch:
@@ -748,7 +755,7 @@ def _diff_slot(self, items: list[tuple], ready: torch.cuda.Event, release: Calla
                         begins += [begin + position for begin in segments[0]]
                         ends += [end + position for end in segments[1]]
                         shifts += [shift - position for shift in segments[2]]
-                        written.append((fqn, numel, offset, current_bits[position : position + numel]))
+                        written.append((key, fqn, numel, offset, current_bits[position : position + numel]))
                         position += numel
                     part = self._diff_flat(
                         tuple(key for key, *_ in batch),
@@ -764,8 +771,9 @@ def _diff_slot(self, items: list[tuple], ready: torch.cuda.Event, release: Calla
                     for key, first, last in ranges:
                         self.checksums[key] = part[5][first:last].sum(0, keepdim=True)
                 if on_gpu and self.anchor is not None:
-                    for fqn, numel, offset, values in written:
+                    for key, fqn, numel, offset, values in written:
                         self._write_anchor_slices(fqn, numel, offset, values)
+                        self.anchored.add(key)
                 self.diff_stream.synchronize()
         finally:
             release()
@@ -787,7 +795,7 @@ def _write_anchor_slices(self, fqn: str, numel: int, offset: int, values: torch.
         for at in range(0, data.numel(), window):
             count = min(window, data.numel() - at)
             self.diff_staging[:count].copy_(data[at : at + count], non_blocking=True)
-            self.diff_stream.synchronize()
+            torch.cuda.current_stream().synchronize()
             staged = self.diff_staging[:count].numpy()
             items = []
             for name, start, begin, end in slices:
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index 6ac0b066263..26f4a7ed531 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -269,6 +269,9 @@ def __init__(self, slices: list[tuple[str, list[int], int, int, torch.dtype]], s
         self.tmp_dir.mkdir(parents=True, exist_ok=True)
         self.files: list[tuple[str, int]] = []
         self.offsets: dict[tuple[str, int], tuple[int, int]] = {}
+        self.sizes = {(name, start): numel * dtype.itemsize for name, _, start, numel, dtype in slices}
+        # (slice, bytes) of every write, from any thread
+        self.written: list[tuple[tuple[str, int], int]] = []
         # Round-robin over the files, which are written in parallel
         num_files = max(1, -(-sum(numel * dtype.itemsize for *_, numel, dtype in slices) // 2**30))
         groups = [slices[index::num_files] for index in range(num_files)]
@@ -305,6 +308,7 @@ def write(self, name: str, start: int, data, at: int = 0) -> None:
         """Write ``data`` (bytes) at byte ``at`` of the slice ``"{name}@{start}"``."""
         fd, offset = self.offsets[(name, start)]
         self._pwrite(fd, data, offset + at)
+        self.written.append(((name, start), memoryview(data).nbytes))
 
     def write_parts(
         self, parts: list[tuple[str, list[int], int, torch.Tensor]], to_host: Callable[[torch.Tensor], Iterable]
@@ -323,7 +327,14 @@ def write_all(self, items: list[tuple[str, int, Any]]) -> None:
             future.result()
 
     def close(self, world_size: int) -> None:
+        """Publish the files and the marker; raises if a slice was not written in full."""
         self.pool.shutdown()
+        written = dict.fromkeys(self.sizes, 0)
+        for key, num_bytes in self.written:
+            written[key] += num_bytes
+        missing = [f"{name}@{start}" for (name, start), size in self.sizes.items() if written[(name, start)] != size]
+        if missing:
+            raise RuntimeError(f"Anchor {self.step_dir.name} is missing data of {len(missing)} slices: {missing[:8]}")
         for name, fd in self.files:
             drop_page_cache(fd)
             os.close(fd)
diff --git a/tests/unit/utils/test_sparse_weights.py b/tests/unit/utils/test_sparse_weights.py
index 0d4af944e3d..9e1c08d0029 100644
--- a/tests/unit/utils/test_sparse_weights.py
+++ b/tests/unit/utils/test_sparse_weights.py
@@ -1,11 +1,14 @@
 import json
 from concurrent.futures import ThreadPoolExecutor
 
+import pytest
 import torch
 
 from prime_rl.utils.sparse_weights import (
     DELTA_DIR,
     MANIFEST,
+    AnchorWriter,
+    anchor_complete,
     as_bits,
     bits_checksum,
     decode_gaps,
@@ -73,3 +76,14 @@ def apply(names):
     patched = apply({"a"})
     assert torch.equal(patched["a"], as_bits(current["a"]).view(-1))
     assert torch.equal(patched["b"], as_bits(previous["b"]).view(-1))
+
+
+def test_anchor_with_an_unwritten_slice_is_not_published(tmp_path):
+    step_dir = tmp_path / "step_5"
+    step_dir.mkdir()
+    values = torch.arange(6, dtype=torch.bfloat16)
+    writer = AnchorWriter([("a", [6], 0, 6, torch.bfloat16), ("b", [6], 0, 6, torch.bfloat16)], step_dir, 0)
+    writer.write_parts([("a", [6], 0, values)], lambda tensor: [tensor.view(torch.uint8).numpy()])
+    with pytest.raises(RuntimeError, match="b@0"):
+        writer.close(1)
+    assert not anchor_complete(step_dir)

From d5ea649c3dae286d63ac499ae828a4c8a3d459b6 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sun, 18 Oct 2026 13:50:47 +0000
Subject: [PATCH 56/88] feat(weights): anchor every 50 steps by default

---
 docs/scaling.md                                          | 2 +-
 packages/prime-rl-configs/src/prime_rl/configs/shared.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index f84657a0d60..54500d548d0 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -282,7 +282,7 @@ With `[weight_broadcast] type = "filesystem"` and `sparse = true`, the trainer p
 [weight_broadcast]
 type = "filesystem"
 sparse = true
-anchor_interval = 25  # full checkpoint every N steps
+anchor_interval = 50  # full checkpoint every N steps
 ```
 
 - RL: deltas come from the optimizer step, with every optimizer and offload mode. Right before the step updates a chunk of parameters, their shards still hold the previous broadcast's weights; the trainer diffs the old and new values of each chunk as the step goes and writes the delta out. No copy of the weights is kept, except of the few tensors that are not optimizer parameters (buffers, frozen parameters, fp8 tensors whose 128-row blocks straddle trainer ranks). The trainer checks that no weight changed outside the optimizer step.
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
index c054e9c236b..ed5e5f3fc28 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
@@ -52,7 +52,7 @@ class FileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig):
     """Between anchors, publish only the weight elements that changed since the previous broadcast.
     Engines patch their weights in place and verify every update. Not supported with LoRA."""
 
-    anchor_interval: int = Field(25, ge=1)
+    anchor_interval: int = Field(50, ge=1)
     """With ``sparse``, publish a full checkpoint every this many steps. Engines that join or fall
     off the delta chain load the last one and replay the deltas since."""
 

From bc69bd592e9a386ac956ac975926500521f544bb Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sun, 18 Oct 2026 13:50:48 +0000
Subject: [PATCH 57/88] fix(weights): verify every tensor of anchors and full
 checkpoints on load

- Anchors and full checkpoints carry each rank's per-tensor checksums of the values written (anchor_rank<r>.checksums); engines check the rows they read on every load (join, replay, retry, full checkpoints), summed over the TP group for split tensors.
- Optimizer tensors the step did not update add an empty delta part with their checksums, so delta verification covers them too.
---
 .../inference/vllm/worker/filesystem.py       | 46 +++++++++++++-
 src/prime_rl/transports/weights/sparse.py     | 60 ++++++++++++-------
 src/prime_rl/utils/sparse_weights.py          | 40 ++++++++++++-
 tests/unit/utils/test_sparse_weights.py       | 30 +++++++---
 4 files changed, 139 insertions(+), 37 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 4ab74b37b46..5a29adf0cb1 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -23,13 +23,16 @@
     ENGINE_LAYOUT,
     MANIFEST,
     anchor_complete,
+    as_bits,
     bits_dtype,
     checkpoint_index,
     decode_gaps,
     pin,
+    read_anchor_checksums,
     read_checkpoint_range,
     read_delta,
     read_manifest,
+    rows_checksum,
     unpin,
     write_file,
 )
@@ -292,6 +295,35 @@ def _verify(self, step_dir: Path, checksums: dict[str, int]) -> None:
         if mismatched:
             raise RuntimeError(f"Sparse weight delta {step_dir.name} diverged from the trainer: {mismatched[:8]}")
 
+    def _verify_rows(
+        self, step_dir: Path, sums: dict[str, torch.Tensor], partial: dict[str, bool], expected: dict[str, int]
+    ) -> None:
+        """Check the rows read from an anchor or full checkpoint against the trainer's checksums. Tensors this
+        worker read only some rows of are summed over the TP group first (not for fp8 tensors, read in whole
+        128-row blocks that TP workers may share); without TP they are not checked."""
+        if not sums:
+            return
+        index_map, tp_group = self.index_map, get_tp_group()
+        names = sorted(sums)
+        totals = torch.stack([sums[name] for name in names])
+        partial = torch.tensor([partial[name] for name in names], device=totals.device)
+        if tp_group.world_size > 1 and partial.any():
+            reduced = totals[partial]
+            torch.distributed.all_reduce(reduced, group=tp_group.device_group)
+            totals[partial] = reduced
+        skip = (
+            partial
+            if tp_group.world_size == 1
+            else partial & torch.tensor([name in index_map.fp8 for name in names], device=totals.device)
+        )
+        mismatched = [
+            name
+            for name, value, skipped in zip(names, totals.tolist(), skip.tolist())
+            if not skipped and value != expected[name]
+        ]
+        if mismatched:
+            raise RuntimeError(f"Sparse weight checkpoint {step_dir.name} diverged from the trainer: {mismatched[:8]}")
+
     def _load_checkpoint(self, step_dir: Path) -> None:
         """Reload every weight from the HF checkpoint (or anchor) in ``step_dir``. In a sparse run, each worker reads
         only the rows it holds, straight into its params."""
@@ -312,6 +344,10 @@ def _load_checkpoint(self, step_dir: Path) -> None:
         start = time.perf_counter()
         index_map, index = self.index_map, checkpoint_index(step_dir)
         device = next(model.parameters()).device
+        # Checksums of anchors and full checkpoints written by the trainer (none for the boot checkpoint)
+        expected = read_anchor_checksums(step_dir)
+        sums: dict[str, torch.Tensor] = {}
+        partial: dict[str, bool] = {}
 
         def read(name: str, begin: int, end: int) -> torch.Tensor:
             num_cols, dtype = index_map.num_cols[name], index[name][1]
@@ -324,11 +360,17 @@ def read(name: str, begin: int, end: int) -> torch.Tensor:
         with ThreadPoolExecutor(16) as pool:
             # A bounded window of reads ahead of the GPU writes
             for window in itertools.batched(ranges, 64):
-                for (name, (begin, _)), rows in zip(window, pool.map(lambda item: read(item[0], *item[1]), window)):
-                    index_map.load_rows(name, rows.to(device), begin)
+                for (name, (begin, end)), rows in zip(window, pool.map(lambda item: read(item[0], *item[1]), window)):
+                    rows = rows.to(device)
+                    index_map.load_rows(name, rows, begin)
+                    if name in expected:
+                        num_cols = index_map.num_cols[name]
+                        sums[name] = rows_checksum(as_bits(rows).reshape(-1), begin * num_cols)
+                        partial[name] = (end - begin) * num_cols != index_map.numel[name]
                     num_bytes += rows.nbytes
         update_mla_absorbed_weights(model)
         torch.cuda.synchronize()
+        self._verify_rows(step_dir, sums, partial, expected)
         logger.info(
             f"Loaded this worker's {num_bytes / 2**30:.2f} GiB of {step_dir.name} in {time.perf_counter() - start:.1f}s"
         )
diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 174b40df543..354e4c2732d 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -121,6 +121,7 @@ def __init__(self, model: nn.Module):
         self.step_parts: list[tuple] = []
         self.diffs: list[Future] = []
         self.anchor: AnchorWriter | None = None
+        self.stepped: set[str] = set()
         self.write_staging: torch.Tensor | None = None
         self.device = torch.cuda.current_device()
         self.stream = torch.cuda.Stream()
@@ -253,20 +254,21 @@ def anchor_parts(self, source: dict[str, torch.Tensor]) -> list[tuple[str, list[
     def write_full_checkpoint(self, model: nn.Module, is_master: bool, step_dir: Path, rank: int, world_size: int):
         """Write this rank's slices of a full HF checkpoint of the live weights, one tensor at a time."""
         shards = [
-            (key, shard, self._slices(fqn, shard.numel(), offset))
+            (key, shard, offset, self._slices(fqn, shard.numel(), offset))
             for key, fqn, shard, offset, _ in self.local_shards(model, is_master)
         ]
         writer = AnchorWriter(
             [
                 (name, shape, start, end - begin, self.wire_dtypes[key])
-                for key, _, slices in shards
+                for key, _, _, slices in shards
                 for name, shape, start, begin, end in slices
             ],
             step_dir,
             rank,
         )
-        for key, shard, slices in shards:
+        for key, shard, offset, slices in shards:
             values = shard.reshape(-1).to(self.wire_dtypes[key])
+            writer.add_sums(*self._value_sums(key, offset, as_bits(values)))
             parts = [(name, shape, start, values[begin:end]) for name, shape, start, begin, end in slices]
             writer.write_parts(parts, self.stream_to_host)
         writer.close(world_size)
@@ -422,9 +424,12 @@ def flush() -> None:
             if key in self.checksums:
                 bits = as_bits(shard.to(self.wire_dtypes[key]))
                 checks.append((key, self._checksum(key, offset, bits)))
-                if self.anchor is not None and key not in self.anchored:
-                    # Not updated in this step (e.g. no gradient): its anchor slices come from the live weights
-                    self._write_anchor_slices(fqn, len(bits), offset, bits)
+                if key not in self.stepped:
+                    # Not updated in this step (e.g. no gradient): an empty part carries its checksums, and its anchor
+                    # slices come from the live weights
+                    pending += self._diff_shard(key, fqn, shard, offset, None, bits, None)
+                    if self.anchor is not None:
+                        self._write_anchor_slices(key, fqn, len(bits), offset, bits)
                 continue
             parts = self._diff_shard(key, fqn, shard, offset, gathered, self.previous[key], self.previous[key])
             pending += parts
@@ -558,15 +563,20 @@ def _window_sums(self, bits, start: int, end: int, begins, ends, shifts) -> torc
             )
         return sums
 
-    def _checksum(self, key: str, offset: int, bits: torch.Tensor) -> torch.Tensor:
-        """The sum of the ``bits_checksum``s of the HF tensors in a key's flat shard bits, in chunks of
-        ``HOOK_CHUNK`` elements: what the optimizer step records for the tensors it writes."""
-        _, begins, ends, shifts = self._segments(key, self.segments[self.key_fqn[key]], offset, bits.numel())
-        total = torch.zeros(1, dtype=torch.int64, device=bits.device)
+    def _value_sums(self, key: str, offset: int, bits: torch.Tensor) -> tuple[list[str], torch.Tensor]:
+        """The ``bits_checksum`` of each HF tensor in a key's flat shard bits, in chunks of ``HOOK_CHUNK``
+        elements."""
+        names, begins, ends, shifts = self._segments(key, self.segments[self.key_fqn[key]], offset, bits.numel())
+        sums = torch.zeros(len(names), dtype=torch.int64, device=bits.device)
         for start in range(0, bits.numel(), self.HOOK_CHUNK):
             end = min(start + self.HOOK_CHUNK, bits.numel())
-            total += self._window_sums(bits[start:end], start, end, begins, ends, shifts).sum()
-        return total
+            sums += self._window_sums(bits[start:end], start, end, begins, ends, shifts)
+        return names, sums
+
+    def _checksum(self, key: str, offset: int, bits: torch.Tensor) -> torch.Tensor:
+        """The sum of the ``bits_checksum``s of the HF tensors in a key's flat shard bits: what the optimizer step
+        records for the tensors it writes."""
+        return self._value_sums(key, offset, bits)[1].sum(0, keepdim=True)
 
     def bind(self, model: nn.Module, is_master: bool, params: list[torch.Tensor], full_offload: bool) -> None:
         """Diff the trainer tensors stored in ``params`` (the optimizer's local parameter shards, contiguous) in its
@@ -617,8 +627,8 @@ def arm(self, anchor: AnchorWriter | None, emit: Callable) -> None:
         """Collect the next optimizer step's delta, handing host parts to ``emit`` as the step goes, and write its
         anchor slices if ``anchor``."""
         self.step_parts, self.scale_changes, self.anchor, self.emit = [], 0, anchor, emit
-        # Optimizer tensors whose anchor slices the step wrote
-        self.anchored: set[str] = set()
+        # Optimizer tensors the step updated (and wrote the anchor slices of)
+        self.stepped: set[str] = set()
         self.hook_wait = self.diff_busy = 0.0
         self.armed = True
 
@@ -704,7 +714,6 @@ def optimizer_hook(self, local_param: torch.Tensor, source: torch.Tensor) -> Non
             for key, fqn, numel, offset, (size, stride, at) in entries:
                 host = source.as_strided(size, stride, source.storage_offset() + at).reshape(-1)
                 host = host.to(self.wire_dtypes[key])
-                self.anchored.add(key)
                 self.anchor.write_all(
                     [
                         (name, hf_start, host[begin:end].view(torch.uint8).numpy())
@@ -770,10 +779,12 @@ def _diff_slot(self, items: list[tuple], ready: torch.cuda.Event, release: Calla
                     self.step_parts.append(part)
                     for key, first, last in ranges:
                         self.checksums[key] = part[5][first:last].sum(0, keepdim=True)
-                if on_gpu and self.anchor is not None:
-                    for key, fqn, numel, offset, values in written:
-                        self._write_anchor_slices(fqn, numel, offset, values)
-                        self.anchored.add(key)
+                for key, fqn, numel, offset, values in written if self.anchor is not None else []:
+                    if on_gpu:
+                        self._write_anchor_slices(key, fqn, numel, offset, values)
+                    else:  # written by optimizer_hook
+                        self.anchor.add_sums(*self._value_sums(key, offset, as_bits(values)))
+                self.stepped.update(key for key, *_ in written)
                 self.diff_stream.synchronize()
         finally:
             release()
@@ -783,9 +794,10 @@ def _diff_slot(self, items: list[tuple], ready: torch.cuda.Event, release: Calla
                 self.emit(self._host_part(self.step_parts, self.diff_staging))
             self.step_parts = []
 
-    def _write_anchor_slices(self, fqn: str, numel: int, offset: int, values: torch.Tensor) -> None:
-        """Write a GPU shard's anchor slices through the pinned staging buffer: the slices in each staged window in
-        parallel, in pieces of up to 32 MiB."""
+    def _write_anchor_slices(self, key: str, fqn: str, numel: int, offset: int, values: torch.Tensor) -> None:
+        """Write a GPU shard's anchor slices, and their checksums, through the pinned staging buffer: the slices in
+        each staged window in parallel, in pieces of up to 32 MiB."""
+        self.anchor.add_sums(*self._value_sums(key, offset, as_bits(values)))
         data, itemsize = values.view(torch.uint8), values.element_size()
         slices = [
             (name, start, begin * itemsize, end * itemsize)
@@ -1017,6 +1029,8 @@ def emit(part: tuple) -> None:
         if anchor is not None:
             # The residual's previous bits hold v{step} until the next encode, which runs on this thread
             source = {key: (self.encoder.shards[key][1], bits) for key, bits in encoder.previous.items()}
+            for key, (offset, bits) in source.items():
+                anchor.add_sums(*encoder._value_sums(key, offset, bits))
             anchor.write_parts(encoder.anchor_parts(source), encoder.stream_to_host)
             self._submit(self.anchors, self._close_anchor, step, step_dir, anchor)
 
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index 26f4a7ed531..66462523a55 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -272,6 +272,8 @@ def __init__(self, slices: list[tuple[str, list[int], int, int, torch.dtype]], s
         self.sizes = {(name, start): numel * dtype.itemsize for name, _, start, numel, dtype in slices}
         # (slice, bytes) of every write, from any thread
         self.written: list[tuple[tuple[str, int], int]] = []
+        # (HF names, their checksums) of the written values
+        self.sums: list[tuple[list[str], torch.Tensor]] = []
         # Round-robin over the files, which are written in parallel
         num_files = max(1, -(-sum(numel * dtype.itemsize for *_, numel, dtype in slices) // 2**30))
         groups = [slices[index::num_files] for index in range(num_files)]
@@ -326,8 +328,13 @@ def write_all(self, items: list[tuple[str, int, Any]]) -> None:
         for future in [self.pool.submit(self.write, *item) for item in items]:
             future.result()
 
+    def add_sums(self, names: list[str], sums: torch.Tensor) -> None:
+        """Record the ``bits_checksum``s of the values written for ``names`` (partial sums add up)."""
+        self.sums.append((names, sums))
+
     def close(self, world_size: int) -> None:
-        """Publish the files and the marker; raises if a slice was not written in full."""
+        """Publish the files, the checksums of this rank's values and the marker; raises if a slice was not written
+        in full."""
         self.pool.shutdown()
         written = dict.fromkeys(self.sizes, 0)
         for key, num_bytes in self.written:
@@ -340,6 +347,11 @@ def close(self, world_size: int) -> None:
             os.close(fd)
             (self.tmp_dir / name).rename(self.step_dir / name)
         self.tmp_dir.rmdir()
+        checksums: dict[str, int] = {}
+        for names, sums in self.sums:
+            for name, value in zip(names, sums.tolist()):
+                checksums[name] = add_checksums(checksums.get(name, 0), value)
+        write_file(self.step_dir / anchor_checksums_file(self.rank), json.dumps(checksums).encode())
         write_file(self.step_dir / f"anchor_rank{self.rank}.done", str(world_size).encode())
 
 
@@ -381,6 +393,19 @@ def read_checkpoint_range(entry: tuple, begin: int, end: int, out: np.ndarray) -
             os.posix_fadvise(file.fileno(), offset, size, os.POSIX_FADV_DONTNEED)
 
 
+def anchor_checksums_file(rank: int) -> str:
+    return f"anchor_rank{rank}.checksums"
+
+
+def read_anchor_checksums(step_dir: Path) -> dict[str, int]:
+    """The ``bits_checksum`` per HF tensor of a complete anchor (or full checkpoint), summed over trainer ranks."""
+    checksums: dict[str, int] = {}
+    for path in step_dir.glob("anchor_rank*.checksums"):
+        for name, value in json.loads(path.read_text()).items():
+            checksums[name] = add_checksums(checksums.get(name, 0), value)
+    return checksums
+
+
 def anchor_complete(step_dir: Path) -> bool:
     """Whether every trainer rank has written its share of an anchor in ``step_dir``."""
     revalidate(step_dir)
@@ -493,8 +518,17 @@ def bits_checksum(bits: torch.Tensor, start: int = 0) -> int:
     in the full tensor (the slice starts at ``start``): additive over any split of a tensor, so
     trainer ranks and engine workers can each sum the part they hold, and a change written at a
     wrong position is caught."""
-    positions = torch.arange(start + 1, start + 1 + bits.numel(), device=bits.device)
-    return int((bits.reshape(-1).long() * positions).sum())
+    return int(rows_checksum(bits.reshape(-1), start))
+
+
+def rows_checksum(bits: torch.Tensor, start: int = 0, chunk: int = 1 << 24) -> torch.Tensor:
+    """``bits_checksum`` of a flat tensor slice as a tensor on its device, in chunks of ``chunk`` elements."""
+    total = torch.zeros((), dtype=torch.int64, device=bits.device)
+    for first in range(0, bits.numel(), chunk):
+        part = bits[first : first + chunk]
+        positions = torch.arange(start + first + 1, start + first + 1 + part.numel(), device=bits.device)
+        total += (part.long() * positions).sum()
+    return total
 
 
 def add_checksums(a: int, b: int) -> int:
diff --git a/tests/unit/utils/test_sparse_weights.py b/tests/unit/utils/test_sparse_weights.py
index 9e1c08d0029..f6e7da3e95b 100644
--- a/tests/unit/utils/test_sparse_weights.py
+++ b/tests/unit/utils/test_sparse_weights.py
@@ -15,7 +15,9 @@
     encode_gaps,
     rank_part,
     rank_parts_marker,
+    read_anchor_checksums,
     read_delta,
+    rows_checksum,
     write_rank_file,
     zigzag,
 )
@@ -78,12 +80,22 @@ def apply(names):
     assert torch.equal(patched["b"], as_bits(previous["b"]).view(-1))
 
 
-def test_anchor_with_an_unwritten_slice_is_not_published(tmp_path):
-    step_dir = tmp_path / "step_5"
-    step_dir.mkdir()
-    values = torch.arange(6, dtype=torch.bfloat16)
-    writer = AnchorWriter([("a", [6], 0, 6, torch.bfloat16), ("b", [6], 0, 6, torch.bfloat16)], step_dir, 0)
-    writer.write_parts([("a", [6], 0, values)], lambda tensor: [tensor.view(torch.uint8).numpy()])
-    with pytest.raises(RuntimeError, match="b@0"):
-        writer.close(1)
-    assert not anchor_complete(step_dir)
+def test_anchor_checksums_and_unwritten_slices(tmp_path):
+    values = torch.randn(12, dtype=torch.bfloat16)
+    slices = [("a", [12], 0, 6, torch.bfloat16), ("a", [12], 6, 6, torch.bfloat16)]
+    for step, written in (("step_5", slices), ("step_6", slices[:1])):
+        step_dir = tmp_path / step
+        step_dir.mkdir()
+        writer = AnchorWriter(slices, step_dir, 0)
+        for name, _, start, numel, _ in written:
+            part = values[start : start + numel]
+            writer.write_parts([(name, [12], start, part)], lambda tensor: [tensor.view(torch.uint8).numpy()])
+            writer.add_sums([name], rows_checksum(as_bits(part), start).view(1))
+        if written == slices:
+            writer.close(1)
+            assert anchor_complete(step_dir)
+            assert read_anchor_checksums(step_dir) == {"a": bits_checksum(as_bits(values))}
+        else:
+            with pytest.raises(RuntimeError, match="a@6"):
+                writer.close(1)
+            assert not anchor_complete(step_dir)

From 0ad0147a1d1d080aa5c8308f70824ada4e075cb8 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sun, 18 Oct 2026 10:36:15 +0000
Subject: [PATCH 58/88] fix(weights): diff packed qkv in the optimizer step
 instead of a copy of its gathered checkpoint entries

---
 src/prime_rl/transports/weights/sparse.py | 35 ++++++++++++++++++++---
 1 file changed, 31 insertions(+), 4 deletions(-)

diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 354e4c2732d..675a3dc3c51 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -18,6 +18,7 @@
 from torch.distributed.tensor._utils import compute_local_shape_and_global_offset
 
 from prime_rl.configs.shared import FileSystemWeightBroadcastConfig
+from prime_rl.trainer.models.fusions import get_model_packed_parameters, join_module_path
 from prime_rl.transports.weights.base import WeightSender
 from prime_rl.transports.weights.filesystem import FileSystemWeightSender
 from prime_rl.utils.logger import get_logger
@@ -45,6 +46,25 @@
     from prime_rl.trainer.optim import OptimizerLike
 
 
+def trainer_tensors(model: nn.Module) -> dict[str, torch.Tensor]:
+    """The model's state-dict entries without running its state-dict hooks (no collectives). A packed parameter
+    (``fusions``) whose packing dimension is sharded stays packed: its checkpoint entries would be gathered copies."""
+    tensors = dict(model.named_parameters(remove_duplicate=False))
+    for name, module in model.named_modules(remove_duplicate=False):
+        for buffer_name, buffer in module.named_buffers(recurse=False):
+            if buffer_name not in module._non_persistent_buffers_set:
+                tensors[join_module_path(name, buffer_name)] = buffer
+    keys = {id(tensor): key for key, tensor in tensors.items()}
+    for packed in get_model_packed_parameters(model):
+        param, spec = packed.parameter, packed.spec
+        if isinstance(param, DTensor) and any(placement.is_shard(spec.dim) for placement in param.placements):
+            continue
+        prefix = keys[id(param)].removesuffix(spec.name)
+        del tensors[keys[id(param)]]
+        tensors.update(zip([prefix + name for name in spec.logical_names], spec.split_into_logical_views(param)))
+    return tensors
+
+
 class SparseDeltaEncoder:
     """Rank-local deltas of the trainer's weight shards, in HF coordinates.
 
@@ -78,16 +98,23 @@ def __init__(self, model: nn.Module):
         self.wire_dtypes: dict[str, torch.dtype] = {}
         self.key_fqn: dict[str, str] = {}
         meta: dict[str, torch.Tensor] = {}
-        for key, value in model.state_dict().items():
+        for key, value in trainer_tensors(model).items():
             dtype = resolve_wire_dtype(keep_in_fp32, key, torch.bfloat16) if isinstance(value, DTensor) else value.dtype
             self.wire_dtypes[key] = dtype
             self.key_fqn[key] = resolve_fqn(model, key)
             meta[self.key_fqn[key]] = torch.empty(value.shape, dtype=dtype, device="meta")
         owner = {id(tensor): fqn for fqn, tensor in meta.items()}
+        # The checkpoint names of packed parameters are views of them
+        canonical = dict(meta)
+        for packed in get_model_packed_parameters(model):
+            if packed.fqn not in meta:
+                continue
+            views = packed.spec.split_into_logical_views(canonical.pop(packed.fqn))
+            canonical.update(zip(packed.logical_fqns, views))
 
         segments = defaultdict(list)
         self.hf_tensors: dict[str, tuple[list[int], str]] = {}
-        for name, tensor in convert_state_dict_to_hf(model, dict(meta)).items():
+        for name, tensor in convert_state_dict_to_hf(model, canonical).items():
             base = tensor if tensor._base is None else tensor._base
             if id(base) not in owner or not tensor.is_contiguous():
                 raise NotImplementedError(f"Sparse weight broadcast cannot map {name} to a contiguous trainer slice")
@@ -176,7 +203,7 @@ def _local_shard(self, key: str, value: torch.Tensor, is_master: bool) -> tuple[
     def _find_gathered(self, model: nn.Module, is_master: bool) -> set[str]:
         """fp8 trainer tensors whose 128-row blocks straddle ranks on some rank."""
         keys, straddles = [], []
-        for key, value in model.state_dict().items():
+        for key, value in trainer_tensors(model).items():
             fqn = resolve_fqn(model, key)
             if fqn not in self.fp8:
                 continue
@@ -201,7 +228,7 @@ def local_shards(self, model: nn.Module, is_master: bool, gather: bool = False)
         ``(previous, current)`` full bf16 tensors of fp8 tensors whose blocks straddle ranks, all-gathered here
         because ``encode`` runs no collectives."""
         shards = []
-        for key, value in model.state_dict().items():
+        for key, value in trainer_tensors(model).items():
             fqn = self.key_fqn[key]
             if fqn not in self.segments:
                 continue

From 31407876461e581a55bdd18b9071580b73a6936e Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sun, 18 Oct 2026 13:59:46 +0000
Subject: [PATCH 59/88] perf(weights): diff bf16 updates in place with
 sync-free kernels, without copies on the GPU

---
 src/prime_rl/trainer/optim/base.py            |   2 +-
 src/prime_rl/transports/weights/sparse.py     | 345 ++++++++++++------
 .../transports/weights/sparse_kernels.py      | 153 ++++++++
 3 files changed, 397 insertions(+), 103 deletions(-)
 create mode 100644 src/prime_rl/transports/weights/sparse_kernels.py

diff --git a/src/prime_rl/trainer/optim/base.py b/src/prime_rl/trainer/optim/base.py
index a85a1e3a2dc..b0f002c87a8 100644
--- a/src/prime_rl/trainer/optim/base.py
+++ b/src/prime_rl/trainer/optim/base.py
@@ -15,7 +15,7 @@ class ChunkedUpdates:
     """Optimizer mixin: with ``update_hook`` set, ``step`` updates the parameters a chunk of up to ``CHUNK_NUMEL``
     local elements at a time and calls the hook around each chunk's update."""
 
-    CHUNK_NUMEL = 1 << 28
+    CHUNK_NUMEL = 1 << 26
     update_hook: UpdateHook | None = None
 
     def step(self, closure=None):
diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 675a3dc3c51..05e0e28e496 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -11,6 +11,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING
 
+import numpy as np
 import torch
 import torch.distributed as dist
 import torch.nn as nn
@@ -21,6 +22,7 @@
 from prime_rl.trainer.models.fusions import get_model_packed_parameters, join_module_path
 from prime_rl.transports.weights.base import WeightSender
 from prime_rl.transports.weights.filesystem import FileSystemWeightSender
+from prime_rl.transports.weights.sparse_kernels import Pieces, build_pieces, count_changes, write_changes
 from prime_rl.utils.logger import get_logger
 from prime_rl.utils.sparse_weights import (
     DELTA_DIR,
@@ -69,12 +71,13 @@ class SparseDeltaEncoder:
     """Rank-local deltas of the trainer's weight shards, in HF coordinates.
 
     The deltas of the optimizer's parameters come from its step (``bind``): right before the step updates a chunk of
-    parameters, their shards still hold the previous broadcast's weights. ``update_hook`` (optimizers on the GPU)
-    copies the chunk's old wire bits, ``optimizer_hook`` (``FullCPUOffloadOptimizer``) receives the new values next
-    to the old ones; a diff thread compares the old and new wire bits and streams the changes out as the step goes.
-    No copy of the weights is kept. The other tensors (the residual: buffers, frozen parameters, state-dict copies,
-    fp8 tensors whose 128-row blocks straddle ranks) are diffed in ``encode`` against a copy of their previous wire
-    bits, which also checks that the optimizer's tensors still hold what its step wrote.
+    parameters, their shards still hold the previous broadcast's weights. ``update_hook`` (optimizers on the GPU) and
+    ``optimizer_hook`` (``FullCPUOffloadOptimizer``) copy the chunk's old wire bits; right after the update,
+    ``sparse_kernels`` compare them with the new bits in place and write the changes to pinned memory without a host
+    sync (bf16; fp8 and fp32 tensors are diffed on the diff thread), and the diff thread streams them out as the step
+    goes. No copy of the weights is kept. The other tensors (the residual: buffers, frozen parameters, fp8 tensors
+    whose 128-row blocks straddle ranks) are diffed in ``encode`` against a copy of their previous wire bits, which
+    also checks that the optimizer's tensors still hold what its step wrote.
 
     The prime->HF conversion is replayed once on meta tensors: every HF tensor must be a contiguous view of one trainer
     tensor, which maps a changed trainer element to its HF name and flat index. Tensors the engines hold in fp8 are
@@ -88,10 +91,13 @@ class SparseDeltaEncoder:
     PART = 1 << 26
     HOOK_CHUNK = 1 << 24
     HOOK_PART = 1 << 24
-    # GPU slots of ``optimizer_hook``, each twice the largest parameter shard; bytes of old values that
-    # ``update_hook`` holds for the diff thread (at least one update's)
+    # Updates in flight: GPU slots of their old bf16 bits (full offload: each the size of the largest parameter
+    # shard) and pinned regions of their changes; bytes of other old values that ``update_hook`` holds for the diff
+    # thread (at least one update's)
     HOOK_SLOTS = 2
     UPDATE_BYTES = 1 << 32
+    # Changes per pinned region of the kernels (more take further rounds)
+    REGION = 1 << 24
 
     def __init__(self, model: nn.Module):
         keep_in_fp32 = getattr(model, "keep_in_fp32_for_weight_transfer", None)
@@ -448,15 +454,17 @@ def flush() -> None:
 
         for key, fqn, shard, offset, gathered in shards:
             shard = shard.reshape(-1)
-            if key in self.checksums:
-                bits = as_bits(shard.to(self.wire_dtypes[key]))
-                checks.append((key, self._checksum(key, offset, bits)))
+            if key in self.fast or key in self.checksums:
                 if key not in self.stepped:
-                    # Not updated in this step (e.g. no gradient): an empty part carries its checksums, and its anchor
-                    # slices come from the live weights
-                    pending += self._diff_shard(key, fqn, shard, offset, None, bits, None)
+                    # Not updated in this step (e.g. no gradient): its anchor slices come from the live weights, and
+                    # an empty part carries its checksums (the bf16 tensors' come with the last part)
+                    bits = as_bits(shard.to(self.wire_dtypes[key]))
+                    if key in self.checksums:
+                        pending += self._diff_shard(key, fqn, shard, offset, None, bits, None)
                     if self.anchor is not None:
                         self._write_anchor_slices(key, fqn, len(bits), offset, bits)
+                if key in self.checksums:
+                    checks.append((key, self._checksum(key, offset, as_bits(shard.to(self.wire_dtypes[key])))))
                 continue
             parts = self._diff_shard(key, fqn, shard, offset, gathered, self.previous[key], self.previous[key])
             pending += parts
@@ -466,12 +474,17 @@ def flush() -> None:
                 num_pending = 0
         if pending:
             flush()
+        names = [name for pieces in self.all_pieces for name in pieces.names]
+        ids = [self.name_ids[name] for name in names]
+        sums, live = self.sums[ids].tolist(), self._live_sums()[ids].tolist()
+        self._emit_host_parts(emit, dict(zip(names, sums)))
+        changed = {name for name, expected, value in zip(names, sums, live) if expected != value}
         if checks:
             keys, live = zip(*checks)
             expected = torch.cat([self.checksums[key] for key in keys])
-            changed = [key for key, same in zip(keys, expected.eq(torch.cat(live)).tolist()) if not same]
-            if changed:
-                raise RuntimeError(f"Trainer weights changed outside the optimizer step, not sent: {changed}")
+            changed |= {key for key, same in zip(keys, expected.eq(torch.cat(live)).tolist()) if not same}
+        if changed:
+            raise RuntimeError(f"Trainer weights changed outside the optimizer step, not sent: {sorted(changed)}")
 
     def _host_part(self, pending: list[tuple], staging: torch.Tensor | None = None) -> tuple:
         """One part ``(entries, gaps, value bytes, checksums)`` on the host from GPU part entries."""
@@ -620,12 +633,33 @@ def bind(self, model: nn.Module, is_master: bool, params: list[torch.Tensor], fu
             else:
                 residual.append(key)
         self.residual = set(residual)
+        # Diffed by ``sparse_kernels``: the tensors of parameters that are all bf16 on the wire. The others (fp8,
+        # fp32) by ``_diff_slot``
+        self.fast = {
+            key
+            for entries in self.param_keys.values()
+            if all(self.wire_dtypes[key] == torch.bfloat16 and fqn not in self.fp8 for key, fqn, *_ in entries)
+            for key, *_ in entries
+        }
+        self.params = params
+        self.pieces: dict[tuple, Pieces] = {}
+        self.name_ids = {name: index for index, name in enumerate(self.hf_tensors)}
         device = f"cuda:{self.device}"
-        if full_offload:
-            slot_bytes = 2 * -(-max(param.nbytes for param in params) // 256) * 256
-            self.hook_slots = queue.SimpleQueue()
-            for _ in range(self.HOOK_SLOTS):
-                self.hook_slots.put(torch.empty(slot_bytes, dtype=torch.uint8, device=device))
+        # The bf16 tensors' checksums of their last written bits, by ``name_ids``
+        self.sums = torch.zeros(len(self.name_ids), dtype=torch.int64, device=device)
+        # Old values of the updates in flight (full offload: one parameter each; grown as needed otherwise), and
+        # the pinned regions the kernels write the changes to
+        slot_bytes = -(-max(param.nbytes for param in params) // 256) * 256 if full_offload else 0
+        self.slots, self.regions = queue.SimpleQueue(), queue.SimpleQueue()
+        for _ in range(self.HOOK_SLOTS):
+            self.slots.put(torch.empty(slot_bytes, dtype=torch.uint8, device=device))
+            self.regions.put(
+                (
+                    torch.empty(self.REGION, dtype=torch.int32, pin_memory=True),
+                    torch.empty(self.REGION, dtype=torch.int16, pin_memory=True),
+                    torch.empty(0, dtype=torch.int32, pin_memory=True),
+                )
+            )
         self.update_bytes = threading.Condition()
         self.held_bytes = 0
         self.differ = ThreadPoolExecutor(1)
@@ -642,18 +676,30 @@ def reset(self, model: nn.Module, is_master: bool) -> None:
             if key in self.residual:
                 self.previous[key] = self._bits(key, shard)
                 self.shards[key] = (shard.numel(), offset)
-            else:
+            elif key not in self.fast:
                 self.checksums[key] = self._checksum(key, offset, as_bits(shard.reshape(-1).to(self.wire_dtypes[key])))
+        self.all_pieces = self._group_pieces(
+            [(param, [entry for entry in self._entries(param) if entry[0] in self.fast], 0) for param in self.params]
+        )
+        self.sums = self._live_sums()
         num_bytes = sum(bits.nbytes for bits in self.previous.values())
         get_logger().info(
-            f"Sparse weight deltas of {len(self.checksums)} tensors come from the optimizer step, "
+            f"Sparse weight deltas of {len(self.checksums) + len(self.fast)} tensors come from the optimizer step, "
             f"{len(self.previous)} tensors ({num_bytes / 2**30:.2f} GiB) are diffed against a copy"
         )
 
+    def _live_sums(self) -> torch.Tensor:
+        """The checksums of the bf16 optimizer tensors' live bits, by ``name_ids``."""
+        sums = torch.zeros_like(self.sums)
+        for pieces in self.all_pieces:
+            count_changes(pieces, None, torch.empty_like(pieces.tables[5]), sums)
+        return sums
+
     def arm(self, anchor: AnchorWriter | None, emit: Callable) -> None:
         """Collect the next optimizer step's delta, handing host parts to ``emit`` as the step goes, and write its
         anchor slices if ``anchor``."""
         self.step_parts, self.scale_changes, self.anchor, self.emit = [], 0, anchor, emit
+        self.host_parts: list[tuple] = []
         # Optimizer tensors the step updated (and wrote the anchor slices of)
         self.stepped: set[str] = set()
         self.hook_wait = self.diff_busy = 0.0
@@ -685,26 +731,151 @@ def _entries(self, local_param: torch.Tensor) -> list[tuple]:
             if start <= at < start + local_param.numel()
         ]
 
+    def _group_pieces(self, items: list[tuple]) -> list[Pieces]:
+        """The ``Pieces`` of ``items`` (``(local parameter, entries, position of its old bits)``), one per
+        parameter dtype; cached."""
+        cache_key = tuple((local.data_ptr(), position) for local, _, position in items)
+        if cache_key not in self.pieces:
+            rows = defaultdict(list)
+            for local, entries, position in items:
+                for key, fqn, numel, offset, (size, stride, at) in entries:
+                    names, begins, ends, shifts = self._segments(key, self.segments[fqn], offset, numel)
+                    for name, begin, end, shift in zip(names, begins, ends, shifts):
+                        # A segment is contiguous in the parameter, also in a strided view
+                        first, last = (
+                            at + sum(i * s for i, s in zip(np.unravel_index(j, size), stride)) for j in (begin, end - 1)
+                        )
+                        assert last - first == end - 1 - begin
+                        rows[local.dtype].append(
+                            (
+                                name,
+                                position + first,
+                                local.data_ptr() + first * local.element_size(),
+                                end - begin,
+                                begin + shift,
+                                self.name_ids[name],
+                            )
+                        )
+            self.pieces[cache_key] = [
+                build_pieces(group, dtype == torch.float32, local.device) for dtype, group in rows.items()
+            ]
+        return self.pieces[cache_key]
+
+    def _launch(self, items: list[tuple], old: torch.Tensor, release: Callable, written: list[tuple]) -> None:
+        """Right after the update of ``items`` (``(local parameter, entries, position of its old bits in old``),
+        on the current stream: count and write its changes to pinned regions and record the new bits' checksums;
+        the diff thread collects them, writes the anchor slices of ``written`` and then calls ``release``."""
+        launches = []
+        for pieces in self._group_pieces(items):
+            indices, values, counts_host = self.regions.get()
+            num_blocks = len(pieces.tables[5])
+            if len(counts_host) < num_blocks:
+                counts_host = torch.empty(num_blocks, dtype=torch.int32, pin_memory=True)
+            counts = torch.empty(num_blocks, dtype=torch.int32, device=old.device)
+            self.sums.index_fill_(0, pieces.segments, 0)
+            count_changes(pieces, old, counts, self.sums)
+            ends = counts.cumsum(0)
+            write_changes(pieces, old, ends, indices, values, 0)
+            counts_host[:num_blocks].copy_(counts, non_blocking=True)
+            launches.append((pieces, ends, (indices, values, counts_host)))
+        ready = torch.cuda.Event()
+        ready.record()
+        keys = [key for _, entries, _ in items for key, *_ in entries]
+        self.diffs.append(self.differ.submit(self._collect, launches, keys, old, ready, release, written))
+
+    @torch.no_grad()
+    def _collect(self, launches: list[tuple], keys: list[str], old: torch.Tensor, ready, release, written) -> None:
+        """Diff thread: the host parts of ``_launch``'s changes to ``keys`` (writing the changes past a region's
+        capacity in further rounds), and the anchor slices of ``written``; then ``release`` the old values."""
+        torch.cuda.set_device(self.device)
+        ready.synchronize()
+        start = time.perf_counter()
+        try:
+            with torch.cuda.stream(self.diff_stream):
+                for pieces, ends, region in launches:
+                    indices, values, counts = region
+                    block_counts = counts[: len(pieces.tables[5])].numpy()
+                    total = int(block_counts.sum())
+                    changed, changed_values = [], []
+                    for skip in range(0, total, len(indices)):
+                        if skip:
+                            write_changes(pieces, old, ends, indices, values, skip)
+                            self.diff_stream.synchronize()
+                        num = min(len(indices), total - skip)
+                        changed.append(indices[:num].numpy().copy())
+                        changed_values.append(values[:num].numpy().copy())
+                    self.regions.put(region)
+                    if not total:
+                        continue
+                    piece_counts = np.add.reduceat(block_counts, pieces.first_blocks)
+                    index = np.concatenate(changed)
+                    previous = np.empty_like(index)
+                    previous[1:] = index[:-1]
+                    previous[(np.cumsum(piece_counts) - piece_counts)[piece_counts > 0]] = -1
+                    entries = [(name, int(count)) for name, count in zip(pieces.names, piece_counts) if count]
+                    self.host_parts.append((entries, index - previous - 1, np.concatenate(changed_values)))
+                if self.anchor is not None:
+                    for key, fqn, numel, offset, view in written:
+                        self._write_anchor_slices(key, fqn, numel, offset, view.reshape(-1).to(self.wire_dtypes[key]))
+                    if not written:  # written by optimizer_hook
+                        for pieces, *_ in launches:
+                            self.anchor.add_sums(pieces.names, self.sums[pieces.segments])
+                    self.diff_stream.synchronize()
+                self.stepped.update(keys)
+        finally:
+            release()
+            self.diff_busy += time.perf_counter() - start
+        if sum(len(part[1]) for part in self.host_parts) >= self.HOOK_PART:
+            self._emit_host_parts(self.emit)
+
+    def _emit_host_parts(self, emit: Callable, checksums: dict[str, int] | None = None) -> None:
+        parts, self.host_parts = self.host_parts, []
+        gaps = np.concatenate([part[1] for part in parts]) if parts else np.empty(0, np.int32)
+        values = np.concatenate([part[2] for part in parts]) if parts else np.empty(0, np.int16)
+        entries = [entry for part in parts for entry in part[0]]
+        emit((entries, torch.from_numpy(gaps), torch.from_numpy(values.view(np.uint8)), checksums or {}))
+
     @torch.no_grad()
     def update_hook(self, params: list[torch.Tensor]) -> Callable[[], None]:
         """``update_hook`` of the optimizers on the GPU, right before they update ``params`` on the current stream:
-        copy the old wire bits of their shards, and diff them against the updated shards on the diff thread."""
+        copy the old wire bits of their shards, and diff them against the updated shards (bf16: right after the
+        update, with ``sparse_kernels``; others on the diff thread)."""
         if not self.armed:
             return lambda: None
-        locals_ = [param.to_local() if isinstance(param, DTensor) else param for param in params]
-        num_bytes = sum(local.numel() for local in locals_) * 2
-        start = time.perf_counter()
-        with self.update_bytes:
-            self.update_bytes.wait_for(lambda: self.held_bytes == 0 or self.held_bytes + num_bytes <= self.UPDATE_BYTES)
-            self.held_bytes += num_bytes
-        self.hook_wait += time.perf_counter() - start
-        items = []
-        for local in locals_:
+        fast, other = [], []
+        for param in params:
+            local = param.to_local() if isinstance(param, DTensor) else param
             entries = self._entries(local)
             if entries:
-                dtypes = {self.wire_dtypes[key] for key, *_ in entries}
-                dtype = dtypes.pop() if len(dtypes) == 1 else local.dtype
-                items.append((entries, local.to(dtype, copy=True), local))
+                (fast if entries[0][0] in self.fast else other).append((local, entries))
+        start = time.perf_counter()
+        num_bytes = sum(local.numel() for local, _ in other) * 2
+        if other:
+            with self.update_bytes:
+                self.update_bytes.wait_for(
+                    lambda: self.held_bytes == 0 or self.held_bytes + num_bytes <= self.UPDATE_BYTES
+                )
+                self.held_bytes += num_bytes
+        if fast:
+            slot = self.slots.get()
+            numel = sum(local.numel() for local, _ in fast)
+            if slot.numel() < 2 * numel:
+                slot = torch.empty(2 * numel, dtype=torch.uint8, device=slot.device)
+            old, items, written, position = slot.view(torch.bfloat16), [], [], 0
+            for local, entries in fast:
+                old[position : position + local.numel()].view(local.shape).copy_(local)
+                items.append((local, entries, position))
+                written += [
+                    (key, fqn, numel, offset, local.as_strided(size, stride, local.storage_offset() + at))
+                    for key, fqn, numel, offset, (size, stride, at) in entries
+                ]
+                position += local.numel()
+        self.hook_wait += time.perf_counter() - start
+        items_other = []
+        for local, entries in other:
+            dtypes = {self.wire_dtypes[key] for key, *_ in entries}
+            dtype = dtypes.pop() if len(dtypes) == 1 else local.dtype
+            items_other.append((entries, local.to(dtype, copy=True), local))
 
         def release() -> None:
             with self.update_bytes:
@@ -712,31 +883,31 @@ def release() -> None:
                 self.update_bytes.notify_all()
 
         def finish() -> None:
-            ready = torch.cuda.Event()
-            ready.record()
-            self.diffs.append(self.differ.submit(self._diff_slot, items, ready, release, True))
+            if fast:
+                self._launch(items, slot.view(torch.int16), partial(self.slots.put, slot), written)
+            if other:
+                ready = torch.cuda.Event()
+                ready.record()
+                self.diffs.append(self.differ.submit(self._diff_slot, items_other, ready, release, True))
 
         return finish
 
     @torch.no_grad()
     def optimizer_hook(self, local_param: torch.Tensor, source: torch.Tensor) -> None:
         """``FullCPUOffloadOptimizer.weight_hook``, on its H2D stream: copy the shard (still the previous
-        broadcast's weights) and its new values ``source`` (pinned host) into a GPU slot, write the new values and
-        diff the slot on the diff thread. At anchor steps, also writes the new values to the anchor from the
-        host."""
+        broadcast's weights) into a GPU slot, write the new values ``source`` (pinned host) and diff the shard
+        against the slot (bf16: right away, with ``sparse_kernels``; others on the diff thread). At anchor steps,
+        also writes the new values to the anchor from the host."""
         entries = self._entries(local_param) if self.armed else []
         if not entries:
             local_param.copy_(source, non_blocking=True)
             return
-        slot = self.hook_slots.get()
-        half = slot.numel() // 2
-        previous = slot[: source.nbytes].view(source.dtype)
-        new = slot[half : half + source.nbytes].view(source.dtype)
-        previous.copy_(local_param.reshape(-1))
-        new.copy_(source.reshape(-1), non_blocking=True)
-        local_param.copy_(new.view(local_param.shape))
-        ready = torch.cuda.Event()
-        ready.record()
+        slot = self.slots.get()
+        fast = entries[0][0] in self.fast
+        previous = slot[: local_param.numel() * (2 if fast else local_param.element_size())]
+        previous = previous.view(torch.bfloat16 if fast else local_param.dtype)
+        previous.view(local_param.shape).copy_(local_param)
+        local_param.copy_(source, non_blocking=True)
         if self.anchor is not None:
             for key, fqn, numel, offset, (size, stride, at) in entries:
                 host = source.as_strided(size, stride, source.storage_offset() + at).reshape(-1)
@@ -747,71 +918,41 @@ def optimizer_hook(self, local_param: torch.Tensor, source: torch.Tensor) -> Non
                         for name, _, hf_start, begin, end in self._slices(fqn, numel, offset)
                     ]
                 )
-        release = partial(self.hook_slots.put, slot)
-        self.diffs.append(self.differ.submit(self._diff_slot, [(entries, previous, new)], ready, release, False))
+        release = partial(self.slots.put, slot)
+        if fast:
+            self._launch([(local_param, entries, 0)], previous.view(torch.int16), release, [])
+            return
+        ready = torch.cuda.Event()
+        ready.record()
+        # The shard holds the new values until the next step, which waits for the diffs
+        items = [(entries, previous, local_param.reshape(-1))]
+        self.diffs.append(self.differ.submit(self._diff_slot, items, ready, release, False))
 
     @torch.no_grad()
     def _diff_slot(self, items: list[tuple], ready: torch.cuda.Event, release: Callable, on_gpu: bool) -> None:
         """Diff thread: the parts of the trainer tensors in ``items`` (``(entries, old values, new values)`` of
-        parameter shards), and their anchor slices if the optimizer is ``on_gpu``; then ``release`` the old values.
-        The bf16 tensors are diffed together, as one flat tensor. Hands the parts so far to ``emit`` every
-        ``HOOK_PART`` changes."""
+        parameter shards that are not all bf16), and their anchor slices if the optimizer is ``on_gpu``; then
+        ``release`` the old values. Hands the parts so far to ``emit`` every ``HOOK_PART`` changes."""
         torch.cuda.set_device(self.device)
         ready.synchronize()
         start = time.perf_counter()
         chunk = self.CHUNK if on_gpu else self.HOOK_CHUNK
         try:
             with torch.cuda.stream(self.diff_stream):
-                batch, written = [], []
                 for entries, previous, new in items:
                     for key, fqn, numel, offset, (size, stride, at) in entries:
                         current, old = (
                             values.as_strided(size, stride, values.storage_offset() + at) for values in (new, previous)
                         )
-                        if fqn in self.fp8 or self.wire_dtypes[key] != torch.bfloat16:
-                            current = current.reshape(-1).to(self.wire_dtypes[key])
-                            old = as_bits(old.reshape(-1).to(self.wire_dtypes[key]))
-                            self.step_parts += self._diff_shard(key, fqn, current, offset, None, old, None, chunk)
-                            self.checksums[key] = self._checksum(key, offset, as_bits(current))
-                            written.append((key, fqn, numel, offset, current))
-                        else:
-                            batch.append((key, fqn, numel, offset, current, old))
-                if batch:
-                    total = sum(numel for _, _, numel, *_ in batch)
-                    current_bits, previous_bits = (
-                        torch.empty(total, dtype=torch.bfloat16, device=f"cuda:{self.device}") for _ in range(2)
-                    )
-                    names, begins, ends, shifts, ranges, position = [], [], [], [], [], 0
-                    for key, fqn, numel, offset, current, old in batch:
-                        current_bits[position : position + numel].view(current.shape).copy_(current)
-                        previous_bits[position : position + numel].view(old.shape).copy_(old)
-                        segment_names, *segments = self._segments(key, self.segments[fqn], offset, numel)
-                        ranges.append((key, len(names), len(names) + len(segment_names)))
-                        names += segment_names
-                        begins += [begin + position for begin in segments[0]]
-                        ends += [end + position for end in segments[1]]
-                        shifts += [shift - position for shift in segments[2]]
-                        written.append((key, fqn, numel, offset, current_bits[position : position + numel]))
-                        position += numel
-                    part = self._diff_flat(
-                        tuple(key for key, *_ in batch),
-                        names,
-                        begins,
-                        ends,
-                        shifts,
-                        as_bits(current_bits),
-                        as_bits(previous_bits),
-                        chunk,
-                    )
-                    self.step_parts.append(part)
-                    for key, first, last in ranges:
-                        self.checksums[key] = part[5][first:last].sum(0, keepdim=True)
-                for key, fqn, numel, offset, values in written if self.anchor is not None else []:
-                    if on_gpu:
-                        self._write_anchor_slices(key, fqn, numel, offset, values)
-                    else:  # written by optimizer_hook
-                        self.anchor.add_sums(*self._value_sums(key, offset, as_bits(values)))
-                self.stepped.update(key for key, *_ in written)
+                        current = current.reshape(-1).to(self.wire_dtypes[key])
+                        old = as_bits(old.reshape(-1).to(self.wire_dtypes[key]))
+                        self.step_parts += self._diff_shard(key, fqn, current, offset, None, old, None, chunk)
+                        self.checksums[key] = self._checksum(key, offset, as_bits(current))
+                        if self.anchor is not None and on_gpu:
+                            self._write_anchor_slices(key, fqn, numel, offset, current)
+                        elif self.anchor is not None:  # written by optimizer_hook
+                            self.anchor.add_sums(*self._value_sums(key, offset, as_bits(current)))
+                        self.stepped.add(key)
                 self.diff_stream.synchronize()
         finally:
             release()
diff --git a/src/prime_rl/transports/weights/sparse_kernels.py b/src/prime_rl/transports/weights/sparse_kernels.py
new file mode 100644
index 00000000000..5e80e5d465f
--- /dev/null
+++ b/src/prime_rl/transports/weights/sparse_kernels.py
@@ -0,0 +1,153 @@
+"""Triton kernels that diff a parameter update in place: no host sync, no workspace beyond the old values."""
+
+from typing import NamedTuple
+
+import numpy as np
+import torch
+import triton
+import triton.language as tl
+
+BLOCK = 4096
+
+
+@triton.jit
+def _load(
+    old_base, old_offsets, addresses, lengths, block_pieces, block_starts, NEW_FP32: tl.constexpr, SAME: tl.constexpr
+):
+    pid = tl.program_id(0)
+    piece = tl.load(block_pieces + pid)
+    offsets = tl.load(block_starts + pid) + tl.arange(0, 4096)
+    mask = offsets < tl.load(lengths + piece)
+    if NEW_FP32:
+        new = tl.load(tl.load(addresses + piece).to(tl.pointer_type(tl.float32)) + offsets, mask=mask, other=0)
+        new = new.to(tl.bfloat16).to(tl.int16, bitcast=True)
+    else:
+        new = tl.load(tl.load(addresses + piece).to(tl.pointer_type(tl.int16)) + offsets, mask=mask, other=0)
+    if SAME:
+        old = new
+    else:
+        old = tl.load(old_base + tl.load(old_offsets + piece) + offsets, mask=mask, other=0)
+    return pid, piece, offsets, mask, old, new
+
+
+@triton.jit
+def _count_kernel(
+    old_base,
+    old_offsets,
+    addresses,
+    lengths,
+    firsts,
+    segments,
+    block_pieces,
+    block_starts,
+    counts,
+    sums,
+    NEW_FP32: tl.constexpr,
+    SAME: tl.constexpr,
+):
+    pid, piece, offsets, mask, old, new = _load(
+        old_base, old_offsets, addresses, lengths, block_pieces, block_starts, NEW_FP32, SAME
+    )
+    tl.store(counts + pid, tl.sum(((old != new) & mask).to(tl.int32)))
+    weights = tl.load(firsts + piece) + offsets + 1
+    tl.atomic_add(sums + tl.load(segments + piece), tl.sum(tl.where(mask, weights * new.to(tl.int64), 0)))
+
+
+@triton.jit(do_not_specialize=["skip", "capacity"])
+def _write_kernel(
+    old_base,
+    old_offsets,
+    addresses,
+    lengths,
+    firsts,
+    block_pieces,
+    block_starts,
+    ends,
+    indices,
+    values,
+    skip,
+    capacity,
+    NEW_FP32: tl.constexpr,
+):
+    pid, piece, offsets, mask, old, new = _load(
+        old_base, old_offsets, addresses, lengths, block_pieces, block_starts, NEW_FP32, False
+    )
+    changed = ((old != new) & mask).to(tl.int32)
+    position = tl.load(ends + pid) - tl.sum(changed) + tl.cumsum(changed, 0) - changed - skip
+    store = (changed != 0) & (position >= 0) & (position < capacity)
+    tl.store(indices + position, (tl.load(firsts + piece) + offsets).to(tl.int32), mask=store)
+    delta = new - old
+    tl.store(values + position, (delta << 1) ^ (delta >> 15), mask=store)
+
+
+class Pieces(NamedTuple):
+    """Contiguous runs of bf16 wire bits, each within one HF tensor: ``names``, the first block of each, and the
+    device tables of the kernels."""
+
+    names: list[str]
+    first_blocks: np.ndarray
+    new_fp32: bool
+    segments: torch.Tensor
+    tables: tuple[torch.Tensor, ...]
+
+
+def build_pieces(rows: list[tuple[str, int, int, int, int, int]], new_fp32: bool, device) -> Pieces:
+    """``rows`` of ``(hf_name, offset in the old values, address of the new values, length, HF flat index of the
+    first element, segment)``."""
+    names, old_offsets, addresses, lengths, firsts, segments = zip(*rows)
+    num_blocks = [-(-length // BLOCK) for length in lengths]
+    first_blocks = np.cumsum([0, *num_blocks[:-1]])
+    block_pieces = np.repeat(np.arange(len(rows), dtype=np.int32), num_blocks)
+    block_starts = (np.arange(len(block_pieces)) - first_blocks[block_pieces]) * BLOCK
+
+    def table(values, dtype=torch.int64):
+        return torch.tensor(np.asarray(values), dtype=dtype).to(device)
+
+    segments = table(segments)
+    tables = (table(old_offsets), table(addresses), table(lengths), table(firsts), segments)
+    return Pieces(
+        list(names), first_blocks, new_fp32, segments, (*tables, table(block_pieces, torch.int32), table(block_starts))
+    )
+
+
+def count_changes(pieces: Pieces, old: torch.Tensor | None, counts: torch.Tensor, sums: torch.Tensor) -> None:
+    """Changed elements per block into ``counts``, and add each piece's ``bits_checksum`` of the new values to
+    ``sums[segment]``; without ``old``, only the checksums."""
+    old_offsets, addresses, lengths, firsts, segments, block_pieces, block_starts = pieces.tables
+    _count_kernel[(len(block_pieces),)](
+        addresses if old is None else old,
+        old_offsets,
+        addresses,
+        lengths,
+        firsts,
+        segments,
+        block_pieces,
+        block_starts,
+        counts,
+        sums,
+        NEW_FP32=pieces.new_fp32,
+        SAME=old is None,
+    )
+
+
+def write_changes(
+    pieces: Pieces, old: torch.Tensor, ends: torch.Tensor, indices: torch.Tensor, values: torch.Tensor, skip: int
+) -> None:
+    """Changes ``skip`` to ``skip + len(indices)`` (``ends``: inclusive cumulative block counts): their HF flat
+    indices and ``zigzag`` bit differences."""
+    old_offsets, addresses, lengths, firsts, _, block_pieces, block_starts = pieces.tables
+    _write_kernel[(len(block_pieces),)](
+        old,
+        old_offsets,
+        addresses,
+        lengths,
+        firsts,
+        block_pieces,
+        block_starts,
+        ends,
+        indices,
+        values,
+        skip,
+        len(indices),
+        NEW_FP32=pieces.new_fp32,
+    )

From 6ce93773da8450c1e586f7a07a6676a2672bf478 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sun, 18 Oct 2026 15:36:19 +0000
Subject: [PATCH 60/88] perf(weights): write full-offload anchors from the live
 weights after the delta, not in the optimizer step

---
 src/prime_rl/transports/weights/sparse.py | 76 ++++++++++++-----------
 1 file changed, 39 insertions(+), 37 deletions(-)

diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 05e0e28e496..c0b2a170688 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -308,22 +308,40 @@ def write_full_checkpoint(self, model: nn.Module, is_master: bool, step_dir: Pat
 
     @torch.no_grad()
     def encode(
-        self, shards: list[tuple], ready: torch.cuda.Event, weights_read: threading.Event, emit: Callable
-    ) -> int:
+        self,
+        shards: list[tuple],
+        ready: torch.cuda.Event,
+        weights_read: threading.Event,
+        emit: Callable,
+        diffed: Callable[[int], None],
+    ) -> None:
         """After ``wait_for_diffs``, on the side stream after ``ready``: diff the residual shards against their
         previous bits and advance them, and check the optimizer's tensors. Hands host parts ``(entries, gaps, value
-        bytes, checksums)`` to ``emit`` as they are encoded and sets ``weights_read`` once done reading the weights.
-        Returns the number of changed fp8 block scales."""
+        bytes, checksums)`` to ``emit`` as they are encoded, then calls ``diffed`` with the number of changed fp8
+        block scales. At anchor steps, then writes the anchor slices the optimizer step did not write from the
+        live weights. Sets ``weights_read`` once done reading the weights."""
         torch.cuda.set_device(self.device)
         try:
             with torch.cuda.stream(self.stream):
                 self.stream.wait_event(ready)
                 self._diff(shards, emit)
                 self.stream.synchronize()
-            return self.scale_changes
+                diffed(self.scale_changes)
+                if self.anchor is not None:
+                    self._write_optimizer_anchor(shards)
+                    self.stream.synchronize()
         finally:
             weights_read.set()
 
+    def _write_optimizer_anchor(self, shards: list[tuple]) -> None:
+        """The anchor slices of the optimizer's tensors that the step did not write (with full offload, all of
+        them: writing them in the step stalls its gradient offload, and the gradients pile up on the GPU)."""
+        for key, fqn, shard, offset, _ in shards:
+            if key in self.residual or (key in self.stepped and not self.full_offload):
+                continue
+            values = shard.reshape(-1).to(self.wire_dtypes[key])
+            self._write_anchor_slices(key, fqn, len(values), offset, values)
+
     def wait_for_diffs(self) -> None:
         for diff in self.diffs:
             diff.result()
@@ -454,17 +472,15 @@ def flush() -> None:
 
         for key, fqn, shard, offset, gathered in shards:
             shard = shard.reshape(-1)
-            if key in self.fast or key in self.checksums:
+            if key in self.checksums:
+                bits = as_bits(shard.to(self.wire_dtypes[key]))
+                checks.append((key, self._checksum(key, offset, bits)))
                 if key not in self.stepped:
-                    # Not updated in this step (e.g. no gradient): its anchor slices come from the live weights, and
-                    # an empty part carries its checksums (the bf16 tensors' come with the last part)
-                    bits = as_bits(shard.to(self.wire_dtypes[key]))
-                    if key in self.checksums:
-                        pending += self._diff_shard(key, fqn, shard, offset, None, bits, None)
-                    if self.anchor is not None:
-                        self._write_anchor_slices(key, fqn, len(bits), offset, bits)
-                if key in self.checksums:
-                    checks.append((key, self._checksum(key, offset, as_bits(shard.to(self.wire_dtypes[key])))))
+                    # Not updated in this step (e.g. no gradient): an empty part carries its checksums (the bf16
+                    # tensors' come with the last part)
+                    pending += self._diff_shard(key, fqn, shard, offset, None, bits, None)
+                continue
+            if key in self.fast:
                 continue
             parts = self._diff_shard(key, fqn, shard, offset, gathered, self.previous[key], self.previous[key])
             pending += parts
@@ -641,7 +657,7 @@ def bind(self, model: nn.Module, is_master: bool, params: list[torch.Tensor], fu
             if all(self.wire_dtypes[key] == torch.bfloat16 and fqn not in self.fp8 for key, fqn, *_ in entries)
             for key, *_ in entries
         }
-        self.params = params
+        self.params, self.full_offload = params, full_offload
         self.pieces: dict[tuple, Pieces] = {}
         self.name_ids = {name: index for index, name in enumerate(self.hf_tensors)}
         device = f"cuda:{self.device}"
@@ -817,9 +833,6 @@ def _collect(self, launches: list[tuple], keys: list[str], old: torch.Tensor, re
                 if self.anchor is not None:
                     for key, fqn, numel, offset, view in written:
                         self._write_anchor_slices(key, fqn, numel, offset, view.reshape(-1).to(self.wire_dtypes[key]))
-                    if not written:  # written by optimizer_hook
-                        for pieces, *_ in launches:
-                            self.anchor.add_sums(pieces.names, self.sums[pieces.segments])
                     self.diff_stream.synchronize()
                 self.stepped.update(keys)
         finally:
@@ -896,8 +909,7 @@ def finish() -> None:
     def optimizer_hook(self, local_param: torch.Tensor, source: torch.Tensor) -> None:
         """``FullCPUOffloadOptimizer.weight_hook``, on its H2D stream: copy the shard (still the previous
         broadcast's weights) into a GPU slot, write the new values ``source`` (pinned host) and diff the shard
-        against the slot (bf16: right away, with ``sparse_kernels``; others on the diff thread). At anchor steps,
-        also writes the new values to the anchor from the host."""
+        against the slot (bf16: right away, with ``sparse_kernels``; others on the diff thread)."""
         entries = self._entries(local_param) if self.armed else []
         if not entries:
             local_param.copy_(source, non_blocking=True)
@@ -908,16 +920,6 @@ def optimizer_hook(self, local_param: torch.Tensor, source: torch.Tensor) -> Non
         previous = previous.view(torch.bfloat16 if fast else local_param.dtype)
         previous.view(local_param.shape).copy_(local_param)
         local_param.copy_(source, non_blocking=True)
-        if self.anchor is not None:
-            for key, fqn, numel, offset, (size, stride, at) in entries:
-                host = source.as_strided(size, stride, source.storage_offset() + at).reshape(-1)
-                host = host.to(self.wire_dtypes[key])
-                self.anchor.write_all(
-                    [
-                        (name, hf_start, host[begin:end].view(torch.uint8).numpy())
-                        for name, _, hf_start, begin, end in self._slices(fqn, numel, offset)
-                    ]
-                )
         release = partial(self.slots.put, slot)
         if fast:
             self._launch([(local_param, entries, 0)], previous.view(torch.int16), release, [])
@@ -950,8 +952,6 @@ def _diff_slot(self, items: list[tuple], ready: torch.cuda.Event, release: Calla
                         self.checksums[key] = self._checksum(key, offset, as_bits(current))
                         if self.anchor is not None and on_gpu:
                             self._write_anchor_slices(key, fqn, numel, offset, current)
-                        elif self.anchor is not None:  # written by optimizer_hook
-                            self.anchor.add_sums(*self._value_sums(key, offset, as_bits(current)))
                         self.stepped.add(key)
                 self.diff_stream.synchronize()
         finally:
@@ -1183,7 +1183,12 @@ def emit(part: tuple) -> None:
             if self.world.is_master:
                 self._submit(self.writes, self._write_manifest, step_dir, step, base_step, self.world.world_size)
             self._submit(self.writes, self._move_streamed, self.streamed_dir, step_dir / DELTA_DIR)
-            scale_changes = encoder.encode(shards, ready, weights_read, emit)
+
+            def diffed(scale_changes: int) -> None:
+                args = (step, base_step, step_dir, num_parts, start, time.perf_counter() - start, scale_changes)
+                self._submit(self.writes, self._finish_delta, *args)
+
+            encoder.encode(shards, ready, weights_read, emit, diffed)
         except BaseException as error:
             # Recorded before the trainer wakes up, so it fails at its next wait instead of training on
             self.errors.append(error)
@@ -1191,9 +1196,6 @@ def emit(part: tuple) -> None:
             raise
         # Free the gathered fp8 tensors
         shards.clear()
-        encode_time = time.perf_counter() - start
-        args = (step, base_step, step_dir, num_parts, start, encode_time, scale_changes)
-        self._submit(self.writes, self._finish_delta, *args)
         if anchor is not None:
             # The residual's previous bits hold v{step} until the next encode, which runs on this thread
             source = {key: (self.encoder.shards[key][1], bits) for key, bits in encoder.previous.items()}

From 81064f4da14b808864e0073382124dff96db525b Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sun, 18 Oct 2026 17:26:50 +0000
Subject: [PATCH 61/88] fix(weights): release the diff slot and regions before
 writing the anchor slices, also on errors

---
 src/prime_rl/transports/weights/sparse.py | 22 ++++++++++++++--------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index c0b2a170688..f85b317883b 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -802,14 +802,14 @@ def _launch(self, items: list[tuple], old: torch.Tensor, release: Callable, writ
     @torch.no_grad()
     def _collect(self, launches: list[tuple], keys: list[str], old: torch.Tensor, ready, release, written) -> None:
         """Diff thread: the host parts of ``_launch``'s changes to ``keys`` (writing the changes past a region's
-        capacity in further rounds), and the anchor slices of ``written``; then ``release`` the old values."""
+        capacity in further rounds); then ``release`` the old values and the regions, and write the anchor slices
+        of ``written``."""
         torch.cuda.set_device(self.device)
         ready.synchronize()
         start = time.perf_counter()
         try:
             with torch.cuda.stream(self.diff_stream):
-                for pieces, ends, region in launches:
-                    indices, values, counts = region
+                for pieces, ends, (indices, values, counts) in launches:
                     block_counts = counts[: len(pieces.tables[5])].numpy()
                     total = int(block_counts.sum())
                     changed, changed_values = [], []
@@ -820,7 +820,6 @@ def _collect(self, launches: list[tuple], keys: list[str], old: torch.Tensor, re
                         num = min(len(indices), total - skip)
                         changed.append(indices[:num].numpy().copy())
                         changed_values.append(values[:num].numpy().copy())
-                    self.regions.put(region)
                     if not total:
                         continue
                     piece_counts = np.add.reduceat(block_counts, pieces.first_blocks)
@@ -830,13 +829,20 @@ def _collect(self, launches: list[tuple], keys: list[str], old: torch.Tensor, re
                     previous[(np.cumsum(piece_counts) - piece_counts)[piece_counts > 0]] = -1
                     entries = [(name, int(count)) for name, count in zip(pieces.names, piece_counts) if count]
                     self.host_parts.append((entries, index - previous - 1, np.concatenate(changed_values)))
-                if self.anchor is not None:
+        finally:
+            # Before the anchor writes, which read the live weights: the next updates wait for these
+            for _, _, region in launches:
+                self.regions.put(region)
+            release()
+        try:
+            if self.anchor is not None:
+                with torch.cuda.stream(self.diff_stream):
                     for key, fqn, numel, offset, view in written:
-                        self._write_anchor_slices(key, fqn, numel, offset, view.reshape(-1).to(self.wire_dtypes[key]))
+                        values = view.reshape(-1).to(self.wire_dtypes[key])
+                        self._write_anchor_slices(key, fqn, numel, offset, values)
                     self.diff_stream.synchronize()
-                self.stepped.update(keys)
+            self.stepped.update(keys)
         finally:
-            release()
             self.diff_busy += time.perf_counter() - start
         if sum(len(part[1]) for part in self.host_parts) >= self.HOOK_PART:
             self._emit_host_parts(self.emit)

From 4bc9e2dd25827bdd7c9e79a336e5f0b36c4a1c68 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Mon, 19 Oct 2026 08:51:58 +0000
Subject: [PATCH 62/88] fix(weights): reject Muon and LoRA with the sparse
 weight broadcast in the trainer config

Muon (dion) has no update hook yet, so the sparse broadcast cannot take its deltas from the optimizer step.
---
 docs/scaling.md                                           | 4 ++--
 packages/prime-rl-configs/src/prime_rl/configs/rl.py      | 6 ++++++
 packages/prime-rl-configs/src/prime_rl/configs/shared.py  | 2 +-
 packages/prime-rl-configs/src/prime_rl/configs/trainer.py | 5 +++++
 4 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index 54500d548d0..acfabd0f3c5 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -285,10 +285,10 @@ sparse = true
 anchor_interval = 50  # full checkpoint every N steps
 ```
 
-- RL: deltas come from the optimizer step, with every optimizer and offload mode. Right before the step updates a chunk of parameters, their shards still hold the previous broadcast's weights; the trainer diffs the old and new values of each chunk as the step goes and writes the delta out. No copy of the weights is kept, except of the few tensors that are not optimizer parameters (buffers, frozen parameters, fp8 tensors whose 128-row blocks straddle trainer ranks). The trainer checks that no weight changed outside the optimizer step.
+- RL: deltas come from the optimizer step, with AdamW and SignSGD in every offload mode. Right before the step updates a chunk of parameters, their shards still hold the previous broadcast's weights; the trainer diffs the old and new values of each chunk as the step goes and writes the delta out. No copy of the weights is kept, except of the few tensors that are not optimizer parameters (buffers, frozen parameters, fp8 tensors whose 128-row blocks straddle trainer ranks). The trainer checks that no weight changed outside the optimizer step.
 - A broadcast that does not follow every optimizer step (the startup broadcast of a resumed run, SFT online evals) publishes a full checkpoint.
 - Anchors (full checkpoints) are written every `anchor_interval` steps, from the optimizer step. They slow that step down, so prefer a larger `anchor_interval`. An engine that joins, restarts or fails an update loads the latest complete anchor, reading only the rows it holds, and replays the deltas since.
-- Not supported: LoRA, fp8 tensors whose columns are not a multiple of 128, and engines that transform weights after loading other than online `fp8_per_block` (e.g. FlashInfer TRTLLM/CUTLASS MoE, Marlin). fp8 deltas need the same vLLM and torch versions on the trainer and the engines.
+- Not supported: Muon (yet), LoRA, fp8 tensors whose columns are not a multiple of 128, and engines that transform weights after loading other than online `fp8_per_block` (e.g. FlashInfer TRTLLM/CUTLASS MoE, Marlin). fp8 deltas need the same vLLM and torch versions on the trainer and the engines.
 
 ### Custom Templates
 
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/rl.py b/packages/prime-rl-configs/src/prime_rl/configs/rl.py
index 7ca41f0a959..7d5e18a735f 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py
@@ -440,6 +440,12 @@ def auto_setup_weight_broadcast(self):
                 "PEFT-shaped directory on disk (LoRAModel.from_local_checkpoint) - in-memory transports "
                 "have no disk artifact to load from."
             )
+        if (
+            self.weight_broadcast.type == "filesystem"
+            and self.weight_broadcast.sparse
+            and self.trainer.optim.type == "muon"
+        ):
+            raise ValueError("Sparse weight broadcast does not support Muon yet")
         self.trainer.weight_broadcast = self.weight_broadcast.model_copy()
         self.orchestrator.weight_broadcast = self.weight_broadcast.model_copy()
         if self.inference is not None:
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
index ed5e5f3fc28..e5ee334f763 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
@@ -50,7 +50,7 @@ class FileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig):
 
     sparse: bool = False
     """Between anchors, publish only the weight elements that changed since the previous broadcast.
-    Engines patch their weights in place and verify every update. Not supported with LoRA."""
+    Engines patch their weights in place and verify every update. Not supported with LoRA or Muon."""
 
     anchor_interval: int = Field(50, ge=1)
     """With ``sparse``, publish a full checkpoint every this many steps. Engines that join or fall
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
index ae98bb46f1f..e7d61b93f4e 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
@@ -715,4 +715,9 @@ def validate_lora_broadcast(self):
                 "LoRA requires weight_broadcast.type = 'filesystem': vLLM loads adapters only from a "
                 "PEFT-shaped directory on disk - in-memory transports have no disk artifact to load from."
             )
+        if self.weight_broadcast.type == "filesystem" and self.weight_broadcast.sparse:
+            if self.model.lora is not None:
+                raise ValueError("Sparse weight broadcast does not support LoRA")
+            if self.optim.type == "muon":
+                raise ValueError("Sparse weight broadcast does not support Muon yet")
         return self

From 36e0de765364d3b45b418962dc70f83099fd7ae4 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Mon, 19 Oct 2026 11:32:04 +0000
Subject: [PATCH 63/88] fix(weights): fail engines reading a sparse delta whose
 trainer encode failed

A trainer rank whose encode or delta writes fail leaves a failure marker in the delta dir; read_delta raises on it instead of polling for parts that never come, and the prefetch request returns the error.
---
 src/prime_rl/inference/vllm/worker/filesystem.py |  8 +++++++-
 src/prime_rl/transports/weights/sparse.py        |  7 +++++++
 src/prime_rl/utils/sparse_weights.py             | 12 ++++++++++++
 tests/unit/utils/test_sparse_weights.py          | 10 ++++++++++
 4 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 5a29adf0cb1..8486ccfd799 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -147,7 +147,13 @@ def _read_delta(
         return pinned, checksums
 
     def prefetch_weights_done(self) -> bool:
-        return self.prefetched is None or self.prefetched[1].done()
+        if self.prefetched is None:
+            return True
+        if not self.prefetched[1].done():
+            return False
+        # Raises if the read failed, e.g. on a trainer rank's failure marker
+        self.prefetched[1].result()
+        return True
 
     def update_weights_from_path(self, weight_path: str) -> None:
         """Move to the broadcast step in ``weight_path`` (a full HF checkpoint and/or a sparse delta)."""
diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index f85b317883b..80ac7a97a92 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -34,6 +34,7 @@
     compile_fp8_quantizer,
     encode_gaps,
     quantize_fp8_blocks,
+    rank_failed_marker,
     rank_part,
     rank_parts_marker,
     read_engine_layout,
@@ -1025,6 +1026,8 @@ def __init__(self, output_dir: Path, config: FileSystemWeightBroadcastConfig):
         self.armed_step: int | None = None
         self.streamed_dir: Path | None = None
         self.num_streamed = 0
+        # The delta dir of the step being encoded, where a failed encode or write leaves this rank's failure marker
+        self.delta_dir: Path | None = None
 
     def diff_in_optimizer(self, optimizer: "OptimizerLike") -> None:
         """Take the deltas of the optimizer's parameters from its steps, which must each be followed by a broadcast."""
@@ -1103,6 +1106,9 @@ def run() -> None:
             except BaseException as error:
                 self.logger.exception(f"Sparse weight broadcast failed in {fn.__name__}")
                 self.errors.append(error)
+                if pool is not self.anchors and self.delta_dir is not None:
+                    # Engines reading the delta fail instead of waiting for parts that never come
+                    _write_json(self.delta_dir / rank_failed_marker(self.world.rank), {"error": repr(error)})
                 raise
 
         return pool.submit(run)
@@ -1176,6 +1182,7 @@ def _encode(
         encoder = self.encoder
         start = time.perf_counter()
         num_parts = 0
+        self.delta_dir = step_dir / DELTA_DIR
 
         def emit(part: tuple) -> None:
             nonlocal num_parts
diff --git a/src/prime_rl/utils/sparse_weights.py b/src/prime_rl/utils/sparse_weights.py
index 66462523a55..e808c5290cf 100644
--- a/src/prime_rl/utils/sparse_weights.py
+++ b/src/prime_rl/utils/sparse_weights.py
@@ -441,6 +441,11 @@ def rank_parts_marker(rank: int) -> str:
     return f"rank_{rank}.parts"
 
 
+def rank_failed_marker(rank: int) -> str:
+    """Written by a trainer rank whose encode or writes of the delta failed."""
+    return f"rank_{rank}.failed"
+
+
 def read_delta(
     step_dir: Path, pool: ThreadPoolExecutor, names: set[str] | None = None
 ) -> tuple[list[tuple[list[tuple[str, int]], np.ndarray, np.ndarray]], dict[str, int]]:
@@ -475,6 +480,13 @@ def read_delta(
                     next_part[rank] += 1
                     found = True
             if not found:
+                for rank in range(len(num_parts)):
+                    failed = delta_dir / rank_failed_marker(rank)
+                    if failed.exists():
+                        raise RuntimeError(
+                            f"Trainer rank {rank} failed to write the sparse weight delta in {step_dir}: "
+                            f"{json.loads(failed.read_text())['error']}"
+                        )
                 time.sleep(0.02)
         frames = [frame for frame in (job.result() for job in jobs) if frame[0]]
         for file in files:
diff --git a/tests/unit/utils/test_sparse_weights.py b/tests/unit/utils/test_sparse_weights.py
index f6e7da3e95b..067746da67b 100644
--- a/tests/unit/utils/test_sparse_weights.py
+++ b/tests/unit/utils/test_sparse_weights.py
@@ -13,6 +13,7 @@
     bits_checksum,
     decode_gaps,
     encode_gaps,
+    rank_failed_marker,
     rank_part,
     rank_parts_marker,
     read_anchor_checksums,
@@ -99,3 +100,12 @@ def test_anchor_checksums_and_unwritten_slices(tmp_path):
             with pytest.raises(RuntimeError, match="a@6"):
                 writer.close(1)
             assert not anchor_complete(step_dir)
+
+
+def test_read_delta_fails_on_a_failed_trainer_rank(tmp_path):
+    delta_dir = tmp_path / DELTA_DIR
+    delta_dir.mkdir()
+    (delta_dir / MANIFEST).write_text(json.dumps({"ranks": 2, "tensors": {}}))
+    (delta_dir / rank_failed_marker(1)).write_text(json.dumps({"error": "RuntimeError('boom')"}))
+    with ThreadPoolExecutor(1) as pool, pytest.raises(RuntimeError, match="rank 1 failed.*boom"):
+        read_delta(tmp_path, pool)

From 2e99ce3c8bd856290a0ee7e6caecc9cffb05bbf1 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Tue, 27 Oct 2026 16:10:24 +0000
Subject: [PATCH 64/88] fix(weights): engine layouts from every world rank, so
 pipeline stages under EP do not overwrite each other

---
 src/prime_rl/inference/vllm/worker/filesystem.py | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 8486ccfd799..76496b81dff 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -9,7 +9,7 @@
 import numpy as np
 import torch
 from torch.nn import Module
-from vllm.distributed.parallel_state import get_ep_group, get_tp_group, get_world_group
+from vllm.distributed.parallel_state import get_tp_group, get_world_group
 from vllm.logger import init_logger
 from vllm.model_executor.model_loader import DefaultModelLoader, get_model_loader
 
@@ -229,11 +229,8 @@ def read_tensor(name: str) -> torch.Tensor:
             read_checkpoint_range(boot[name], 0, prod(shape), data)
             return torch.from_numpy(data).view(dtype).view(shape)
 
-        # The trainer waits for every worker's layout
-        if self.vllm_config.parallel_config.enable_expert_parallel:
-            group = get_ep_group()
-        else:
-            group = get_world_group()
+        # The trainer waits for every worker's layout. The world group spans every DP, PP and TP rank
+        group = get_world_group()
         layout = {"world_size": group.world_size, "fp8": index_map.fp8_layout(model, read_tensor)}
         self.index_map = index_map
         path = step_dir.parent / ENGINE_LAYOUT / f"rank_{group.rank_in_group}.json"

From 2299e35116b4e603446e4cd70c82752dce32168d Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Tue, 27 Oct 2026 16:10:24 +0000
Subject: [PATCH 65/88] fix(weights): engines report their layout at every full
 broadcast, and the trainer clears the last run's at its first

---
 .../inference/vllm/worker/filesystem.py       | 21 ++++++++++++-------
 src/prime_rl/transports/weights/sparse.py     |  5 +++++
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 76496b81dff..66c7bef615d 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -160,14 +160,17 @@ def update_weights_from_path(self, weight_path: str) -> None:
         step_dir = Path(weight_path)
         step = int(step_dir.name.removeprefix("step_"))
         has_delta = (step_dir / DELTA_DIR / MANIFEST).exists()
+        base_step = read_manifest(step_dir)["base_step"] if has_delta else None
         if has_delta or anchor_complete(step_dir):
             self._build_index_map(step_dir)
+            if base_step is None:
+                # A full broadcast: the trainer takes the layout at its first one, also after a restart
+                self._write_layout(step_dir.parent)
         # Unknown until this update succeeds: after a failure, the next update reloads and replays
         held_step, self.weight_step = self.weight_step, -1
         if not has_delta:
             self._load_checkpoint(step_dir)
         else:
-            base_step = read_manifest(step_dir)["base_step"]
             if base_step is None:
                 # The run's first broadcast: the weights the engine booted with
                 if held_step is not None:
@@ -208,8 +211,7 @@ def _apply_chunk(
 
     def _build_index_map(self, step_dir: Path) -> None:
         """At a sparse run's first update, while the engine still holds its boot weights: map checkpoint tensors to
-        the params, and report which ones this worker holds in fp8 (and their scale format), so the trainer sends
-        their deltas as fp8."""
+        the params, and find which ones this worker holds in fp8 (and their scale format)."""
         if self.index_map is not None:
             return
         if (step_dir / DELTA_DIR / MANIFEST).exists():
@@ -229,11 +231,16 @@ def read_tensor(name: str) -> torch.Tensor:
             read_checkpoint_range(boot[name], 0, prod(shape), data)
             return torch.from_numpy(data).view(dtype).view(shape)
 
-        # The trainer waits for every worker's layout. The world group spans every DP, PP and TP rank
-        group = get_world_group()
-        layout = {"world_size": group.world_size, "fp8": index_map.fp8_layout(model, read_tensor)}
+        index_map.fp8_layout(model, read_tensor)
         self.index_map = index_map
-        path = step_dir.parent / ENGINE_LAYOUT / f"rank_{group.rank_in_group}.json"
+
+    def _write_layout(self, broadcast_dir: Path) -> None:
+        """Report which checkpoint tensors this worker holds in fp8, so the trainer sends their deltas as fp8. The
+        trainer waits for every worker of the world group (every DP, PP and TP rank)."""
+        assert self.index_map is not None
+        group = get_world_group()
+        layout = {"world_size": group.world_size, "fp8": self.index_map.fp8}
+        path = broadcast_dir / ENGINE_LAYOUT / f"rank_{group.rank_in_group}.json"
         path.parent.mkdir(exist_ok=True)
         write_file(path, json.dumps(layout).encode())
         logger.info(f"Wrote the engine layout: {len(layout['fp8'])} checkpoint tensors held in fp8")
diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 80ac7a97a92..e1bd35e2446 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -1,5 +1,6 @@
 import json
 import queue
+import shutil
 import threading
 import time
 from bisect import bisect_left, bisect_right
@@ -26,6 +27,7 @@
 from prime_rl.utils.logger import get_logger
 from prime_rl.utils.sparse_weights import (
     DELTA_DIR,
+    ENGINE_LAYOUT,
     FP8_BLOCK,
     MANIFEST,
     AnchorWriter,
@@ -1139,6 +1141,9 @@ def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
     def _broadcast_full(self, model: nn.Module, step: int, step_dir: Path) -> None:
         """Publish v{step} as a full checkpoint of the live weights; the run's first broadcast at step 0 is the
         checkpoint the engines booted from. The next delta is against it."""
+        if self.last_step is None and self.world.is_master:
+            # Engines report their layout again as they load it; the last run's may differ
+            shutil.rmtree(step_dir.parent / ENGINE_LAYOUT, ignore_errors=True)
         if self.last_step is None and step == 0:
             self.anchor_step = step
             if self.world.is_master:

From b5a4f3cc321bd7a0c2abb9333d6f9b006ee544bb Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 28 Oct 2026 14:16:13 +0000
Subject: [PATCH 66/88] feat(weights): sparse weight broadcast with Muon

Muon steps its batches of matrices as update chunks (the same on every rank, as they run collectives), so the
sparse broadcast takes Muon's deltas from its step like AdamW's. Every chunk steps a copy of its param group, so
per-group state such as Muon's step count advances once per step.
---
 src/prime_rl/trainer/optim/__init__.py | 15 +++++++++-
 src/prime_rl/trainer/optim/base.py     | 41 +++++++++++++++-----------
 2 files changed, 37 insertions(+), 19 deletions(-)

diff --git a/src/prime_rl/trainer/optim/__init__.py b/src/prime_rl/trainer/optim/__init__.py
index 030d2d0c956..40cfaebe94f 100644
--- a/src/prime_rl/trainer/optim/__init__.py
+++ b/src/prime_rl/trainer/optim/__init__.py
@@ -1,6 +1,7 @@
+import dion
 import torch
 import torch.distributed as dist
-from dion import Muon
+from dion.opt_utils import create_param_batches
 from torch import nn
 from torch.distributed.device_mesh import DeviceMesh
 from torch.optim import Optimizer
@@ -29,6 +30,18 @@ class SignSGD(ChunkedUpdates, _SignSGD):
     pass
 
 
+class Muon(ChunkedUpdates, dion.Muon):
+    def update_units(self, group: dict):
+        if group["algorithm"] != "muon":
+            yield from super().update_units(group)
+            return
+        # Muon's batches of matrices it orthogonalizes together, the same on every rank: they run collectives
+        world_size = self._get_group_mesh_info(group)[1]
+        params = [param for param in group["params"] if param.grad is not None]
+        for batch in create_param_batches(params, world_size, self._matrix_partitions):
+            yield batch, sum(param.numel() for param in batch) // world_size
+
+
 def _warmup_muon_mesh(mesh: DeviceMesh) -> None:
     """Establish NCCL peer connections before Muon's first optimizer step.
 
diff --git a/src/prime_rl/trainer/optim/base.py b/src/prime_rl/trainer/optim/base.py
index b0f002c87a8..e24a760c6c5 100644
--- a/src/prime_rl/trainer/optim/base.py
+++ b/src/prime_rl/trainer/optim/base.py
@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
 from typing import TypeAlias
 
 import torch
@@ -12,12 +12,18 @@
 
 
 class ChunkedUpdates:
-    """Optimizer mixin: with ``update_hook`` set, ``step`` updates the parameters a chunk of up to ``CHUNK_NUMEL``
+    """Optimizer mixin: with ``update_hook`` set, ``step`` updates the parameters a chunk of about ``CHUNK_NUMEL``
     local elements at a time and calls the hook around each chunk's update."""
 
     CHUNK_NUMEL = 1 << 26
     update_hook: UpdateHook | None = None
 
+    def update_units(self, group: dict) -> Iterator[tuple[list[torch.Tensor], int]]:
+        """The parameters of a group to update, in units that a chunk does not split, with their local numel."""
+        for param in group["params"]:
+            if param.grad is not None:
+                yield [param], (param.to_local() if isinstance(param, DTensor) else param).numel()
+
     def step(self, closure=None):
         if self.update_hook is None:
             return super().step(closure)
@@ -25,26 +31,25 @@ def step(self, closure=None):
         groups = self.param_groups
         try:
             for group in groups:
-                chunk, numel = [], 0
-                for param in group["params"]:
-                    if param.grad is None:
-                        continue
-                    chunk.append(param)
-                    numel += (param.to_local() if isinstance(param, DTensor) else param).numel()
+                chunks, numel = [[]], 0
+                for params, unit_numel in self.update_units(group):
                     if numel >= self.CHUNK_NUMEL:
-                        self._step_params(group, chunk)
-                        chunk, numel = [], 0
-                if chunk:
-                    self._step_params(group, chunk)
+                        chunks.append([])
+                        numel = 0
+                    chunks[-1] += params
+                    numel += unit_numel
+                stepped = group
+                for chunk in filter(None, chunks):
+                    # Every chunk steps a copy of the group, so per-group state (e.g. Muon's step count) advances once
+                    stepped = {**group, "params": chunk}
+                    self.param_groups = [stepped]
+                    finish = self.update_hook(chunk)
+                    super().step()
+                    finish()
+                group.update({key: value for key, value in stepped.items() if key != "params"})
         finally:
             self.param_groups = groups
 
-    def _step_params(self, group: dict, params: list[torch.Tensor]) -> None:
-        self.param_groups = [{**group, "params": params}]
-        finish = self.update_hook(params)
-        super().step()
-        finish()
-
 
 class OffloadOptimizer(ABC):
     @property

From b8ecec127dccad1875f0e4578d9e687337235d65 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 28 Oct 2026 14:16:13 +0000
Subject: [PATCH 67/88] refactor(weights):
 SparseFileSystemWeightBroadcastConfig, 1:1 with SparseFileSystemWeightSender

`[weight_broadcast] type = "sparse_filesystem"` replaces `type = "filesystem"` + `sparse = true`. LoRA is rejected
by the existing non-filesystem check; Muon is supported. Engines use the filesystem worker for both types. The
worker's delta chunks get a named type.
---
 docs/scaling.md                                |  9 ++++-----
 .../src/prime_rl/configs/inference.py          |  2 +-
 .../src/prime_rl/configs/rl.py                 |  8 +-------
 .../src/prime_rl/configs/shared.py             | 18 ++++++++++++------
 .../src/prime_rl/configs/trainer.py            |  7 +------
 src/prime_rl/inference/vllm/server.py          |  1 +
 .../inference/vllm/worker/filesystem.py        | 18 +++++++++++++-----
 src/prime_rl/transports/weights/__init__.py    |  6 +++---
 src/prime_rl/transports/weights/sparse.py      |  9 ++++-----
 9 files changed, 40 insertions(+), 38 deletions(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index 0c6681e04d0..fdae880f1fd 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -290,19 +290,18 @@ By default, the trainer and inference worker each allocate one transfer arena. S
 
 ### Sparse filesystem weight broadcast
 
-With `[weight_broadcast] type = "filesystem"` and `sparse = true`, the trainer publishes only the weight elements that changed since the previous broadcast, in the engine's weight format (bf16, or fp8 for engines that quantize online to `fp8_per_block`). Engines read and decode each delta while they keep serving, patch their weights in place during the pause, and verify every update against trainer checksums.
+With `[weight_broadcast] type = "sparse_filesystem"`, the trainer publishes only the weight elements that changed since the previous broadcast, in the engine's weight format (bf16, or fp8 for engines that quantize online to `fp8_per_block`). Engines read and decode each delta while they keep serving, patch their weights in place during the pause, and verify every update against trainer checksums.
 
 ```toml
 [weight_broadcast]
-type = "filesystem"
-sparse = true
+type = "sparse_filesystem"
 anchor_interval = 50  # full checkpoint every N steps
 ```
 
-- RL: deltas come from the optimizer step, with AdamW and SignSGD in every offload mode. Right before the step updates a chunk of parameters, their shards still hold the previous broadcast's weights; the trainer diffs the old and new values of each chunk as the step goes and writes the delta out. No copy of the weights is kept, except of the few tensors that are not optimizer parameters (buffers, frozen parameters, fp8 tensors whose 128-row blocks straddle trainer ranks). The trainer checks that no weight changed outside the optimizer step.
+- RL: deltas come from the optimizer step, with every optimizer and offload mode. Right before the step updates a chunk of parameters, their shards still hold the previous broadcast's weights; the trainer diffs the old and new values of each chunk as the step goes and writes the delta out. No copy of the weights is kept, except of the few tensors that are not optimizer parameters (buffers, frozen parameters, fp8 tensors whose 128-row blocks straddle trainer ranks). The trainer checks that no weight changed outside the optimizer step.
 - A broadcast that does not follow every optimizer step (the startup broadcast of a resumed run, SFT online evals) publishes a full checkpoint.
 - Anchors (full checkpoints) are written every `anchor_interval` steps, from the optimizer step. They slow that step down, so prefer a larger `anchor_interval`. An engine that joins, restarts or fails an update loads the latest complete anchor, reading only the rows it holds, and replays the deltas since.
-- Not supported: Muon (yet), LoRA, fp8 tensors whose columns are not a multiple of 128, and engines that transform weights after loading other than online `fp8_per_block` (e.g. FlashInfer TRTLLM/CUTLASS MoE, Marlin). fp8 deltas need the same vLLM and torch versions on the trainer and the engines.
+- Not supported: LoRA, fp8 tensors whose columns are not a multiple of 128, and engines that transform weights after loading other than online `fp8_per_block` (e.g. FlashInfer TRTLLM/CUTLASS MoE, Marlin). fp8 deltas need the same vLLM and torch versions on the trainer and the engines.
 
 ### Custom Templates
 
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/inference.py b/packages/prime-rl-configs/src/prime_rl/configs/inference.py
index 3b1bc706b07..2a0bf1993a9 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/inference.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/inference.py
@@ -186,7 +186,7 @@ def auto_setup_api_server_count(self):
 
 
 class WeightBroadcastConfig(BaseConfig):
-    type: Literal["nccl", "filesystem", "nixl"] = "filesystem"
+    type: Literal["nccl", "filesystem", "sparse_filesystem", "nixl"] = "filesystem"
     """Weight broadcast transport."""
 
 
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/rl.py b/packages/prime-rl-configs/src/prime_rl/configs/rl.py
index 328cc601f6a..888ea3f8977 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py
@@ -442,12 +442,6 @@ def auto_setup_weight_broadcast(self):
                 "PEFT-shaped directory on disk (LoRAModel.from_local_checkpoint) - in-memory transports "
                 "have no disk artifact to load from."
             )
-        if (
-            self.weight_broadcast.type == "filesystem"
-            and self.weight_broadcast.sparse
-            and self.trainer.optim.type == "muon"
-        ):
-            raise ValueError("Sparse weight broadcast does not support Muon yet")
         self.trainer.weight_broadcast = self.weight_broadcast.model_copy()
         self.orchestrator.weight_broadcast = self.weight_broadcast.model_copy()
         if self.inference is not None:
@@ -638,7 +632,7 @@ def auto_setup_deployment(self):
                 dp = self.inference.vllm.data_parallel_size
                 if self.inference.vllm.api_server_count < dp and not self.inference.vllm.enable_lora:
                     self.inference.vllm.api_server_count = dp
-                if self.weight_broadcast.type != "filesystem":
+                if self.weight_broadcast.type in ("nccl", "nixl"):
                     # Every inference GPU joins the in-memory transfer group, so size it after the DP fill.
                     world_size = dp * self.inference.vllm.tensor_parallel_size
                     self.trainer.weight_broadcast.inference_world_size = world_size
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
index dce6c05dac0..f21be1cb677 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
@@ -48,13 +48,16 @@ class BaseWeightBroadcastConfig(BaseConfig):
 class FileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig):
     type: Literal["filesystem"] = "filesystem"
 
-    sparse: bool = False
-    """Between anchors, publish only the weight elements that changed since the previous broadcast.
-    Engines patch their weights in place and verify every update. Not supported with LoRA or Muon."""
+
+class SparseFileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig):
+    """Filesystem broadcast that, between anchors, publishes only the weight elements that changed since the
+    previous broadcast. Engines patch their weights in place and verify every update. Not supported with LoRA."""
+
+    type: Literal["sparse_filesystem"] = "sparse_filesystem"
 
     anchor_interval: int = Field(50, ge=1)
-    """With ``sparse``, publish a full checkpoint every this many steps. Engines that join or fall
-    off the delta chain load the last one and replay the deltas since."""
+    """Publish a full checkpoint every this many steps. Engines that join or fall off the delta chain load
+    the last one and replay the deltas since."""
 
 
 class InMemoryWeightBroadcastConfig(BaseWeightBroadcastConfig):
@@ -89,7 +92,10 @@ class NIXLWeightBroadcastConfig(InMemoryWeightBroadcastConfig):
 
 
 WeightBroadcastConfig: TypeAlias = Annotated[
-    FileSystemWeightBroadcastConfig | NCCLWeightBroadcastConfig | NIXLWeightBroadcastConfig,
+    FileSystemWeightBroadcastConfig
+    | SparseFileSystemWeightBroadcastConfig
+    | NCCLWeightBroadcastConfig
+    | NIXLWeightBroadcastConfig,
     Field(discriminator="type"),
 ]
 """Weight transfer from trainer to inference. ``rl`` copies its ``[weight_broadcast]`` to the trainer and orchestrator."""
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
index ab847f806a9..7c8732abe44 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
@@ -716,16 +716,11 @@ class TrainerConfig(BaseTrainerConfig):
 
     @model_validator(mode="after")
     def validate_lora_broadcast(self):
-        if self.model.lora is not None and self.weight_broadcast.type in ("nccl", "nixl"):
+        if self.model.lora is not None and self.weight_broadcast.type != "filesystem":
             raise ValueError(
                 "LoRA requires weight_broadcast.type = 'filesystem': vLLM loads adapters only from a "
                 "PEFT-shaped directory on disk - in-memory transports have no disk artifact to load from."
             )
-        if self.weight_broadcast.type == "filesystem" and self.weight_broadcast.sparse:
-            if self.model.lora is not None:
-                raise ValueError("Sparse weight broadcast does not support LoRA")
-            if self.optim.type == "muon":
-                raise ValueError("Sparse weight broadcast does not support Muon yet")
         return self
 
     @model_validator(mode="after")
diff --git a/src/prime_rl/inference/vllm/server.py b/src/prime_rl/inference/vllm/server.py
index a06d1ff3bf9..b74c32fd44e 100644
--- a/src/prime_rl/inference/vllm/server.py
+++ b/src/prime_rl/inference/vllm/server.py
@@ -36,6 +36,7 @@ def models(request: Request) -> OpenAIServingModels:
 WORKER_EXTENSION_CLS = {
     "nccl": "prime_rl.inference.vllm.worker.nccl.NCCLWeightUpdateWorker",
     "filesystem": "prime_rl.inference.vllm.worker.filesystem.FileSystemWeightUpdateWorker",
+    "sparse_filesystem": "prime_rl.inference.vllm.worker.filesystem.FileSystemWeightUpdateWorker",
     "nixl": "prime_rl.inference.vllm.worker.nixl.NIXLWeightUpdateWorker",
 }
 
diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 66c7bef615d..a272988ee63 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -4,7 +4,7 @@
 from concurrent.futures import Future, ThreadPoolExecutor
 from math import prod
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, NamedTuple
 
 import numpy as np
 import torch
@@ -47,6 +47,16 @@
 
 logger = init_logger("vllm.inference.vllm.worker_filesystem")
 
+
+class DeltaChunk(NamedTuple):
+    """Changes to apply at once, in pinned host memory: ``(checkpoint tensor, count)`` entries, their gaps
+    (``decode_gaps``) and value bytes."""
+
+    entries: list[tuple[str, int]]
+    gaps: torch.Tensor
+    values: torch.Tensor
+
+
 _prefetch_pool = ThreadPoolExecutor(1)
 _reader_pool: ThreadPoolExecutor | None = None
 
@@ -80,9 +90,7 @@ def prefetch_weights_from_path(self, weight_path: str) -> None:
         """Start decoding the sparse delta in ``weight_path`` into pinned host memory."""
         self.prefetched = (weight_path, _prefetch_pool.submit(self._read_delta, Path(weight_path)))
 
-    def _read_delta(
-        self, step_dir: Path
-    ) -> tuple[list[tuple[list[tuple[str, int]], torch.Tensor, torch.Tensor]], dict[str, int]]:
+    def _read_delta(self, step_dir: Path) -> tuple[list[DeltaChunk], dict[str, int]]:
         """Read this worker's part of a delta into a pinned host buffer reused across steps, in chunks of up to 2^26
         changes of consecutive frames."""
         global _reader_pool
@@ -132,7 +140,7 @@ def _read_delta(
                 values_end += values.nbytes
             if chunk:
                 pinned.append(
-                    (
+                    DeltaChunk(
                         [entry for entries, _, _ in chunk for entry in entries],
                         self.host_buffer[gaps_at:gaps_end].view(torch.int32),
                         self.host_buffer[values_at:values_end],
diff --git a/src/prime_rl/transports/weights/__init__.py b/src/prime_rl/transports/weights/__init__.py
index c5150bf4682..d967cd89a42 100644
--- a/src/prime_rl/transports/weights/__init__.py
+++ b/src/prime_rl/transports/weights/__init__.py
@@ -29,10 +29,10 @@ def setup_weight_sender(
 ) -> WeightSender:
     if config.type == "nccl":
         return NCCLWeightSender(output_dir, config, torch.cuda.current_device())
-    elif config.type == "filesystem" and config.sparse:
-        return SparseFileSystemWeightSender(output_dir, config)
     elif config.type == "filesystem":
         return FileSystemWeightSender(output_dir, config, lora_config)
+    elif config.type == "sparse_filesystem":
+        return SparseFileSystemWeightSender(output_dir, config)
     elif config.type == "nixl":
         return NIXLWeightSender(output_dir, config, parallel_dims)
     else:
@@ -47,7 +47,7 @@ def setup_weight_receiver(
 ) -> WeightReceiver:
     if config.type == "nccl":
         return NCCLWeightReceiver(broadcast_dir, config, admin_plane, model_name)
-    elif config.type == "filesystem":
+    elif config.type in ("filesystem", "sparse_filesystem"):
         return FileSystemWeightReceiver(broadcast_dir, config, admin_plane, model_name)
     elif config.type == "nixl":
         return NIXLWeightReceiver(broadcast_dir, config, admin_plane, model_name)
diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index e1bd35e2446..2b8b5e08a7d 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -19,10 +19,9 @@
 from torch.distributed.tensor import DTensor
 from torch.distributed.tensor._utils import compute_local_shape_and_global_offset
 
-from prime_rl.configs.shared import FileSystemWeightBroadcastConfig
+from prime_rl.configs.shared import SparseFileSystemWeightBroadcastConfig
 from prime_rl.trainer.models.fusions import get_model_packed_parameters, join_module_path
 from prime_rl.transports.weights.base import WeightSender
-from prime_rl.transports.weights.filesystem import FileSystemWeightSender
 from prime_rl.transports.weights.sparse_kernels import Pieces, build_pieces, count_changes, write_changes
 from prime_rl.utils.logger import get_logger
 from prime_rl.utils.sparse_weights import (
@@ -998,7 +997,7 @@ def _write_json(path: Path, value) -> None:
     write_file(path, json.dumps(value).encode())
 
 
-class SparseFileSystemWeightSender(FileSystemWeightSender):
+class SparseFileSystemWeightSender(WeightSender):
     """Filesystem broadcast where steps carry only the delta against the previous broadcast (format in
     ``prime_rl.utils.sparse_weights``). The delta comes from the optimizer step (``diff_in_optimizer``) and is
     written as the step goes; the rest is finished in the background while the next step trains, and the trainer
@@ -1009,8 +1008,8 @@ class SparseFileSystemWeightSender(FileSystemWeightSender):
 
     gathers_full_weights = False
 
-    def __init__(self, output_dir: Path, config: FileSystemWeightBroadcastConfig):
-        super().__init__(output_dir, config)
+    def __init__(self, output_dir: Path, config: SparseFileSystemWeightBroadcastConfig):
+        super().__init__(output_dir, config.timeout)
         self.config = config
         self.encoder: SparseDeltaEncoder | None = None
         self.last_step: int | None = None

From 0cd553e8eb2da10a52541bbfe5ff7e2b07243ced Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 28 Oct 2026 14:16:15 +0000
Subject: [PATCH 68/88] fix(weights): fail sparse commits on failed ranks, fail
 every rank on a layout timeout, keep chains engines may replay

- The master's commit fails on any rank's failure marker instead of waiting for its parts marker.
- The master broadcasts an engine-layout timeout to every rank, so none hangs in the collective.
- Old delta chains are removed only for an anchor that was complete one broadcast earlier: by then every engine
  update that may have started replaying an older chain is done.
---
 src/prime_rl/transports/weights/sparse.py | 38 ++++++++++++++++++-----
 1 file changed, 31 insertions(+), 7 deletions(-)

diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 2b8b5e08a7d..520554cb231 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -1013,7 +1013,11 @@ def __init__(self, output_dir: Path, config: SparseFileSystemWeightBroadcastConf
         self.config = config
         self.encoder: SparseDeltaEncoder | None = None
         self.last_step: int | None = None
+        # The anchor the master keeps every broadcast since (``_clean``), and the latest complete anchor as of the
+        # current broadcast and as of now
         self.anchor_step: int | None = None
+        self.seen_anchor: int | None = None
+        self.complete_anchor: int | None = None
         # Each in order: encodes, delta writes and commits, anchor writes
         self.encodes = ThreadPoolExecutor(1)
         self.writes = ThreadPoolExecutor(1)
@@ -1086,9 +1090,16 @@ def _stream_part(self, part: tuple) -> None:
     def _set_fp8_layout(self) -> None:
         if self.encoder.fp8_layout is None:
             # Every engine worker reports which tensors it holds in fp8 at its first update
-            broadcast_dir = self.step_dir(0).parent
-            layout = [self._wait_for_engine_layout(broadcast_dir) if self.world.is_master else None]
+            layout = [None]
+            if self.world.is_master:
+                try:
+                    layout[0] = self._wait_for_engine_layout(self.step_dir(0).parent)
+                except Exception as error:
+                    # Every rank joins the broadcast and fails with the master
+                    layout[0] = error
             dist.broadcast_object_list(layout, src=0)
+            if isinstance(layout[0], Exception):
+                raise layout[0]
             self.encoder.set_fp8_layout(self.model, layout[0], self.world.is_master)
 
     def close(self) -> None:
@@ -1116,6 +1127,11 @@ def run() -> None:
 
     def _broadcast(self, model: nn.Module, step: int, step_dir: Path) -> None:
         self._wait_for_encode()
+        # Old chains are dropped only for an anchor that was complete one broadcast ago: by this broadcast's
+        # handshake, every engine update that may have started replaying an older chain before it was complete is done
+        if self.seen_anchor is not None:
+            self.anchor_step = max(self.anchor_step, self.seen_anchor)
+        self.seen_anchor = self.complete_anchor
         if self.world.is_master:
             (step_dir / DELTA_DIR).mkdir(exist_ok=True)
         if self.encoder is None:
@@ -1152,7 +1168,8 @@ def _broadcast_full(self, model: nn.Module, step: int, step_dir: Path) -> None:
             self.encoder.write_full_checkpoint(
                 model, self.world.is_master, step_dir, self.world.rank, self.world.world_size
             )
-            self._wait_for_anchor(step, step_dir)
+            self._wait_for_anchor(step_dir)
+            self.anchor_step = step
             if self.world.is_master:
                 WeightSender._finish(self, step, step_dir)
         self.last_step = step
@@ -1260,7 +1277,14 @@ def _finish_delta(
         if not self.world.is_master:
             return
         markers = [delta_dir / rank_parts_marker(rank) for rank in range(self.world.world_size)]
-        while not all(marker.exists() for marker in markers):
+        # A rank whose encode or writes failed leaves its failure marker (before its parts marker, if any)
+        failures = [delta_dir / rank_failed_marker(rank) for rank in range(self.world.world_size)]
+        while True:
+            done = all(marker.exists() for marker in markers)
+            if any(failure.exists() for failure in failures):
+                raise RuntimeError(f"A trainer rank failed to write sparse weight delta v{step}")
+            if done:
+                break
             time.sleep(0.05)
             revalidate(delta_dir)
         WeightSender._finish(self, step, step_dir)
@@ -1278,13 +1302,13 @@ def _finish_delta(
 
     def _close_anchor(self, step: int, step_dir: Path, anchor: AnchorWriter) -> None:
         anchor.close(self.world.world_size)
-        self._wait_for_anchor(step, step_dir)
+        self._wait_for_anchor(step_dir)
+        self.complete_anchor = step
 
-    def _wait_for_anchor(self, step: int, step_dir: Path) -> None:
+    def _wait_for_anchor(self, step_dir: Path) -> None:
         if self.world.is_master:
             while not anchor_complete(step_dir):
                 time.sleep(0.1)
-            self.anchor_step = max(self.anchor_step or 0, step)
 
     def _finish(self, step: int, step_dir: Path) -> None:
         """Committed in ``_broadcast_full`` or in the background (``_finish_delta``)."""

From 14bd1bd4481a3038ab9c3b620749440a6e71fab8 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 28 Oct 2026 14:29:22 +0000
Subject: [PATCH 69/88] test(optim): chunked optimizer steps match a plain step
 and advance per-group step counts once

---
 tests/unit/train/test_chunked_updates.py | 40 ++++++++++++++++++++++++
 1 file changed, 40 insertions(+)
 create mode 100644 tests/unit/train/test_chunked_updates.py

diff --git a/tests/unit/train/test_chunked_updates.py b/tests/unit/train/test_chunked_updates.py
new file mode 100644
index 00000000000..b3712c9375e
--- /dev/null
+++ b/tests/unit/train/test_chunked_updates.py
@@ -0,0 +1,40 @@
+import torch
+
+from prime_rl.trainer.optim.base import ChunkedUpdates
+
+
+class CountingSGD(torch.optim.SGD):
+    """Counts its steps per param group, like Muon."""
+
+    def step(self, closure=None):
+        for group in self.param_groups:
+            group["step"] = group.get("step", 0) + 1
+        return super().step(closure)
+
+
+class ChunkedSGD(ChunkedUpdates, CountingSGD):
+    CHUNK_NUMEL = 10
+
+
+def test_chunked_step_matches_plain_step():
+    params = [torch.nn.Parameter(torch.randn(numel)) for numel in (4, 8, 16, 2)]
+    reference = [torch.nn.Parameter(param.detach().clone()) for param in params]
+    for param in params + reference:
+        param.grad = torch.ones_like(param)
+    params[3].grad = reference[3].grad = None
+    optimizer, plain = ChunkedSGD(params, lr=0.1), CountingSGD(reference, lr=0.1)
+    chunks = []
+
+    def hook(chunk):
+        chunks.append([param.numel() for param in chunk])
+        return lambda: None
+
+    optimizer.update_hook = hook
+    for _ in range(2):
+        optimizer.step()
+        plain.step()
+
+    assert chunks == [[4, 8], [16]] * 2
+    assert all(torch.equal(param, ref) for param, ref in zip(params, reference))
+    assert optimizer.param_groups[0]["step"] == plain.param_groups[0]["step"] == 2
+    assert optimizer.param_groups[0]["params"] == params

From 2ad7455113dffbd714275386a09cec5e98401a27 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 28 Oct 2026 14:45:25 +0000
Subject: [PATCH 70/88] refactor(weights): sparse deltas as
 `[weight_broadcast.delta]` on the filesystem broadcast config

`FileSystemWeightBroadcastConfig.delta: DeltaWeightBroadcastConfig | None` replaces the `sparse_filesystem` type:
a set `delta` selects SparseFileSystemWeightSender, which takes the delta config (`anchor_interval`). The trainer
rejects `delta` with LoRA.
---
 docs/scaling.md                               |  6 +++--
 .../src/prime_rl/configs/inference.py         |  2 +-
 .../src/prime_rl/configs/rl.py                |  2 +-
 .../src/prime_rl/configs/shared.py            | 24 +++++++++----------
 .../src/prime_rl/configs/trainer.py           |  4 +++-
 src/prime_rl/inference/vllm/server.py         |  1 -
 src/prime_rl/transports/weights/__init__.py   |  6 ++---
 src/prime_rl/transports/weights/sparse.py     |  6 ++---
 8 files changed, 26 insertions(+), 25 deletions(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index fdae880f1fd..267a1f3d5a5 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -290,11 +290,13 @@ By default, the trainer and inference worker each allocate one transfer arena. S
 
 ### Sparse filesystem weight broadcast
 
-With `[weight_broadcast] type = "sparse_filesystem"`, the trainer publishes only the weight elements that changed since the previous broadcast, in the engine's weight format (bf16, or fp8 for engines that quantize online to `fp8_per_block`). Engines read and decode each delta while they keep serving, patch their weights in place during the pause, and verify every update against trainer checksums.
+With `[weight_broadcast.delta]` on the filesystem broadcast, the trainer publishes only the weight elements that changed since the previous broadcast, in the engine's weight format (bf16, or fp8 for engines that quantize online to `fp8_per_block`). Engines read and decode each delta while they keep serving, patch their weights in place during the pause, and verify every update against trainer checksums.
 
 ```toml
 [weight_broadcast]
-type = "sparse_filesystem"
+type = "filesystem"
+
+[weight_broadcast.delta]
 anchor_interval = 50  # full checkpoint every N steps
 ```
 
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/inference.py b/packages/prime-rl-configs/src/prime_rl/configs/inference.py
index 2a0bf1993a9..3b1bc706b07 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/inference.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/inference.py
@@ -186,7 +186,7 @@ def auto_setup_api_server_count(self):
 
 
 class WeightBroadcastConfig(BaseConfig):
-    type: Literal["nccl", "filesystem", "sparse_filesystem", "nixl"] = "filesystem"
+    type: Literal["nccl", "filesystem", "nixl"] = "filesystem"
     """Weight broadcast transport."""
 
 
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/rl.py b/packages/prime-rl-configs/src/prime_rl/configs/rl.py
index 888ea3f8977..f6a74a97554 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py
@@ -632,7 +632,7 @@ def auto_setup_deployment(self):
                 dp = self.inference.vllm.data_parallel_size
                 if self.inference.vllm.api_server_count < dp and not self.inference.vllm.enable_lora:
                     self.inference.vllm.api_server_count = dp
-                if self.weight_broadcast.type in ("nccl", "nixl"):
+                if self.weight_broadcast.type != "filesystem":
                     # Every inference GPU joins the in-memory transfer group, so size it after the DP fill.
                     world_size = dp * self.inference.vllm.tensor_parallel_size
                     self.trainer.weight_broadcast.inference_world_size = world_size
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
index f21be1cb677..ca94226f675 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
@@ -45,21 +45,22 @@ class BaseWeightBroadcastConfig(BaseConfig):
     offered version in time. Raise it for large models on slow shared filesystems."""
 
 
-class FileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig):
-    type: Literal["filesystem"] = "filesystem"
-
-
-class SparseFileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig):
-    """Filesystem broadcast that, between anchors, publishes only the weight elements that changed since the
-    previous broadcast. Engines patch their weights in place and verify every update. Not supported with LoRA."""
-
-    type: Literal["sparse_filesystem"] = "sparse_filesystem"
+class DeltaWeightBroadcastConfig(BaseConfig):
+    """Between anchors, publish only the weight elements that changed since the previous broadcast. Engines patch
+    their weights in place and verify every update. Not supported with LoRA."""
 
     anchor_interval: int = Field(50, ge=1)
     """Publish a full checkpoint every this many steps. Engines that join or fall off the delta chain load
     the last one and replay the deltas since."""
 
 
+class FileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig):
+    type: Literal["filesystem"] = "filesystem"
+
+    delta: DeltaWeightBroadcastConfig | None = None
+    """Publish sparse deltas instead of full checkpoints."""
+
+
 class InMemoryWeightBroadcastConfig(BaseWeightBroadcastConfig):
     host: str = "localhost"
     """Weight transfer host."""
@@ -92,10 +93,7 @@ class NIXLWeightBroadcastConfig(InMemoryWeightBroadcastConfig):
 
 
 WeightBroadcastConfig: TypeAlias = Annotated[
-    FileSystemWeightBroadcastConfig
-    | SparseFileSystemWeightBroadcastConfig
-    | NCCLWeightBroadcastConfig
-    | NIXLWeightBroadcastConfig,
+    FileSystemWeightBroadcastConfig | NCCLWeightBroadcastConfig | NIXLWeightBroadcastConfig,
     Field(discriminator="type"),
 ]
 """Weight transfer from trainer to inference. ``rl`` copies its ``[weight_broadcast]`` to the trainer and orchestrator."""
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
index 7c8732abe44..428bc60a733 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
@@ -716,11 +716,13 @@ class TrainerConfig(BaseTrainerConfig):
 
     @model_validator(mode="after")
     def validate_lora_broadcast(self):
-        if self.model.lora is not None and self.weight_broadcast.type != "filesystem":
+        if self.model.lora is not None and self.weight_broadcast.type in ("nccl", "nixl"):
             raise ValueError(
                 "LoRA requires weight_broadcast.type = 'filesystem': vLLM loads adapters only from a "
                 "PEFT-shaped directory on disk - in-memory transports have no disk artifact to load from."
             )
+        if self.model.lora is not None and self.weight_broadcast.type == "filesystem" and self.weight_broadcast.delta:
+            raise ValueError("weight_broadcast.delta does not support LoRA")
         return self
 
     @model_validator(mode="after")
diff --git a/src/prime_rl/inference/vllm/server.py b/src/prime_rl/inference/vllm/server.py
index b74c32fd44e..a06d1ff3bf9 100644
--- a/src/prime_rl/inference/vllm/server.py
+++ b/src/prime_rl/inference/vllm/server.py
@@ -36,7 +36,6 @@ def models(request: Request) -> OpenAIServingModels:
 WORKER_EXTENSION_CLS = {
     "nccl": "prime_rl.inference.vllm.worker.nccl.NCCLWeightUpdateWorker",
     "filesystem": "prime_rl.inference.vllm.worker.filesystem.FileSystemWeightUpdateWorker",
-    "sparse_filesystem": "prime_rl.inference.vllm.worker.filesystem.FileSystemWeightUpdateWorker",
     "nixl": "prime_rl.inference.vllm.worker.nixl.NIXLWeightUpdateWorker",
 }
 
diff --git a/src/prime_rl/transports/weights/__init__.py b/src/prime_rl/transports/weights/__init__.py
index d967cd89a42..32f991781e7 100644
--- a/src/prime_rl/transports/weights/__init__.py
+++ b/src/prime_rl/transports/weights/__init__.py
@@ -29,10 +29,10 @@ def setup_weight_sender(
 ) -> WeightSender:
     if config.type == "nccl":
         return NCCLWeightSender(output_dir, config, torch.cuda.current_device())
+    elif config.type == "filesystem" and config.delta is not None:
+        return SparseFileSystemWeightSender(output_dir, config.timeout, config.delta)
     elif config.type == "filesystem":
         return FileSystemWeightSender(output_dir, config, lora_config)
-    elif config.type == "sparse_filesystem":
-        return SparseFileSystemWeightSender(output_dir, config)
     elif config.type == "nixl":
         return NIXLWeightSender(output_dir, config, parallel_dims)
     else:
@@ -47,7 +47,7 @@ def setup_weight_receiver(
 ) -> WeightReceiver:
     if config.type == "nccl":
         return NCCLWeightReceiver(broadcast_dir, config, admin_plane, model_name)
-    elif config.type in ("filesystem", "sparse_filesystem"):
+    elif config.type == "filesystem":
         return FileSystemWeightReceiver(broadcast_dir, config, admin_plane, model_name)
     elif config.type == "nixl":
         return NIXLWeightReceiver(broadcast_dir, config, admin_plane, model_name)
diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 520554cb231..0ced3c8a1a6 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -19,7 +19,7 @@
 from torch.distributed.tensor import DTensor
 from torch.distributed.tensor._utils import compute_local_shape_and_global_offset
 
-from prime_rl.configs.shared import SparseFileSystemWeightBroadcastConfig
+from prime_rl.configs.shared import DeltaWeightBroadcastConfig
 from prime_rl.trainer.models.fusions import get_model_packed_parameters, join_module_path
 from prime_rl.transports.weights.base import WeightSender
 from prime_rl.transports.weights.sparse_kernels import Pieces, build_pieces, count_changes, write_changes
@@ -1008,8 +1008,8 @@ class SparseFileSystemWeightSender(WeightSender):
 
     gathers_full_weights = False
 
-    def __init__(self, output_dir: Path, config: SparseFileSystemWeightBroadcastConfig):
-        super().__init__(output_dir, config.timeout)
+    def __init__(self, output_dir: Path, timeout: int, config: DeltaWeightBroadcastConfig):
+        super().__init__(output_dir, timeout)
         self.config = config
         self.encoder: SparseDeltaEncoder | None = None
         self.last_step: int | None = None

From 2616dfa764377e4e50977cb3bbe60624a3a923aa Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 28 Oct 2026 16:58:56 +0000
Subject: [PATCH 71/88] fix(weights): a tied weight's other names carry no
 sparse deltas; `delta.interval` names the anchor interval

- Tied weights (e.g. a tied lm_head) appear under several names; only the first carries deltas, checksums and anchor
  slices, so their changes are not sent twice. Engines hold a tied weight once.
- `DeltaWeightBroadcastConfig.anchor_interval` is now `interval`: a full checkpoint (anchor) every this many steps,
  deltas against the previous step in between.
---
 docs/scaling.md                                        |  4 ++--
 .../prime-rl-configs/src/prime_rl/configs/shared.py    |  7 ++++---
 src/prime_rl/transports/weights/sparse.py              | 10 ++++++----
 3 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index 267a1f3d5a5..282057a3ab2 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -297,12 +297,12 @@ With `[weight_broadcast.delta]` on the filesystem broadcast, the trainer publish
 type = "filesystem"
 
 [weight_broadcast.delta]
-anchor_interval = 50  # full checkpoint every N steps
+interval = 50  # full checkpoint (anchor) every N steps
 ```
 
 - RL: deltas come from the optimizer step, with every optimizer and offload mode. Right before the step updates a chunk of parameters, their shards still hold the previous broadcast's weights; the trainer diffs the old and new values of each chunk as the step goes and writes the delta out. No copy of the weights is kept, except of the few tensors that are not optimizer parameters (buffers, frozen parameters, fp8 tensors whose 128-row blocks straddle trainer ranks). The trainer checks that no weight changed outside the optimizer step.
 - A broadcast that does not follow every optimizer step (the startup broadcast of a resumed run, SFT online evals) publishes a full checkpoint.
-- Anchors (full checkpoints) are written every `anchor_interval` steps, from the optimizer step. They slow that step down, so prefer a larger `anchor_interval`. An engine that joins, restarts or fails an update loads the latest complete anchor, reading only the rows it holds, and replays the deltas since.
+- Anchors (full checkpoints) are written every `interval` steps, from the optimizer step. They slow that step down, so prefer a larger `interval`. An engine that joins, restarts or fails an update loads the latest complete anchor, reading only the rows it holds, and replays the deltas since.
 - Not supported: LoRA, fp8 tensors whose columns are not a multiple of 128, and engines that transform weights after loading other than online `fp8_per_block` (e.g. FlashInfer TRTLLM/CUTLASS MoE, Marlin). fp8 deltas need the same vLLM and torch versions on the trainer and the engines.
 
 ### Custom Templates
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
index ca94226f675..4fab10acf31 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py
@@ -49,9 +49,10 @@ class DeltaWeightBroadcastConfig(BaseConfig):
     """Between anchors, publish only the weight elements that changed since the previous broadcast. Engines patch
     their weights in place and verify every update. Not supported with LoRA."""
 
-    anchor_interval: int = Field(50, ge=1)
-    """Publish a full checkpoint every this many steps. Engines that join or fall off the delta chain load
-    the last one and replay the deltas since."""
+    interval: int = Field(50, ge=1)
+    """Publish a full checkpoint (anchor) every this many steps; the steps in between publish their delta
+    against the previous step. Engines that join or fall off the delta chain load the last anchor and replay
+    the deltas since."""
 
 
 class FileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig):
diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index 0ced3c8a1a6..bba287a2d78 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -235,11 +235,13 @@ def local_shards(self, model: nn.Module, is_master: bool, gather: bool = False)
         module parameters during forward, but the sharded storage stays put). With ``gather``, ``gathered`` is
         ``(previous, current)`` full bf16 tensors of fp8 tensors whose blocks straddle ranks, all-gathered here
         because ``encode`` runs no collectives."""
-        shards = []
+        shards, seen = [], set()
         for key, value in trainer_tensors(model).items():
             fqn = self.key_fqn[key]
-            if fqn not in self.segments:
+            # A tied weight's other names (e.g. a tied lm_head) carry nothing: engines hold it once
+            if fqn not in self.segments or id(value) in seen:
                 continue
+            seen.add(id(value))
             shard = self._local_shard(key, value, is_master)
             empty = shard is None or shard[0].numel() == 0
             gathered = None
@@ -1003,7 +1005,7 @@ class SparseFileSystemWeightSender(WeightSender):
     written as the step goes; the rest is finished in the background while the next step trains, and the trainer
     waits only before its next final backward, until that work has read the weights. A broadcast that does not
     follow an optimizer step collected this way (the first one, or any broadcast without ``diff_in_optimizer``)
-    publishes a full checkpoint. Anchors (full checkpoints every ``anchor_interval`` steps) serve engines that fall
+    publishes a full checkpoint. Anchors (full checkpoints every ``interval`` steps) serve engines that fall
     off the delta chain. The writer thread commits a step once every rank's files are on disk."""
 
     gathers_full_weights = False
@@ -1072,7 +1074,7 @@ def _arm(self, step: int) -> None:
         """Collect v{step}'s delta, and its anchor slices at anchor steps, in the coming optimizer step."""
         encoder = self.encoder
         anchor = None
-        if step % self.config.anchor_interval == 0:
+        if step % self.config.interval == 0:
             # Written next to the step dir, which the broadcast handshake recreates, and moved in once complete
             anchor = AnchorWriter(encoder.anchor_slices(), self.step_dir(step), self.world.rank)
         # Parts written during the step go next to the step dir too, and are moved in when it is published

From bbfe2f152bb9806c7fc666aba85a8a3377e8c355 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 28 Oct 2026 17:07:55 +0000
Subject: [PATCH 72/88] perf(weights): read the sparse delta before the
 dispatcher stops for the weight update

`WeightReceiver.prepare` does the work of an update that needs no engine pause; the watcher awaits it before
`on_version_pending`. The filesystem receiver acknowledges the version and has the engines read and decode a
sparse delta there, so rollouts keep being dispatched meanwhile. Before, train dispatch stopped for the whole
receive (0.37-0.75 s on Qwen3-0.6B/8B vs 0.33-0.40 s with NCCL), which left fewer rollouts in flight.
---
 src/prime_rl/orchestrator/watcher.py          |  2 ++
 src/prime_rl/transports/weights/base.py       |  4 ++++
 src/prime_rl/transports/weights/filesystem.py | 17 ++++++++++++++---
 3 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/src/prime_rl/orchestrator/watcher.py b/src/prime_rl/orchestrator/watcher.py
index baef29e0415..8aea5494031 100644
--- a/src/prime_rl/orchestrator/watcher.py
+++ b/src/prime_rl/orchestrator/watcher.py
@@ -89,6 +89,8 @@ async def apply_policy_update(self, next_step: int) -> None:
             # Record the published version before notifying pending observers.
             # ``policy.version`` advances only after inference applies it.
             self.ckpt_step = next_step
+            # Rollouts keep being dispatched while the engines can still serve
+            await self.receiver.prepare(next_step)
 
             # Let observers cancel rollouts BEFORE pausing the inference engines.
             # Aborting a rollout triggers vLLM's KV-connector cleanup (NIXL's
diff --git a/src/prime_rl/transports/weights/base.py b/src/prime_rl/transports/weights/base.py
index 41d57d10886..a915409f0fa 100644
--- a/src/prime_rl/transports/weights/base.py
+++ b/src/prime_rl/transports/weights/base.py
@@ -174,6 +174,10 @@ def _ack(self, step: int) -> None:
         """Acknowledge the offered version — unblocks the waiting trainer."""
         (self.step_dir(step) / RECEIVER_READY_MARKER).touch()
 
+    async def prepare(self, step: int) -> None:
+        """Work on the offered v{step} that needs no engine pause, done before the consumer stops dispatching for
+        the update. ``receive`` does it too if it was not done."""
+
     @abstractmethod
     async def receive(self, step: int) -> None:
         """Acknowledge the offered v{step} and move the engines onto it."""
diff --git a/src/prime_rl/transports/weights/filesystem.py b/src/prime_rl/transports/weights/filesystem.py
index 7b4fe538a52..4d59ab4f69c 100644
--- a/src/prime_rl/transports/weights/filesystem.py
+++ b/src/prime_rl/transports/weights/filesystem.py
@@ -66,12 +66,17 @@ class FileSystemWeightReceiver(WeightReceiver):
     live traffic — an in-place adapter reload is a vLLM-native op that needs
     no engine pause; a full checkpoint pauses the engines for the load."""
 
-    async def receive(self, step: int) -> None:
+    prepared_step: int | None = None
+
+    async def prepare(self, step: int) -> None:
+        """Acknowledge v{step} and wait for its weights while the engines keep serving: a sparse delta's manifest
+        comes first and its part files as the trainer encodes, so the engines read and decode them now and the
+        pause only applies the delta."""
+        if self.prepared_step == step:
+            return
         weights_dir = self.step_dir(step)
         manifest, finished = weights_dir / DELTA_DIR / MANIFEST, weights_dir / FINISHED_MARKER
         self._ack(step)
-        # A sparse delta's manifest comes first and its part files as the trainer encodes: the engines
-        # read and decode them while they keep serving, and the pause only applies the delta.
         while not (manifest.exists() or finished.exists()):
             await asyncio.sleep(0.05)
             revalidate(manifest.parent)
@@ -79,6 +84,12 @@ async def receive(self, step: int) -> None:
             start = time.perf_counter()
             await self.admin_plane.prefetch_weights(weights_dir)
             self.logger.debug(f"Prefetched sparse weight delta v{step} in {time.perf_counter() - start:.2f}s")
+        self.prepared_step = step
+
+    async def receive(self, step: int) -> None:
+        await self.prepare(step)
+        weights_dir = self.step_dir(step)
+        finished = weights_dir / FINISHED_MARKER
         while not finished.exists():
             await asyncio.sleep(0.05)
             revalidate(weights_dir)

From 97be3054fc6bc8edeae627dba0868f342f72b8e7 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 28 Oct 2026 19:14:38 +0000
Subject: [PATCH 73/88] fix(weights): anchors and full checkpoints keep every
 name of a tied weight

Only the deltas go out once under a tied weight's first name. Its other names are written to full checkpoints and
to anchors (and their checksums), from the same values.
---
 src/prime_rl/transports/weights/sparse.py | 46 ++++++++++++++---------
 1 file changed, 29 insertions(+), 17 deletions(-)

diff --git a/src/prime_rl/transports/weights/sparse.py b/src/prime_rl/transports/weights/sparse.py
index bba287a2d78..0d4476faad5 100644
--- a/src/prime_rl/transports/weights/sparse.py
+++ b/src/prime_rl/transports/weights/sparse.py
@@ -106,7 +106,13 @@ def __init__(self, model: nn.Module):
         self.wire_dtypes: dict[str, torch.dtype] = {}
         self.key_fqn: dict[str, str] = {}
         meta: dict[str, torch.Tensor] = {}
+        # A tied weight's other names: its deltas go out once, under the first name, and anchors carry every name
+        self.aliases: dict[str, list[str]] = defaultdict(list)
+        first_key: dict[int, str] = {}
         for key, value in trainer_tensors(model).items():
+            if id(value) in first_key:
+                self.aliases[first_key[id(value)]].append(key)
+            first_key.setdefault(id(value), key)
             dtype = resolve_wire_dtype(keep_in_fp32, key, torch.bfloat16) if isinstance(value, DTensor) else value.dtype
             self.wire_dtypes[key] = dtype
             self.key_fqn[key] = resolve_fqn(model, key)
@@ -230,18 +236,19 @@ def _find_gathered(self, model: nn.Module, is_master: bool) -> set[str]:
         return {key for key, flag in zip(keys, flags.tolist()) if flag}
 
     @torch.no_grad()
-    def local_shards(self, model: nn.Module, is_master: bool, gather: bool = False) -> list[tuple]:
+    def local_shards(
+        self, model: nn.Module, is_master: bool, gather: bool = False, aliases: bool = False
+    ) -> list[tuple]:
         """This rank's ``(key, fqn, shard, global flat offset, gathered)``, taken on the training thread (FSDP swaps
         module parameters during forward, but the sharded storage stays put). With ``gather``, ``gathered`` is
         ``(previous, current)`` full bf16 tensors of fp8 tensors whose blocks straddle ranks, all-gathered here
-        because ``encode`` runs no collectives."""
-        shards, seen = [], set()
+        because ``encode`` runs no collectives. A tied weight's other names only with ``aliases``."""
+        shards = []
+        alias_keys = {alias for names in self.aliases.values() for alias in names}
         for key, value in trainer_tensors(model).items():
             fqn = self.key_fqn[key]
-            # A tied weight's other names (e.g. a tied lm_head) carry nothing: engines hold it once
-            if fqn not in self.segments or id(value) in seen:
+            if fqn not in self.segments or (key in alias_keys and not aliases):
                 continue
-            seen.add(id(value))
             shard = self._local_shard(key, value, is_master)
             empty = shard is None or shard[0].numel() == 0
             gathered = None
@@ -292,7 +299,7 @@ def write_full_checkpoint(self, model: nn.Module, is_master: bool, step_dir: Pat
         """Write this rank's slices of a full HF checkpoint of the live weights, one tensor at a time."""
         shards = [
             (key, shard, offset, self._slices(fqn, shard.numel(), offset))
-            for key, fqn, shard, offset, _ in self.local_shards(model, is_master)
+            for key, fqn, shard, offset, _ in self.local_shards(model, is_master, aliases=True)
         ]
         writer = AnchorWriter(
             [
@@ -344,7 +351,7 @@ def _write_optimizer_anchor(self, shards: list[tuple]) -> None:
             if key in self.residual or (key in self.stepped and not self.full_offload):
                 continue
             values = shard.reshape(-1).to(self.wire_dtypes[key])
-            self._write_anchor_slices(key, fqn, len(values), offset, values)
+            self._write_anchor_slices(key, len(values), offset, values)
 
     def wait_for_diffs(self) -> None:
         for diff in self.diffs:
@@ -734,6 +741,8 @@ def anchor_slices(self) -> list[tuple[str, list[int], int, int, torch.dtype]]:
             key: (fqn, numel, offset) for entries in self.param_keys.values() for key, fqn, numel, offset, _ in entries
         }
         shards.update({key: (self.key_fqn[key], numel, offset) for key, (numel, offset) in self.shards.items()})
+        for key, (_, numel, offset) in list(shards.items()):
+            shards.update({alias: (self.key_fqn[alias], numel, offset) for alias in self.aliases[key]})
         return [
             (name, shape, start, end - begin, self.wire_dtypes[key])
             for key, (fqn, numel, offset) in shards.items()
@@ -843,7 +852,7 @@ def _collect(self, launches: list[tuple], keys: list[str], old: torch.Tensor, re
                 with torch.cuda.stream(self.diff_stream):
                     for key, fqn, numel, offset, view in written:
                         values = view.reshape(-1).to(self.wire_dtypes[key])
-                        self._write_anchor_slices(key, fqn, numel, offset, values)
+                        self._write_anchor_slices(key, numel, offset, values)
                     self.diff_stream.synchronize()
             self.stepped.update(keys)
         finally:
@@ -961,7 +970,7 @@ def _diff_slot(self, items: list[tuple], ready: torch.cuda.Event, release: Calla
                         self.step_parts += self._diff_shard(key, fqn, current, offset, None, old, None, chunk)
                         self.checksums[key] = self._checksum(key, offset, as_bits(current))
                         if self.anchor is not None and on_gpu:
-                            self._write_anchor_slices(key, fqn, numel, offset, current)
+                            self._write_anchor_slices(key, numel, offset, current)
                         self.stepped.add(key)
                 self.diff_stream.synchronize()
         finally:
@@ -972,15 +981,17 @@ def _diff_slot(self, items: list[tuple], ready: torch.cuda.Event, release: Calla
                 self.emit(self._host_part(self.step_parts, self.diff_staging))
             self.step_parts = []
 
-    def _write_anchor_slices(self, key: str, fqn: str, numel: int, offset: int, values: torch.Tensor) -> None:
+    def _write_anchor_slices(self, key: str, numel: int, offset: int, values: torch.Tensor) -> None:
         """Write a GPU shard's anchor slices, and their checksums, through the pinned staging buffer: the slices in
-        each staged window in parallel, in pieces of up to 32 MiB."""
-        self.anchor.add_sums(*self._value_sums(key, offset, as_bits(values)))
+        each staged window in parallel, in pieces of up to 32 MiB. Also under the names of a tied weight."""
         data, itemsize = values.view(torch.uint8), values.element_size()
-        slices = [
-            (name, start, begin * itemsize, end * itemsize)
-            for name, _, start, begin, end in self._slices(fqn, numel, offset)
-        ]
+        slices = []
+        for name_key in (key, *self.aliases[key]):
+            self.anchor.add_sums(*self._value_sums(name_key, offset, as_bits(values)))
+            slices += [
+                (name, start, begin * itemsize, end * itemsize)
+                for name, _, start, begin, end in self._slices(self.key_fqn[name_key], numel, offset)
+            ]
         window, piece = self.diff_staging.numel(), 1 << 25
         for at in range(0, data.numel(), window):
             count = min(window, data.numel() - at)
@@ -1235,6 +1246,7 @@ def diffed(scale_changes: int) -> None:
         if anchor is not None:
             # The residual's previous bits hold v{step} until the next encode, which runs on this thread
             source = {key: (self.encoder.shards[key][1], bits) for key, bits in encoder.previous.items()}
+            source.update({alias: source[key] for key in list(source) for alias in encoder.aliases[key]})
             for key, (offset, bits) in source.items():
                 anchor.add_sums(*encoder._value_sums(key, offset, bits))
             anchor.write_parts(encoder.anchor_parts(source), encoder.stream_to_host)

From adbf95d520ddfb256b0af706882542d394add773 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 28 Oct 2026 20:46:39 +0000
Subject: [PATCH 74/88] fix(weights): anchors and full checkpoints load tensors
 the engine casts on load

A tensor vLLM casts on load (e.g. GLM-5's fp32 `indexer.k_norm` from a bf16 checkpoint) was left out of
`row_ranges`, so an anchor or replay load kept the engine's stale values without a word. Its writes are now
recorded like the others and loaded by value with the same cast; a delta that changes it still raises.
---
 .../inference/vllm/worker/weight_transfer.py  | 20 ++++++-----
 .../inference/test_checkpoint_index_map.py    | 33 +++++++++++++++++++
 2 files changed, 45 insertions(+), 8 deletions(-)
 create mode 100644 tests/unit/inference/test_checkpoint_index_map.py

diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py
index fa6a5458e34..a68837a663c 100644
--- a/src/prime_rl/inference/vllm/worker/weight_transfer.py
+++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py
@@ -44,7 +44,8 @@ class CheckpointIndexMap:
     checkpoint format; their live params must keep that layout (one byte per
     checkpoint element). Supported: every write is one ``copy_`` of a row/column
     block of a 1-D or 2-D checkpoint tensor. Tensors the loader casts (e.g. an
-    fp32 router from a bf16 checkpoint) are only allowed to stay unchanged.
+    fp32 norm from a bf16 checkpoint) are loaded by value from full checkpoints
+    and anchors; a delta that changes them raises.
     """
 
     def __init__(self, model: Module, tensors: dict[str, tuple[list[int], str]]):
@@ -63,8 +64,8 @@ def record_copy_(destination: torch.Tensor, source: torch.Tensor, non_blocking:
             source_2d = source.view(1, -1) if source.dim() == 1 else source
             destination_2d = destination.view(1, -1) if destination.dim() == 1 else destination
             if name is not None and source.dtype != destination.dtype:
+                # Cast on load (e.g. an fp32 norm from a bf16 checkpoint): loaded by value, no deltas
                 unpatchable.add(name)
-                return destination
             if (
                 name is None
                 or param_name is None
@@ -124,7 +125,7 @@ def placeholders():
         self.fp8: dict[str, str] = {}
         self.unpatchable = unpatchable
         # Checkpoint tensors this worker holds; unpatchable ones too, so a change to them is caught in apply_delta
-        self.names = set(self.records) | unpatchable
+        self.names = set(self.records)
 
     def held_checksums(self) -> tuple[list[str], torch.Tensor, torch.Tensor]:
         """``bits_checksum`` of the live param elements each checkpoint tensor was written to, weighted by
@@ -278,11 +279,10 @@ def _view(self, index: int, rows: int, cols: int, offset: int, strides) -> torch
 
     def row_ranges(self) -> dict[str, tuple[int, int]]:
         """Checkpoint tensor -> the rows ``[begin, end)`` of its ``[rows, num_cols]`` view this worker holds
-        (fp8 tensors: whole 128-row blocks), for ``load_rows``. Tensors the engine casts on load are left out:
-        they never change."""
+        (fp8 tensors: whole 128-row blocks), for ``load_rows``."""
         ranges = {}
         for name, writes in self.records.items():
-            if name.endswith("_scale_inv") or name in self.unpatchable:
+            if name.endswith("_scale_inv"):
                 continue  # scales are set with their weights
             begin, end = min(write[0] for write in writes), max(write[0] + write[1] for write in writes)
             if name in self.fp8:
@@ -295,7 +295,7 @@ def row_ranges(self) -> dict[str, tuple[int, int]]:
     def load_rows(self, name: str, rows: torch.Tensor, begin: int) -> None:
         """Set the params from rows ``begin:`` of checkpoint tensor ``name`` (the ``row_ranges`` rows, on the
         params' device): the counterpart of ``apply_delta`` for full checkpoints. fp8 tensors are quantized
-        like the engine does, block scales included."""
+        like the engine does, block scales included; tensors the engine casts on load are cast the same way."""
         sources = [(name, rows, begin)]
         if name in self.fp8:
             quantized, scales = quantize_fp8_blocks(rows, self.fp8[name] == "ue8m0")
@@ -303,7 +303,11 @@ def load_rows(self, name: str, rows: torch.Tensor, begin: int) -> None:
         for source_name, source, first in sources:
             for row0, rows_, col0, cols, index, offset, *strides in self.records[source_name]:
                 block = source[row0 - first : row0 - first + rows_, col0 : col0 + cols]
-                self._view(index, rows_, cols, offset, strides).copy_(as_bits(block))
+                if source_name in self.unpatchable:
+                    flat = self.params[index].data.view(-1)
+                    flat.as_strided((rows_, cols), tuple(strides), offset).copy_(block)
+                else:
+                    self._view(index, rows_, cols, offset, strides).copy_(as_bits(block))
 
     @torch.no_grad()
     def apply_delta(self, entries: list[tuple[str, int]], indices: torch.Tensor, delta: torch.Tensor) -> None:
diff --git a/tests/unit/inference/test_checkpoint_index_map.py b/tests/unit/inference/test_checkpoint_index_map.py
new file mode 100644
index 00000000000..d2cd0dbf2b2
--- /dev/null
+++ b/tests/unit/inference/test_checkpoint_index_map.py
@@ -0,0 +1,33 @@
+import pytest
+import torch
+from torch import nn
+
+from prime_rl.inference.vllm.worker.weight_transfer import CheckpointIndexMap
+
+
+class Toy(nn.Module):
+    """An engine model that keeps a norm in fp32 (cast on load from the bf16 checkpoint)."""
+
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(4, 8, dtype=torch.bfloat16))
+        self.norm = nn.Parameter(torch.zeros(8, dtype=torch.float32))
+
+    def load_weights(self, weights):
+        for name, tensor in weights:
+            getattr(self, name).data.copy_(tensor)
+
+
+def test_cast_tensors_load_by_value_and_reject_deltas():
+    model = Toy()
+    index_map = CheckpointIndexMap(model, {"weight": ([4, 8], "bfloat16"), "norm": ([8], "bfloat16")})
+    assert index_map.row_ranges() == {"weight": (0, 4), "norm": (0, 1)}
+
+    norm, weight = torch.randn(1, 8, dtype=torch.bfloat16), torch.randn(4, 8, dtype=torch.bfloat16)
+    index_map.load_rows("norm", norm, 0)
+    index_map.load_rows("weight", weight, 0)
+    assert torch.equal(model.norm.data, norm.view(-1).float())
+    assert torch.equal(model.weight.data, weight)
+
+    with pytest.raises(NotImplementedError):
+        index_map.apply_delta([("norm", 1)], torch.tensor([0]), torch.tensor([1], dtype=torch.int16))

From 5286d0a6b2c193bc7fbb5717c26b6a193a2f7bad Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Mon, 2 Nov 2026 19:34:07 +0000
Subject: [PATCH 75/88] chore: bump vLLM to 0.33.0 (torch 2.14.0)

vLLM 0.33.0 is compiled against and pins torch 2.14.0, so this also moves
the torch stack: torch 2.14.0, torchvision 0.29.0, triton 3.8.0, NCCL 2.30.7,
flashinfer 0.7.2.post1, transformers 5.20.2 (what vLLM 0.33 CI tests),
nvidia-cudnn-frontend 1.31.0, and NIXL 1.5.0 (required by vLLM's NIXL
connector from 0.33).

Breakages fixed:
- serving_tokens: upstream reads sampling_params.logprobs before consuming
  outputs and passes derender_request; decide logprob packing up front.
- server: cli_args moved to vllm.entrypoints.common.
- qwen38 PLE overlap is now a method; GPT-OSS _load_weights_other lost
  heads_per_rank/head_start and sinks now load via their own loader.
- routed experts: _get_local_prefix_cache_hit takes max_cache_hit_length;
  P/D decode requests skip the new retained-row caps.
- sampling_order: accept has_topp_only_rows.

Removed (fixed upstream): DeepSeek V4 C128 boundary (#59709), DeepSeek V4
attn_sink loading (#57797), aux-output store eviction rescan (#62392).

flash-mla has no torch 2.14 build; the trainer now uses the FlashMLA vLLM
bundles (also on aarch64). flash-attn moves to its torch 2.14 wheels.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 Dockerfile.cuda                               |   8 +-
 .../src/prime_rl/configs/trainer.py           |   2 +-
 pyproject.toml                                |  55 +--
 scripts/install_nixl_from_source.sh           |   5 +-
 src/prime_rl/inference/patches.py             | 136 ++----
 .../inference/vllm/gpt_oss_weight_loading.py  |  26 +-
 .../inference/vllm/qwen38_weight_loading.py   |   8 +-
 src/prime_rl/inference/vllm/sampling_order.py |   9 +-
 src/prime_rl/inference/vllm/server.py         |   2 +-
 src/prime_rl/inference/vllm/serving_tokens.py |  50 +-
 .../trainer/activation_checkpointing.py       |   3 +
 .../trainer/models/kernels/__init__.py        |  13 +-
 .../kernels/deepseek_v4/dsv4_sparse_attn.py   |   4 +-
 .../trainer/models/kernels/sparse_mla_fwd.py  |   2 +-
 tests/unit/inference/test_routed_experts.py   |  37 +-
 .../test_routed_experts_kv_offload.py         |  33 ++
 .../unit/inference/test_routed_experts_pd.py  |  31 +-
 tests/unit/inference/test_serving_tokens.py   |  38 ++
 .../train/models/test_deepseek_v4_kernels.py  |  11 +-
 uv.lock                                       | 446 +++++++++---------
 20 files changed, 467 insertions(+), 452 deletions(-)

diff --git a/Dockerfile.cuda b/Dockerfile.cuda
index baa6c45820f..c3e657a2775 100644
--- a/Dockerfile.cuda
+++ b/Dockerfile.cuda
@@ -110,7 +110,8 @@ RUN --mount=type=cache,target=/usr/local/share/uv/cache \
 # build RPATH in wheel libraries, which dropped RUNPATH=/opt/ucx/lib and failed
 # the linkage check below on every build. Pin the validated backend.
 ARG INCLUDE_NIXL_FROM_SOURCE=1
-ARG NIXL_VERSION=0.10.1
+# NIXL tags carry a `v` prefix from 1.0 on.
+ARG NIXL_VERSION=v1.5.0
 RUN if [ "$INCLUDE_NIXL_FROM_SOURCE" = "1" ]; then \
         VIRTUAL_ENV=/app/.venv uv pip install pip \
         && git clone --depth 1 --branch ${NIXL_VERSION} https://github.com/ai-dynamo/nixl.git /tmp/nixl \
@@ -243,10 +244,7 @@ COPY --from=builder --chown=appuser:appuser \
     --exclude=lib/python3.12/site-packages/flash_attn \
     --exclude=lib/python3.12/site-packages/flash_attn_2_cuda*.so \
     /app/.venv /app/.venv
-RUN /app/.venv/bin/python -c 'from importlib.metadata import version; [version(name) for name in ("deep-ep", "deep-gemm", "mamba-ssm", "nixl", "prime-kernels", "quack-kernels", "vllm", "vllm-router")]' \
-    && if [ "$(uname -m)" = "x86_64" ]; then \
-        /app/.venv/bin/python -c 'from importlib.metadata import version; version("flash-mla")'; \
-    fi
+RUN /app/.venv/bin/python -c 'from importlib.metadata import version; [version(name) for name in ("deep-ep", "deep-gemm", "mamba-ssm", "nixl", "prime-kernels", "quack-kernels", "vllm", "vllm-router")]'
 
 # Place executables in the environment at the front of the path
 ENV PATH="/app/.venv/bin:$PATH"
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
index 4be7a8daaeb..fb0d0c90c7b 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
@@ -324,7 +324,7 @@ class ModelConfig(BaseModelConfig):
     quantization: QuantizationConfig | None = None
 
     dsa_backend: DSABackend = "cudnn_flashmla"
-    """Kernels for the sparse attention of ``glm_moe_dsa`` (DSA sparse MLA, e.g. GLM-5) and ``deepseek_v4``. ``cudnn_flashmla`` (default) runs the FlashMLA sparse prefill forward with the cuDNN frontend CuTe-DSL backward; it needs an SM90, SM100 or SM103 GPU, ``flash-mla`` (installed with the ``gpu`` extra on x86_64), 64 or 128 attention heads and head_dim 512, and setup fails otherwise. ``tilelang`` runs the TileLang forward and backward."""
+    """Kernels for the sparse attention of ``glm_moe_dsa`` (DSA sparse MLA, e.g. GLM-5) and ``deepseek_v4``. ``cudnn_flashmla`` (default) runs the FlashMLA sparse prefill forward with the cuDNN frontend CuTe-DSL backward; it needs an SM90, SM100 or SM103 GPU, vLLM's bundled FlashMLA (installed with the ``gpu`` extra), 64 or 128 attention heads and head_dim 512, and setup fails otherwise. ``tilelang`` runs the TileLang forward and backward."""
 
     dsv4_kv_precision: Literal["bf16", "fp8"] = "fp8"
     """Precision DeepSeek V4 attention reads keys and values at. ``fp8`` rounds them as vLLM's ``fp8_ds_mla`` KV cache stores them, with a straight-through gradient."""
diff --git a/pyproject.toml b/pyproject.toml
index 34f23c524b3..f56e284f0e3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -13,7 +13,7 @@ dependencies = [
     "loguru>=0.7.3",
     "numpy>=2.2.6",
     "sympy>=1.13.3",
-    "transformers==5.19.0",
+    "transformers==5.20.2",
     "wandb>=0.26.1",
     "wandb-workspaces>=0.4.3",
     "prime>=0.18.0",
@@ -52,30 +52,27 @@ prime_rl = "prime_rl.inference.patches:apply_shared_vllm_patches"
 
 [project.optional-dependencies]
 gpu = [
-    "torch>=2.13.0 ; sys_platform == 'linux'",
+    "torch>=2.14.0 ; sys_platform == 'linux'",
     "torchvision ; sys_platform == 'linux'",
     "torchaudio ; sys_platform == 'linux'",
     "torchdata>=0.11.0 ; sys_platform == 'linux'",
-    "vllm>=0.32.0 ; sys_platform == 'linux'",
+    "vllm>=0.33.0 ; sys_platform == 'linux'",
     "mooncake-transfer-engine>=0.3.10.post2 ; sys_platform == 'linux'",
     "torchtitan ; sys_platform == 'linux'",
     "dion ; sys_platform == 'linux'",
     "tilelang>=0.1.8 ; sys_platform == 'linux'",
     # DSA sparse MLA backward kernel (CuTe DSL); >=1.29 matches the pinned CUTLASS DSL 4.7.
-    "nvidia-cudnn-frontend>=1.30.0 ; sys_platform == 'linux'",
-    # DSA sparse MLA forward for the default `model.dsa_backend = "cudnn_flashmla"` (sm_90a + sm_100f).
-    # FlashMLA is not on PyPI and only an x86_64 wheel is published.
-    "flash-mla ; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "nvidia-cudnn-frontend>=1.31.0 ; sys_platform == 'linux'",
     "deep-select ; sys_platform == 'linux'",
     "flash-linear-attention>=0.6.0 ; sys_platform == 'linux'",
     "mamba-ssm>=2.3.0 ; sys_platform == 'linux'",
     "nvidia-ml-py>=12.575.51 ; sys_platform == 'linux'",
     "torchao>=0.18.0 ; sys_platform == 'linux'",
-    "nixl-cu13==0.10.1 ; sys_platform == 'linux'",
+    "nixl-cu13==1.5.0 ; sys_platform == 'linux'",
     # Precompiled FlashInfer kernels for vLLM's flashinfer-python pin. Without them every fresh node
     # downloads the trtllm-gen cubins and JIT-builds its modules into its node-local cache at startup.
-    "flashinfer-cubin==0.7.0.post1 ; sys_platform == 'linux'",
-    "flashinfer-jit-cache==0.7.0.post1+cu130 ; sys_platform == 'linux'",
+    "flashinfer-cubin==0.7.2.post1 ; sys_platform == 'linux'",
+    "flashinfer-jit-cache==0.7.2.post1+cu130 ; sys_platform == 'linux'",
     "modelexpress==0.8.0.dev20261022 ; sys_platform == 'linux'",
 ]
 dashboard = [
@@ -83,13 +80,13 @@ dashboard = [
     "uvicorn>=0.30",
 ]
 dynamo = ["ai-dynamo==1.4.2 ; sys_platform == 'linux'"]
-# Both wheels are the same flash-attn 2.8.3 / cu130 / torch 2.13 / cp312 build and carry
+# Both wheels are the same flash-attn 2.8.3 / cu130 / torch 2.14 / cp312 build and carry
 # sm_80 / sm_90 / sm_100 / sm_120 cubins. `attn="auto"` picks FA4 on sm_100 and FA3 on
 # sm_90, so FA2 is the resolved backend on sm_120 (see `_resolve_attn_impl` in
 # trainer/model.py).
 flash-attn = [
-    "flash-attn @ https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.47/flash_attn-2.8.3+cu130torch2.13-cp312-cp312-linux_x86_64.whl ; platform_machine == 'x86_64'",
-    "flash-attn @ https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.48/flash_attn-2.8.3+cu130torch2.13-cp312-cp312-linux_aarch64.whl ; platform_machine == 'aarch64'",
+    "flash-attn @ https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.10.0/flash_attn-2.8.3+cu130torch2.14-cp312-cp312-linux_x86_64.whl ; platform_machine == 'x86_64'",
+    "flash-attn @ https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.10.1/flash_attn-2.8.3+cu130torch2.14-cp312-cp312-linux_aarch64.whl ; platform_machine == 'aarch64'",
 ]
 flash-attn-3 = ["flash_attn_3 ; sys_platform == 'linux'"]
 
@@ -99,7 +96,7 @@ flash-attn-cute = [
 disagg = [
     "deep-ep ; sys_platform == 'linux'",
     "deep-gemm ; sys_platform == 'linux'",
-    "nixl ; sys_platform == 'linux'",
+    "nixl==1.5.0 ; sys_platform == 'linux'",
     "vllm-router ; sys_platform == 'linux'",
 ]
 quack = [
@@ -153,10 +150,11 @@ override-dependencies = [
     # Shared by vLLM, TileLang, and the pinned FlashAttention 4 runtime.
     "apache-tvm-ffi==0.1.11",
     "nvidia-cudnn-cu13>=9.15 ; sys_platform == 'linux'",
-    # vLLM 0.32.0 requires this exact CUTLASS DSL version and quack-kernels 0.6.5.
+    # vLLM 0.33.0 requires this exact CUTLASS DSL version and quack-kernels 0.6.5.
     # Keep this pin aligned with quack-kernels and the flash-attn-4 revision below.
     "nvidia-cutlass-dsl==4.8.0",
-    "torch>=2.13.0",
+    # vLLM 0.33.0 is compiled against and pins exactly torch 2.14.0.
+    "torch==2.14.0",
     "openenv-core",
     # TODO: PyPI release of prime-traces with PrimeIntellect-ai/prime#1109 (orjson records) asap.
     "prime-traces @ git+https://github.com/PrimeIntellect-ai/prime@adf8d25c3e34df8e039674fae9e0e4bf9cee27f3#subdirectory=packages/prime-traces",
@@ -181,8 +179,8 @@ members = [
 ]
 # Tau tasksets require litellm<1.82.7, while Harbor requires litellm>=1.92.0.
 # Apex pins a verifiers git revision instead of the editable workspace member.
-# ALE CLI requires numpy>=2.4.4, while vLLM 0.32.0 (via mistral-common) requires numpy<2.4.
-# IFBench requires setuptools<78, while vLLM 0.32.0 requires setuptools>=78.1.1.
+# ALE CLI requires numpy>=2.4.4, while vLLM 0.33.0 (via mistral-common) requires numpy<2.4.
+# IFBench requires setuptools<78, while vLLM 0.33.0 requires setuptools>=78.1.1.
 # A workspace shares one lock, so these packages need separate environments.
 exclude = [
     "deps/prime-envs/environments/if/ifbench",
@@ -206,8 +204,8 @@ requires-dist = [
 # caller requests its CUDA 13 extra. Use the requested CUDA 13 package only.
 [[tool.uv.dependency-metadata]]
 name = "nixl"
-version = "0.10.1"
-requires-dist = ["nixl-cu13==0.10.1"]
+version = "1.5.0"
+requires-dist = ["nixl-cu13==1.5.0"]
 
 # CUTLASS DSL has the same CUDA 12 default. vLLM requests its CUDA 13 extra,
 # so resolve that variant without installing both runtime packages.
@@ -228,9 +226,9 @@ harbor = "2026-10-08T00:00:00Z"
 # CA bundle: security fixes (e.g. 2026.10.19 for the eMudhra compromise) must not wait out the cooldown
 certifi = false
 tokenspeed-mla = false
-# vLLM 0.32.0 pins flashinfer-python 0.7.0.post1, published inside the rolling window
+# vLLM 0.33.0 pins flashinfer-python 0.7.2.post1, published inside the rolling window
 flashinfer-python = false
-# vLLM 0.32.0 pins these, published inside the rolling window
+# vLLM 0.33.0 pins these, published inside the rolling window
 instanttensor = false
 humming-kernels = false
 fastokens = false
@@ -256,7 +254,6 @@ deep-gemm = false
 nixl-cu13 = false
 torchao = false
 prime-kernels = false
-flash-mla = false
 deep-select = false
 
 [tool.uv.sources]
@@ -286,8 +283,8 @@ vllm-router = [
     { url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.2.3/vllm_router-0.2.3-cp38-abi3-manylinux_2_28_aarch64.whl", marker = "platform_machine == 'aarch64'" },
 ]
 vllm = [
-    { url = "https://github.com/vllm-project/vllm/releases/download/v0.32.0/vllm-0.32.0-cp38-abi3-manylinux_2_28_x86_64.whl", marker = "platform_machine == 'x86_64'" },
-    { url = "https://github.com/vllm-project/vllm/releases/download/v0.32.0/vllm-0.32.0-cp38-abi3-manylinux_2_28_aarch64.whl", marker = "platform_machine == 'aarch64'" },
+    { url = "https://github.com/vllm-project/vllm/releases/download/v0.33.0rc3/vllm-0.33.0-cp38-abi3-manylinux_2_28_x86_64.whl", marker = "platform_machine == 'x86_64'" },
+    { url = "https://github.com/vllm-project/vllm/releases/download/v0.33.0rc3/vllm-0.33.0-cp38-abi3-manylinux_2_28_aarch64.whl", marker = "platform_machine == 'aarch64'" },
 ]
 deep-ep = [
     { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl", marker = "platform_machine == 'x86_64'" },
@@ -305,22 +302,20 @@ deep-select = [
     { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121630-cp312-cp312-linux_x86_64.whl", marker = "platform_machine == 'x86_64'" },
     { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121543-cp312-cp312-linux_aarch64.whl", marker = "platform_machine == 'aarch64'" },
 ]
-# FlashMLA nv_dev @ b7643bd built against torch 2.13.0+cu130 (see the release notes for the recipe).
-flash-mla = { url = "https://github.com/NovaSky-AI/skyrl-wheels/releases/download/cu13torch2.13-flashmla-b7643bd/flash_mla-1.0.0%2Bcu.13.0.torch.2.13.b7643bd-cp312-cp312-linux_x86_64.whl", marker = "platform_machine == 'x86_64'" }
 flash-linear-attention = { git = "https://github.com/fla-org/flash-linear-attention" }
 mamba-ssm = { git = "https://github.com/state-spaces/mamba", rev = "f577286" }
 flash_attn_3 = { index = "pytorch-cu130-test" }
 flashinfer-cubin = { index = "flashinfer" }
 # Per-platform wheels: the aarch64 build also requires sm110a/sm121a wheels that x86_64 lacks.
 flashinfer-jit-cache = [
-    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache-0.7.0.post1%2Bcu130-cp39-abi3-manylinux_2_28_x86_64.whl", marker = "platform_machine == 'x86_64'" },
-    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache-0.7.0.post1%2Bcu130-cp39-abi3-manylinux_2_28_aarch64.whl", marker = "platform_machine == 'aarch64'" },
+    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache-0.7.2.post1%2Bcu130-cp39-abi3-manylinux_2_28_x86_64.whl", marker = "platform_machine == 'x86_64'" },
+    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache-0.7.2.post1%2Bcu130-cp39-abi3-manylinux_2_28_aarch64.whl", marker = "platform_machine == 'aarch64'" },
 ]
 
 [tool.uv.extra-build-dependencies]
 flash-attn = [{ requirement = "torch", match-runtime = true }]
 flash-attn-3 = [{ requirement = "torch", match-runtime = true }]
-mamba-ssm = ["torch>=2.13.0"]
+mamba-ssm = ["torch>=2.14.0"]
 
 [tool.uv.extra-build-variables]
 flash-attn = { FLASH_ATTENTION_SKIP_CUDA_BUILD = "TRUE" }
diff --git a/scripts/install_nixl_from_source.sh b/scripts/install_nixl_from_source.sh
index 387afa113aa..09db8e4b65a 100755
--- a/scripts/install_nixl_from_source.sh
+++ b/scripts/install_nixl_from_source.sh
@@ -28,7 +28,8 @@ mkdir -p "$WORKSPACE"
 UCX_SRC="$WORKSPACE/ucx_source"
 UCX_INSTALL="$PROJECT_DIR/third_party/ucx"
 NIXL_SRC="$WORKSPACE/nixl_source"
-NIXL_VERSION="${NIXL_VERSION:-0.10.1}"
+# Git tag; NIXL tags carry a `v` prefix from 1.0 on, the wheel version does not.
+NIXL_VERSION="${NIXL_VERSION:-v1.5.0}"
 CUDA_PATH="${CUDA_HOME:-/usr/local/cuda}"
 NPROC=$(nproc)
 
@@ -102,5 +103,5 @@ printf 'meson-python==0.21.1\n' > "$WORKSPACE/nixl-build-constraints.txt"
 "$PYTHON" -m pip wheel . --no-deps --wheel-dir="$WHEEL_DIR" \
     --build-constraint "$WORKSPACE/nixl-build-constraints.txt"
 
-WHEEL=$(ls "$WHEEL_DIR"/nixl_cu13-"$NIXL_VERSION"-*.whl | head -1)
+WHEEL=$(ls "$WHEEL_DIR"/nixl_cu13-"${NIXL_VERSION#v}"-*.whl | head -1)
 echo "=== NIXL wheel built at: $WHEEL ==="
diff --git a/src/prime_rl/inference/patches.py b/src/prime_rl/inference/patches.py
index 34b9067f02a..31157d5fed7 100644
--- a/src/prime_rl/inference/patches.py
+++ b/src/prime_rl/inference/patches.py
@@ -44,8 +44,6 @@ def apply_shared_vllm_patches():
     # Set by `server()` when the LoRA target modules include no expert layers.
     if os.environ.get("PRIME_NO_MOE_LORA") == "1":
         monkey_patch_no_moe_lora()
-    monkey_patch_deepseek_v4_c128_boundary()
-    monkey_patch_deepseek_v4_attn_sink_loading()
 
 
 @dataclass
@@ -146,115 +144,24 @@ def process_output(self, pending):
 
 
 def monkey_patch_aux_output_store():
-    """Keep routed-experts store eviction from rescanning referenced keys, and size the store above the KV cache.
+    """Size the routed-experts store above the KV cache.
 
-    Eviction skips keys referenced by running requests, and they pile up at the front of the LRU, so every ``put``
-    on a full store rescans them. Their position never affects which keys are evicted, so move them to the end.
-    At full KV the referenced keys alone can exceed a store of exactly KV size and ``put`` raises, so the worker
-    builds it at ``AUX_OUTPUT_STORE_HEADROOM`` times vLLM's size (anonymous mmap: unused slots cost no memory).
+    With a KV connector the scheduler retains rows for up to twice the KV cache's blocks
+    (vllm-project/vllm#62591), and a store of exactly KV size then grows by copying its whole arena on the
+    engine path. The worker builds it at ``AUX_OUTPUT_STORE_HEADROOM`` times vLLM's size instead (anonymous
+    mmap: unused slots cost no memory). vllm-project/vllm#62392 keeps referenced keys out of the eviction LRU,
+    which this patch used to work around.
     """
-    from itertools import takewhile
-
     from vllm.distributed.aux_output_connector import worker as aux_worker
     from vllm.distributed.aux_output_connector.store import BlockObjectStore
 
-    original_evict_to_fit = BlockObjectStore._evict_to_fit
-
-    def _evict_to_fit(self, protected: set[str]) -> None:
-        if len(self._lru) > self.num_slots:
-            for key in list(takewhile(self._references.__contains__, self._lru)):
-                self._lru.move_to_end(key)
-        original_evict_to_fit(self, protected)
-
     class _RoomyBlockObjectStore(BlockObjectStore):
         def __init__(self, *, max_bytes: int, object_nbytes: int):
             super().__init__(max_bytes=AUX_OUTPUT_STORE_HEADROOM * max_bytes, object_nbytes=object_nbytes)
 
-    BlockObjectStore._evict_to_fit = _evict_to_fit
     aux_worker.BlockObjectStore = _RoomyBlockObjectStore
 
 
-def monkey_patch_deepseek_v4_c128_boundary():
-    """Fix DeepSeek V4 long-context KV corruption under piecewise CUDA graphs.
-
-    On vLLM 0.30+, including the pinned 0.32.0, the piecewise CUDA graphs are captured on a step
-    that skips writing C128 compressed KV, so every replay skips it too. Mixed batches of up to
-    512 tokens never store their C128 entries, later queries attend over stale compressed KV, and
-    long-context rollouts drift from the trainer, inflating mismatch KL.
-
-    The skip is an optimization: ``DeepseekCompressor.forward`` returns early when
-    ``_get_c128_boundary`` reports that no request completes a 128-token block this step.
-    vllm-project/vllm#55353 made that helper return a bool under Model Runner V2, so capture on
-    dummy batches records the early return. Returning None instead disables the skip, which is
-    the path vLLM already takes under async speculative decoding and the one 0.29 always took.
-
-    The upstream fix, vllm-project/vllm#59709, keeps the skip on eager steps and only stops it
-    from being captured. That change sits inline in ``forward``, so matching it here would mean
-    copying the whole method. Patching the helper costs one extra compress kernel launch on eager
-    steps that complete no block, and the kernel already skips those tokens on the GPU.
-
-    Remove this patch once vllm-project/vllm#59709 lands in the pinned vLLM version.
-    """
-    from vllm.logger import init_logger
-    from vllm.models.deepseek_v4 import compressor
-
-    if getattr(compressor._get_c128_boundary, "_prime_rl_always_compresses", False):
-        return
-
-    def _get_c128_boundary(metadata):
-        return None
-
-    _get_c128_boundary._prime_rl_always_compresses = True
-    compressor._get_c128_boundary = _get_c128_boundary
-    init_logger("vllm.prime_rl.deepseek_v4").info("DeepSeek V4 C128 compressor store runs on every step.")
-
-
-def monkey_patch_deepseek_v4_attn_sink_loading():
-    """Fix DeepSeek V4 weight updates silently keeping the boot-time attention sinks.
-
-    ``DeepseekV4Model.load_weights`` writes the sinks with a bare slice copy instead of
-    ``param.weight_loader``. Layerwise reload moves parameters to meta and only buffers writes
-    that go through a loader, so the copy is silently dropped. This pads this rank's heads back
-    to the full parameter with ``-inf`` (no sink) and loads that through ``param.weight_loader``.
-
-    Remove this patch once vllm-project/vllm#57797 lands in the pinned vLLM version. Only
-    ``nvidia/model.py`` is covered, not the MTP, DSpark or non-NVIDIA variants.
-    """
-    from vllm.logger import init_logger
-    from vllm.model_executor.model_loader.weight_utils import default_weight_loader
-    from vllm.models.deepseek_v4.nvidia import model as dsv4_model
-
-    original_load_weights = dsv4_model.DeepseekV4Model.load_weights
-    if getattr(original_load_weights, "_prime_rl_uses_weight_loaders", False):
-        return
-
-    def load_weights(self, weights):
-        params = dict(self.named_parameters())
-        tp_size = dsv4_model.get_tensor_model_parallel_world_size()
-        heads_per_rank = self.config.num_attention_heads // tp_size
-        head_start = heads_per_rank * dsv4_model.get_tensor_model_parallel_rank()
-        loaded_params: set[str] = set()
-
-        def remaining_weights():
-            for name, weight in weights:
-                if "attn_sink" not in name or dsv4_model.is_pp_missing_parameter(name, self):
-                    yield name, weight
-                    continue
-                param = params[name]
-                sink = weight.new_full(tuple(param.shape), -float("inf"))
-                sink[:heads_per_rank] = weight[head_start : head_start + heads_per_rank]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, sink)
-                loaded_params.add(name)
-
-        loaded_params.update(original_load_weights(self, remaining_weights()))
-        return loaded_params
-
-    load_weights._prime_rl_uses_weight_loaders = True
-    dsv4_model.DeepseekV4Model.load_weights = load_weights
-    init_logger("vllm.prime_rl.deepseek_v4").info("DeepSeek V4 attn_sink loads through param.weight_loader.")
-
-
 def monkey_patch_online_mxfp8_moe_keep_kernel_on_reload():
     """Keep the boot-time online-MXFP8 MoE kernel across weight reloads.
 
@@ -450,12 +357,16 @@ def _patched_config():
 def monkey_patch_routed_experts_with_pd_connectors():
     """Let vLLM's routed-experts capture (AuxOutput connector) run on NIXL P/D instances.
 
-    vLLM (0.32 included) rejects routed-experts capture with P/D connectors, because a decode instance never
+    vLLM (0.33 included) rejects routed-experts capture with P/D connectors, because a decode instance never
     computes a remotely prefilled prompt, so it has no rows for it. This carries our upstream fix
     (AuxOutput patch 02): a request on a decode instance emits rows only from its first forward on
     this instance, and the worker does not buffer the
-    partially known first block. The router takes the prompt rows from the prefill instance. Delete
-    once vLLM supports routed experts with P/D connectors.
+    partially known first block. The router takes the prompt rows from the prefill instance.
+
+    Since vllm-project/vllm#62591 the scheduler also caps local and external prefix hits at the blocks whose
+    rows this instance retains (``max_cache_hit`` / ``cap_external_hit``). A decode instance never has rows for
+    a remotely prefilled prompt, so the cap would zero the NIXL load and recompute the whole prompt here; decode
+    requests skip both caps. Delete once vLLM supports routed experts with P/D connectors.
     """
     from vllm.config.vllm import VllmConfig
     from vllm.distributed.aux_output_connector.connector import AuxOutputSchedulerConnector
@@ -501,6 +412,19 @@ def request_finished(self, request):
             _emit_floors(self).pop(request.request_id, None)
         return original_request_finished(self, request)
 
+    original_max_cache_hit = AuxOutputSchedulerConnector.max_cache_hit
+
+    def max_cache_hit(self, request):
+        # None is the scheduler's "no cap", as without an AuxOutput connector.
+        return None if _is_pd_decode_request(request) else original_max_cache_hit(self, request)
+
+    original_cap_external_hit = AuxOutputSchedulerConnector.cap_external_hit
+
+    def cap_external_hit(self, request, num_local, num_external):
+        if _is_pd_decode_request(request):
+            return num_external
+        return original_cap_external_hit(self, request, num_local, num_external)
+
     original_worker_init = AuxOutputWorkerConnector.__init__
 
     def worker_init(self, *args, **kwargs):
@@ -512,6 +436,8 @@ def worker_init(self, *args, **kwargs):
     VllmConfig._verify_aux_output_compatibility = _verify_aux_output_compatibility
     AuxOutputSchedulerConnector.build_connector_meta = build_connector_meta
     AuxOutputSchedulerConnector.request_finished = request_finished
+    AuxOutputSchedulerConnector.max_cache_hit = max_cache_hit
+    AuxOutputSchedulerConnector.cap_external_hit = cap_external_hit
     AuxOutputWorkerConnector.__init__ = worker_init
 
 
@@ -677,8 +603,8 @@ def __init__(self, *args, **kwargs):
             _track_externally_loaded_blocks(self.kv_cache_manager, gid, block_size),
         )
 
-    def _get_local_prefix_cache_hit(self, request):
-        hit = original_get_local_prefix_cache_hit(self, request)
+    def _get_local_prefix_cache_hit(self, request, max_cache_hit_length=None):
+        hit = original_get_local_prefix_cache_hit(self, request, max_cache_hit_length)
         if self._prime_rl_loaded_blocks is None or request.num_output_tokens or _is_pd_decode_request(request):
             return hit
         gid, block_size, loaded = self._prime_rl_loaded_blocks
@@ -794,7 +720,7 @@ def _patch_qwen35_moe_lora_format():
 def _patch_lora_key_prefix():
     """Accept both bare-suffix and fully-qualified expert module names in LoRA adapters.
 
-    Copy of vLLM 0.32.0's ``LoRAModel.from_local_checkpoint`` with one change: the
+    Copy of vLLM 0.33.0's ``LoRAModel.from_local_checkpoint`` with one change: the
     ``.experts`` branch of ``check_unexpected_modules`` accepts either the bare suffix
     (``down_proj``) or the qualified per-expert name (``experts.N.down_proj``), where
     upstream only accepts the qualified form. Our trainer's 2D per-expert adapters
diff --git a/src/prime_rl/inference/vllm/gpt_oss_weight_loading.py b/src/prime_rl/inference/vllm/gpt_oss_weight_loading.py
index 247c516f4c9..04962da1939 100644
--- a/src/prime_rl/inference/vllm/gpt_oss_weight_loading.py
+++ b/src/prime_rl/inference/vllm/gpt_oss_weight_loading.py
@@ -4,16 +4,14 @@
 
 
 def patch_gpt_oss_weight_loading() -> None:
-    """Route GPT-OSS BF16 expert and sink weights through vLLM loaders.
+    """Route GPT-OSS BF16 expert weights through vLLM loaders.
 
     Remove this patch when the pinned vLLM version uses weight loaders for these
-    prepared tensors in ``GptOssModel._load_weights_other``.
+    prepared tensors in ``GptOssModel._load_weights_other``. Attention sinks already load
+    through their own loader (vllm-project/vllm#57797), which shards them per rank itself.
     """
     from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
-    from vllm.model_executor.model_loader.weight_utils import (
-        default_weight_loader,
-        maybe_remap_moe_expert_param_name,
-    )
+    from vllm.model_executor.model_loader.weight_utils import maybe_remap_moe_expert_param_name
     from vllm.model_executor.models import gpt_oss
 
     original_load_weights = gpt_oss.GptOssModel._load_weights_other
@@ -57,8 +55,6 @@ def load_weights(
         self,
         ep_rank_end: int,
         ep_rank_start: int,
-        heads_per_rank: int,
-        head_start: int,
         weights: Iterable[tuple[str, torch.Tensor]],
         stacked_params_mapping: list[tuple[str, ...]],
     ) -> set[str]:
@@ -79,9 +75,7 @@ def load_weights(
 
         def remaining_weights():
             for name, weight in weights:
-                if not (
-                    name.endswith((".w13_weight", ".w2_weight", ".w13_bias", ".w2_bias")) or name.endswith(".sinks")
-                ):
+                if not name.endswith((".w13_weight", ".w2_weight", ".w13_bias", ".w2_bias")):
                     yield name, weight
                     continue
 
@@ -107,17 +101,11 @@ def remaining_weights():
                         weight = weight[ep_rank_start:ep_rank_end]
                     else:
                         weight = weight[:, 2 * tp_rank_start : 2 * tp_rank_end]
-                elif name.endswith(".w2_bias"):
+                else:  # .w2_bias
                     if use_ep:
                         weight = weight[ep_rank_start:ep_rank_end]
                     elif tp_rank != 0:
                         weight = weight.zero_()
-                else:
-                    weight = weight.narrow(0, head_start, heads_per_rank)
-                    weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                    weight_loader(param, weight)
-                    loaded_params.add(name)
-                    continue
 
                 param.weight_loader(
                     param,
@@ -133,8 +121,6 @@ def remaining_weights():
                 self,
                 ep_rank_end,
                 ep_rank_start,
-                heads_per_rank,
-                head_start,
                 remaining_weights(),
                 stacked_params_mapping,
             )
diff --git a/src/prime_rl/inference/vllm/qwen38_weight_loading.py b/src/prime_rl/inference/vllm/qwen38_weight_loading.py
index 28ec3a2c224..a62648bf06d 100644
--- a/src/prime_rl/inference/vllm/qwen38_weight_loading.py
+++ b/src/prime_rl/inference/vllm/qwen38_weight_loading.py
@@ -12,7 +12,6 @@ def patch_qwen38_weight_loading() -> None:
     calling the embedding parameter's weight_loader.
     """
     from vllm.model_executor.model_loader.reload.layerwise import get_layerwise_info
-    from vllm.models.qwen4_exp.common.ple import compute_ple_shard_overlap
     from vllm.models.qwen4_exp.nvidia.ngram_embedding import Qwen4ExpNGramEmbedding
 
     original_load_weights = Qwen4ExpNGramEmbedding.load_weights
@@ -46,11 +45,8 @@ def local_weights():
                         expected_rows,
                         embedding.embedding_dim,
                     ):
-                        overlap = compute_ple_shard_overlap(
-                            checkpoint_start=checkpoint_start,
-                            checkpoint_rows=expected_rows,
-                            tp_start=embedding.shard_indices.org_vocab_start_index,
-                            tp_end=embedding.shard_indices.org_vocab_end_index,
+                        overlap = embedding._compute_shard_overlap(
+                            checkpoint_start=checkpoint_start, checkpoint_rows=expected_rows
                         )
                         if overlap is None:
                             skipped.add("ngram_embedding.weight")
diff --git a/src/prime_rl/inference/vllm/sampling_order.py b/src/prime_rl/inference/vllm/sampling_order.py
index bb3c4ac6d41..c7cf41271a2 100644
--- a/src/prime_rl/inference/vllm/sampling_order.py
+++ b/src/prime_rl/inference/vllm/sampling_order.py
@@ -25,8 +25,13 @@ def reference_top_p_then_top_k(logits: Tensor, k: Tensor | None, p: Tensor | Non
     return logits.scatter_(dim=-1, index=indices, src=sorted_logits)
 
 
-def apply_top_p_then_top_k(logits: Tensor, k: Tensor | None, p: Tensor | None) -> Tensor:
-    """Full-vocabulary nucleus mass, sorting only bounded candidates."""
+def apply_top_p_then_top_k(
+    logits: Tensor, k: Tensor | None, p: Tensor | None, has_topp_only_rows: bool = True
+) -> Tensor:
+    """Full-vocabulary nucleus mass, sorting only bounded candidates.
+
+    ``has_topp_only_rows`` matches vLLM's signature; it only selects a split kernel there.
+    """
     # For a single row, the full sort is cheaper than bounded-path dispatch.
     if k is None or logits.shape[0] == 1:
         return reference_top_p_then_top_k(logits, k, p)
diff --git a/src/prime_rl/inference/vllm/server.py b/src/prime_rl/inference/vllm/server.py
index 70db458de2f..00a8e0df6dd 100644
--- a/src/prime_rl/inference/vllm/server.py
+++ b/src/prime_rl/inference/vllm/server.py
@@ -7,8 +7,8 @@
 from fastapi.responses import JSONResponse
 from starlette.datastructures import State
 from vllm.engine.protocol import EngineClient
+from vllm.entrypoints.common.cli_args import make_arg_parser, validate_parsed_serve_args
 from vllm.entrypoints.launchers.api_server.app_state import init_app_state
-from vllm.entrypoints.launchers.cli_args import make_arg_parser, validate_parsed_serve_args
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 from vllm.entrypoints.serve.lora.protocol import LoadLoRAAdapterRequest
diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py
index 55a3ae7b386..0e20a181dcd 100644
--- a/src/prime_rl/inference/vllm/serving_tokens.py
+++ b/src/prime_rl/inference/vllm/serving_tokens.py
@@ -34,6 +34,7 @@
 import pybase64
 from fastapi import Request
 from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
 from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
     GenerateRequest,
     GenerateResponse,
@@ -84,16 +85,33 @@ def pack_sampled_logprobs(logprobs: FlatLogprobs) -> np.ndarray:
 
 class _PackedOutputs:
     """Wraps the result generator: takes the per-token payloads off each final
-    output (so upstream builds no per-token objects) and keeps them packed."""
+    output (so upstream builds no per-token objects) and keeps them packed.
+
+    ``request`` is the engine's request; ``serve_tokens_full_generator`` hands upstream a copy
+    whose logprobs are already off when they are packed here, since upstream reads
+    ``sampling_params.logprobs`` before it consumes the generator."""
 
     def __init__(self, generator: AsyncIterator[RequestOutput], request: GenerateRequest):
+        sampling_params = request.sampling_params
+        extra_args = sampling_params.extra_args or {}
         self._generator = generator
-        self._request = request
-        self._num_logprobs = request.sampling_params.logprobs
-        self._compact = (request.sampling_params.extra_args or {}).get("prl_compact_logprobs", False)
-        self._paired = (request.sampling_params.extra_args or {}).get("prl_paired_capture", False)
+        self._num_logprobs = sampling_params.logprobs
+        self._paired = extra_args.get("prl_paired_capture", False)
+        self._compact = not self._paired and extra_args.get("prl_compact_logprobs", False)
+        # ``serve_tokens`` asks for flat logprobs, so the engine returns FlatLogprobs here.
+        self._flat = (
+            not self._paired
+            and not self._compact
+            and sampling_params.flat_logprobs
+            and sampling_params.logprobs is not None
+            and sampling_params.logprobs <= 1
+        )
         self.fields: dict[int, dict[str, Any]] = {}
 
+    @property
+    def strips_logprobs(self) -> bool:
+        return self._compact or self._flat
+
     async def __aiter__(self):
         async for request_output in self._generator:
             for output in request_output.outputs:
@@ -120,12 +138,10 @@ async def __aiter__(self):
                         output.logprobs, self._num_logprobs, drop_filtered=output.sampling_mask is not None
                     )
                     output.logprobs = None
-                    self._request.sampling_params.logprobs = None
-                elif isinstance(output.logprobs, FlatLogprobs) and (self._num_logprobs or 0) <= 1:
+                elif self._flat:
+                    assert isinstance(output.logprobs, FlatLogprobs), "Did not output flat logprobs"
                     fields["completion_logprobs"] = encode_array(pack_sampled_logprobs(output.logprobs))
                     output.logprobs = None
-                    # Upstream would build ``logprobs.content`` from the logprobs we took.
-                    self._request.sampling_params.logprobs = None
                 mask = output.sampling_mask
                 if mask is not None:
                     fields["sampling_mask"] = {
@@ -171,7 +187,7 @@ async def serve_tokens(
         if paired:
             # Full support evidence supplies both sampled and candidate logprobs.
             sampling_params.logprobs = None
-        if not request.stream:
+        if not request.stream and request.output_mode == "tokens":
             # Upstream's prompt-logprobs response path expects the per-token dicts.
             sampling_params.flat_logprobs = sampling_params.prompt_logprobs is None
             if not sampling_params.stop:
@@ -185,16 +201,26 @@ async def serve_tokens_full_generator(  # type: ignore[override]
         request_id: str,
         model_name: str,
         request_metadata: RequestResponseMetadata,
+        derender_request: ChatCompletionRequest | None = None,
     ) -> ErrorResponse | GenerateResponse:
-        # The engine still needs the original sampling params while formatting strips payloads.
-        request = request.model_copy(update={"sampling_params": copy(request.sampling_params)})
+        if request.output_mode != "tokens":
+            return await super().serve_tokens_full_generator(
+                request, result_generator, request_id, model_name, request_metadata, derender_request=derender_request
+            )
         packed = _PackedOutputs(result_generator, request)
+        # The engine still needs the original sampling params; upstream formats from a copy without
+        # the logprobs we pack, or it would build ``logprobs.content`` from them.
+        sampling_params = copy(request.sampling_params)
+        if packed.strips_logprobs:
+            sampling_params.logprobs = None
+        request = request.model_copy(update={"sampling_params": sampling_params})
         response = await super().serve_tokens_full_generator(
             request,
             packed,
             request_id,
             model_name,
             request_metadata,
+            derender_request=derender_request,
         )
         if isinstance(response, GenerateTokensResponse):
             response = packed.post_process(response)
diff --git a/src/prime_rl/trainer/activation_checkpointing.py b/src/prime_rl/trainer/activation_checkpointing.py
index 436f07a3942..887ebc5d457 100644
--- a/src/prime_rl/trainer/activation_checkpointing.py
+++ b/src/prime_rl/trainer/activation_checkpointing.py
@@ -130,6 +130,9 @@ def get_activation_checkpoint_wrapper(config: ActivationCheckpointConfig) -> Cal
         checkpoint_wrapper,
         checkpoint_impl=CheckpointImpl.NO_REENTRANT,
         context_fn=partial(create_selective_checkpoint_contexts, policy),
+        # Keep SAC-saved tensors out of surrounding saved-tensor hooks (activation offloading), as before
+        # torch 2.14 made this explicit; leaving it unset warns whenever a hook is in scope.
+        respect_saved_tensors_hooks=False,
     )
 
 
diff --git a/src/prime_rl/trainer/models/kernels/__init__.py b/src/prime_rl/trainer/models/kernels/__init__.py
index 87a8b50656a..bb01c38b6e0 100644
--- a/src/prime_rl/trainer/models/kernels/__init__.py
+++ b/src/prime_rl/trainer/models/kernels/__init__.py
@@ -16,8 +16,13 @@ def cudnn_flashmla_error(heads: int, kv_group: int, dim: int) -> str | None:
     capability = torch.cuda.get_device_capability()
     if capability[0] != 9 and capability not in ((10, 0), (10, 3)):
         return f"the cudnn_flashmla backend runs on SM90, SM100 and SM103, but this GPU is SM{capability[0]}{capability[1]}"
-    if importlib.util.find_spec("flash_mla") is None:
-        return (
-            "`flash-mla` is not installed; install it with `uv sync --extra gpu` (the published wheel is x86_64 only)"
-        )
+    if importlib.util.find_spec("vllm") is None or importlib.util.find_spec("vllm._flashmla_C") is None:
+        return "vLLM's bundled FlashMLA (`vllm._flashmla_C`) is not installed; install it with `uv sync --extra gpu`"
     return None
+
+
+def flash_mla_sparse_fwd(*args, **kwargs):
+    """FlashMLA's sparse prefill forward, as bundled in the vLLM wheel (the `gpu` extra, x86_64 and aarch64)."""
+    from vllm.third_party.flashmla.flash_mla_interface import flash_mla_sparse_fwd
+
+    return flash_mla_sparse_fwd(*args, **kwargs)
diff --git a/src/prime_rl/trainer/models/kernels/deepseek_v4/dsv4_sparse_attn.py b/src/prime_rl/trainer/models/kernels/deepseek_v4/dsv4_sparse_attn.py
index 45aba995114..54c4c8cf59f 100644
--- a/src/prime_rl/trainer/models/kernels/deepseek_v4/dsv4_sparse_attn.py
+++ b/src/prime_rl/trainer/models/kernels/deepseek_v4/dsv4_sparse_attn.py
@@ -3,7 +3,7 @@
 # autograd rule that ties them together and forms the attention-sink gradient in torch. The
 # TileLang kernels themselves live in `dsv4_sparse_attn_fwd.py` and `dsv4_sparse_attn_bwd.py`; this
 # module is the only place that knows both exist. The `cudnn_flashmla` backend calls FlashMLA's
-# sparse prefill forward and the cuDNN frontend's DSA backward from their wheels.
+# sparse prefill forward (bundled in the vLLM wheel) and the cuDNN frontend's DSA backward.
 
 # TileLang ships a libcudart stub that proxies to the real CUDA runtime via
 # dlsym(RTLD_DEFAULT, ...).  If the stub's own symbols are the first ones found
@@ -172,7 +172,7 @@ def cudnn_flashmla_forward(
     _num_stages: int,
     _threads: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    from flash_mla import flash_mla_sparse_fwd
+    from prime_rl.trainer.models.kernels import flash_mla_sparse_fwd
 
     batch, seq_len, heads, dim = q.shape
     n_positions = kv.shape[1]
diff --git a/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py b/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py
index f0789fff48e..5945bc65b73 100644
--- a/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py
+++ b/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py
@@ -203,7 +203,7 @@ def flashmla_sparse_mla_forward(
     q: torch.Tensor, kv: torch.Tensor, indices: torch.Tensor, sm_scale: float, d_v: int
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """FlashMLA sparse prefill forward; returns a natural-log LSE."""
-    from flash_mla import flash_mla_sparse_fwd
+    from prime_rl.trainer.models.kernels import flash_mla_sparse_fwd
 
     B, S, H, D_qk = q.shape
     S_kv, kv_group = kv.shape[1:3]
diff --git a/tests/unit/inference/test_routed_experts.py b/tests/unit/inference/test_routed_experts.py
index 97b7bdc97e5..3f9d1b1760f 100644
--- a/tests/unit/inference/test_routed_experts.py
+++ b/tests/unit/inference/test_routed_experts.py
@@ -1,12 +1,11 @@
 import pickle
-import random
 
 import numpy as np
 from vllm.distributed.aux_output_connector import worker as aux_worker
 from vllm.distributed.aux_output_connector.connector import AuxRequestOutput
-from vllm.distributed.aux_output_connector.store import BlockObject, BlockObjectStore, BlockObjectStoreError
+from vllm.distributed.aux_output_connector.store import BlockObjectStore
 
-from prime_rl.inference.patches import monkey_patch_aux_output_store
+from prime_rl.inference.patches import AUX_OUTPUT_STORE_HEADROOM, monkey_patch_aux_output_store
 from prime_rl.inference.vllm.routed_experts import PackedAuxOutputs
 
 
@@ -26,33 +25,9 @@ def test_packed_aux_outputs_round_trip():
         np.testing.assert_array_equal(packed[request_id].rows, output.rows)
 
 
-def test_aux_output_store_eviction_matches_upstream(monkeypatch):
-    upstream_evict = BlockObjectStore._evict_to_fit
-    monkeypatch.setattr(BlockObjectStore, "_evict_to_fit", upstream_evict)
+def test_aux_output_store_has_headroom(monkeypatch):
     monkeypatch.setattr(aux_worker, "BlockObjectStore", BlockObjectStore)
     monkey_patch_aux_output_store()
-    upstream = type("Upstream", (BlockObjectStore,), {"_evict_to_fit": upstream_evict})
-    rng = random.Random(0)
-    for _ in range(50):
-        stores = [upstream(max_bytes=40, object_nbytes=1), BlockObjectStore(max_bytes=40, object_nbytes=1)]
-        references: dict[str, int] = {}
-        for _ in range(200):
-            objects = [BlockObject(f"k{rng.randrange(120)}", b"x") for _ in range(rng.randrange(6))]
-            retain = [f"k{rng.randrange(120)}" for _ in range(rng.randrange(4))]
-            release = [key for key, count in references.items() if count and rng.random() < 0.15]
-            for key in retain:
-                references[key] = references.get(key, 0) + 1
-            for key in release:
-                references[key] -= 1
-            failed = []
-            for store in stores:
-                try:
-                    store.put(objects, retain_keys=retain, release_keys=release)
-                    failed.append(False)
-                except BlockObjectStoreError:
-                    failed.append(True)
-            a, b = stores
-            assert failed[0] == failed[1]
-            assert set(a._lru) == set(b._lru)
-            # Only the order of unreferenced keys decides future evictions.
-            assert [k for k in a._lru if k not in a._references] == [k for k in b._lru if k not in b._references]
+    store = aux_worker.BlockObjectStore(max_bytes=40, object_nbytes=1)
+    assert isinstance(store, BlockObjectStore)
+    assert store.num_slots == AUX_OUTPUT_STORE_HEADROOM * 40
diff --git a/tests/unit/inference/test_routed_experts_kv_offload.py b/tests/unit/inference/test_routed_experts_kv_offload.py
index d69cef24626..5e93b085033 100644
--- a/tests/unit/inference/test_routed_experts_kv_offload.py
+++ b/tests/unit/inference/test_routed_experts_kv_offload.py
@@ -66,3 +66,36 @@ def request(rid, tokens, start):
     # Reallocating the blocks to a new owner clears them.
     assert manager.allocate_slots(request("c", list(range(1000, 1176)), start=0), 176) is not None
     assert not loaded
+
+
+def test_cut_patch_takes_upstream_max_cache_hit_length():
+    import torch
+    from vllm.sampling_params import SamplingParams
+    from vllm.utils.hashing import sha256
+    from vllm.v1.core.kv_cache_manager import KVCacheManager
+    from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
+    from vllm.v1.core.sched.scheduler import Scheduler
+    from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
+    from vllm.v1.request import Request
+
+    from prime_rl.inference.patches import monkey_patch_routed_experts_cut_loaded_prefix_hits
+
+    monkey_patch_routed_experts_cut_loaded_prefix_hits()
+    bs = 16
+    init_none_hash(sha256)
+    spec = FullAttentionSpec(block_size=bs, num_kv_heads=1, head_size=8, dtype=torch.float16)
+    config = KVCacheConfig(num_blocks=12, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["l0"], spec)])
+    manager = KVCacheManager(config, max_model_len=1024, scheduler_block_size=bs, hash_block_size=bs)
+    hasher = get_request_block_hasher(bs, sha256)
+    a = Request("a", list(range(70)), SamplingParams(max_tokens=1), None, block_hasher=hasher)
+    manager.allocate_slots(a, 70)
+    a.num_computed_tokens = a.num_tokens
+    manager.free(a)
+
+    # The scheduler passes the AuxOutput connector's cap positionally.
+    scheduler = SimpleNamespace(connector=None, kv_cache_manager=manager, _prime_rl_loaded_blocks=None)
+    b = Request("b", list(range(70)), SamplingParams(max_tokens=1), None, block_hasher=hasher)
+    _, num_local, _, _ = Scheduler._get_local_prefix_cache_hit(scheduler, b, 32)
+    assert num_local == 32
+    _, num_local, _, _ = Scheduler._get_local_prefix_cache_hit(scheduler, b)
+    assert num_local == 64
diff --git a/tests/unit/inference/test_routed_experts_pd.py b/tests/unit/inference/test_routed_experts_pd.py
index 0ea1a1c1be0..866e948d4a8 100644
--- a/tests/unit/inference/test_routed_experts_pd.py
+++ b/tests/unit/inference/test_routed_experts_pd.py
@@ -9,11 +9,21 @@
 )
 
 
+def _kv_cache_manager(block_size=4):
+    import torch
+    from vllm.v1.core.kv_cache_manager import KVCacheManager
+    from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
+
+    spec = FullAttentionSpec(block_size=block_size, num_kv_heads=1, head_size=8, dtype=torch.float16)
+    config = KVCacheConfig(num_blocks=12, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["l0"], spec)])
+    return KVCacheManager(config, max_model_len=1024, scheduler_block_size=block_size, hash_block_size=block_size)
+
+
 def test_remote_prefill_rows_start_at_first_forward():
     from vllm.distributed.aux_output_connector.connector import AuxOutputSchedulerConnector
 
     monkey_patch_routed_experts_with_pd_connectors()
-    connector = AuxOutputSchedulerConnector()
+    connector = AuxOutputSchedulerConnector(_kv_cache_manager())
     request = SimpleNamespace(
         request_id="r",
         block_hashes=[b"a" * 32],
@@ -40,6 +50,25 @@ def emit_start():
     assert not _emit_floors(connector)
 
 
+def test_pd_decode_requests_skip_retained_row_caps():
+    from vllm.distributed.aux_output_connector.connector import AuxOutputSchedulerConnector
+
+    monkey_patch_routed_experts_with_pd_connectors()
+    connector = AuxOutputSchedulerConnector(_kv_cache_manager(), kv_connector=object())
+
+    def request(kv_transfer_params):
+        return SimpleNamespace(request_id="r", block_hashes=[], num_tokens=10, kv_transfer_params=kv_transfer_params)
+
+    # A decode instance has no rows for a remotely prefilled prompt: the NIXL load must not be zeroed.
+    decode = request({"do_remote_prefill": True, "remote_engine_id": "p0"})
+    assert connector.cap_external_hit(decode, 0, 9) == 9
+    assert connector.max_cache_hit(decode) is None
+    # Other requests keep upstream's caps (no retained rows -> no hit).
+    local = request(None)
+    assert connector.cap_external_hit(local, 0, 8) == 0
+    assert connector.max_cache_hit(local) == 0
+
+
 def test_remote_prefill_partial_first_block_is_not_buffered():
     from vllm.distributed.aux_output_connector.routed_experts import RoutedExpertsBuffer
     from vllm.distributed.aux_output_connector.worker import _WorkerRequestState
diff --git a/tests/unit/inference/test_serving_tokens.py b/tests/unit/inference/test_serving_tokens.py
index 0d7d2b08733..080a4f49ec1 100644
--- a/tests/unit/inference/test_serving_tokens.py
+++ b/tests/unit/inference/test_serving_tokens.py
@@ -244,3 +244,41 @@ def test_pack_sampled_logprobs_takes_first_entry_and_clamps():
     logprobs.append({3: Logprob(float("-inf"), 9, None)})
 
     np.testing.assert_array_equal(pack_sampled_logprobs(logprobs), np.array([-0.5, -9999.0, -9999.0], np.float32))
+
+
+@pytest.mark.asyncio
+async def test_flat_sampled_logprobs_are_packed():
+    """The default RL path (flat logprobs, logprobs<=1): upstream reads ``logprobs`` before it
+    consumes the outputs, so the override must turn them off up front."""
+    serving = object.__new__(PrimeRlServingTokens)
+    serving.engine_client = SimpleNamespace(
+        vllm_config=SimpleNamespace(aux_output_config=SimpleNamespace(enable_return_routed_experts=False))
+    )
+    serving.enable_prompt_tokens_details = False
+    serving.enable_log_outputs = False
+    logprobs = FlatLogprobs()
+    for sampled, value in [(2, -0.25), (7, -20000.0)]:
+        logprobs.append_fast([sampled], [value], iter(range(1)), [None])
+    output = RequestOutput(
+        "test",
+        None,
+        [10, 11],
+        None,
+        [CompletionOutput(0, "", [2, 7], None, logprobs, finish_reason="length")],
+        True,
+    )
+
+    async def generate_outputs():
+        yield output
+
+    params = SamplingParams(logprobs=1, flat_logprobs=True)
+    request = GenerateRequest(model="test", token_ids=[10, 11], sampling_params=params, return_token_ids=True)
+    metadata = RequestResponseMetadata(request_id="test")
+    response = await serving.serve_tokens_full_generator(request, generate_outputs(), "test", "test", metadata)
+    assert isinstance(response, GenerateTokensResponse)
+    assert params.logprobs == 1
+    [choice] = response.model_dump()["choices"]
+    assert choice["logprobs"] is None
+    packed = choice["completion_logprobs"]
+    values = np.frombuffer(pybase64.b64decode(packed["data"]), dtype=packed["dtype"])
+    np.testing.assert_array_equal(values, np.array([-0.25, -9999.0], dtype=np.float32))
diff --git a/tests/unit/train/models/test_deepseek_v4_kernels.py b/tests/unit/train/models/test_deepseek_v4_kernels.py
index 6747b5c9c69..d57afb2718b 100644
--- a/tests/unit/train/models/test_deepseek_v4_kernels.py
+++ b/tests/unit/train/models/test_deepseek_v4_kernels.py
@@ -170,21 +170,16 @@ def _compare_accumulated_grads(module: nn.Module, expected: dict[str, torch.Tens
 SM_SCALE = DIM**-0.5
 
 requires_cudnn_flashmla = pytest.mark.skipif(
-    importlib.util.find_spec("flash_mla") is None
+    importlib.util.find_spec("vllm") is None
+    or importlib.util.find_spec("vllm._flashmla_C") is None
     or importlib.util.find_spec("cudnn") is None
     or not torch.cuda.is_available()
     or (torch.cuda.get_device_capability()[0] != 9 and torch.cuda.get_device_capability() not in ((10, 0), (10, 3))),
-    reason="the cudnn_flashmla backend needs `flash-mla` (the `gpu` extra), the cuDNN frontend and an SM90, SM100 or SM103 GPU",
+    reason="the cudnn_flashmla backend needs vLLM's FlashMLA (the `gpu` extra), the cuDNN frontend and an SM90, SM100 or SM103 GPU",
 )
 
 DSV4_BACKENDS = ["tilelang", pytest.param("cudnn_flashmla", marks=requires_cudnn_flashmla)]
 
-if importlib.util.find_spec("flash_mla") is not None:
-    # vLLM's bundled FlashMLA (`vllm._flashmla_C`) and the flash-mla wheel both define the GNU-unique
-    # symbol `Fwd_Sm90_Impl::features`, so the copy loaded first in the process serves both. vLLM's
-    # older list lacks TOPK_LENGTH, so load the wheel at collection, before any vLLM test loads its copy.
-    import flash_mla  # noqa: F401
-
 # What a real query with a short window or a saturated top-k looks like; a masked slot still
 # costs a GEMM column.
 MASKED_FRACTION = 0.25
diff --git a/uv.lock b/uv.lock
index 69b2a5b4a74..62221e7a921 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1,5 +1,5 @@
 version = 1
-revision = 3
+revision = 5
 requires-python = "==3.12.*"
 resolution-markers = [
     "platform_machine == 'x86_64' and sys_platform == 'linux'",
@@ -24,7 +24,6 @@ deep-select = false
 dion = false
 fastokens = false
 flash-attn-3 = false
-flash-mla = false
 flashinfer-python = false
 harbor = "2026-10-08T00:00:00Z"
 humming-kernels = false
@@ -166,8 +165,8 @@ overrides = [
     { name = "openenv-core" },
     { name = "prime-traces", git = "https://github.com/PrimeIntellect-ai/prime?subdirectory=packages%2Fprime-traces&rev=adf8d25c3e34df8e039674fae9e0e4bf9cee27f3" },
     { name = "tilelang", specifier = "==0.1.12" },
-    { name = "torch", marker = "sys_platform != 'linux'", specifier = ">=2.13.0" },
-    { name = "torch", marker = "sys_platform == 'linux'", specifier = ">=2.13.0", index = "https://download.pytorch.org/whl/cu130" },
+    { name = "torch", marker = "sys_platform != 'linux'", specifier = "==2.14.0" },
+    { name = "torch", marker = "sys_platform == 'linux'", specifier = "==2.14.0", index = "https://download.pytorch.org/whl/cu130" },
 ]
 
 [[manifest.dependency-metadata]]
@@ -177,8 +176,8 @@ requires-dist = ["pydantic>=2.10.6,<=2.13.5", "uvloop>=0.21.0"]
 
 [[manifest.dependency-metadata]]
 name = "nixl"
-version = "0.10.1"
-requires-dist = ["nixl-cu13==0.10.1"]
+version = "1.5.0"
+requires-dist = ["nixl-cu13==1.5.0"]
 
 [[manifest.dependency-metadata]]
 name = "nvidia-cutlass-dsl"
@@ -291,7 +290,7 @@ wheels = [
 
 [[package]]
 name = "aiohttp"
-version = "3.14.3"
+version = "3.14.5"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "aiohappyeyeballs" },
@@ -303,14 +302,15 @@ dependencies = [
     { name = "typing-extensions" },
     { name = "yarl" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/58/d9/22ce5786ac0c1653ae8b6c23bded02c1686d11f0dbb45b31ce128e0df985/aiohttp-3.14.3.tar.gz", hash = "sha256:9491196535a88924a60afd5b5f434b5b203b6cc616250878dbdb223a8f7844bc", size = 7971213, upload-time = "2026-07-23T01:57:27.037Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/6c/4c/bdccd81e9ee225b69c60e7766c9a5b05364f118f4d383713b89a682d772d/aiohttp-3.14.5.tar.gz", hash = "sha256:5558a7f5a05af9ecf744af91e5baefc436f93c9333e656c27ec253f9a6bbe178", size = 8078052, upload-time = "2026-10-11T01:05:12.408Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/18/d4/eb96299230e20acf2efae207cb8d69051f1f68e357e5ea5e479bf6fb097a/aiohttp-3.14.3-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:39aded8c7f3b935b54aab1d8d73c70ec0ee2d3ec3b943e0e86611bc150ba47f5", size = 754690, upload-time = "2026-07-23T01:53:47.332Z" },
-    { url = "https://files.pythonhosted.org/packages/30/07/4bbc222cc8dbe31d4c3e8a5baad2286e4d42026ac0c570027b89afce6344/aiohttp-3.14.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:617105e2c3018ee38d0c8ce5ee3c84f621a6d8b9f723202aacaff28449ca91ee", size = 511949, upload-time = "2026-07-23T01:53:55.083Z" },
-    { url = "https://files.pythonhosted.org/packages/54/b9/42e74c46b7b7c794b995bbc1f573fb48950c38b19d8600c62a6804ee2d67/aiohttp-3.14.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f631fe87a6f30df5fbe6d79640b25e4cffb38c31c7fb6f10871517b84b0f8c1a", size = 1765282, upload-time = "2026-07-23T01:53:59.662Z" },
-    { url = "https://files.pythonhosted.org/packages/52/b7/7cd31f29d6055bd711ae6e669367fba6f5ae9de463910a793e30556a8db7/aiohttp-3.14.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:543906c127fb1d929b95076db19b83fa2d46751006ff1e23b093aa5ac4d8db42", size = 1792122, upload-time = "2026-07-23T01:54:15.752Z" },
-    { url = "https://files.pythonhosted.org/packages/49/ed/3b203fa6de1b338c14acdc06bf6ca9b043b7944f005966958c2ced932cde/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:041badb8f84396357c4d3ad26de6afd7a32b112f43d3c63045c0c8278cfd2043", size = 1725210, upload-time = "2026-07-23T01:54:24.129Z" },
-    { url = "https://files.pythonhosted.org/packages/e0/eb/aad34e897e668424d6e995da5dff8a4a09af93363d3392488772957a63aa/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:d1558173930a5a8d3069cee5c92fc91c87c4dbcb099debbb3622053717145a19", size = 1768629, upload-time = "2026-07-23T01:54:40.103Z" },
+    { url = "https://files.pythonhosted.org/packages/d5/94/6ba86efddcb616c811b40e6a0dfdd862738f647e4e3860a961075d5e9ed8/aiohttp-3.14.5-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:df37b620684e19b5e25724412518ccafc3b1a49cdac706fdbd2f983fad943450", size = 817499, upload-time = "2026-10-11T01:00:18.901Z" },
+    { url = "https://files.pythonhosted.org/packages/64/91/11b89f45ca486252dd67dd5f3231fec04bf5518da39a95cb3997619f17fb/aiohttp-3.14.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d079c0a0135c36e7beb6f1c88087c8f108dc5891cdd0b5eafa778421bda70ed2", size = 548067, upload-time = "2026-10-11T01:00:22.659Z" },
+    { url = "https://files.pythonhosted.org/packages/22/ff/c6615806c14aab34f82b9424ccde8ce6e417315fd57ce1c5b4d4747888e1/aiohttp-3.14.5-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:abfda5cb094a829f7bc25216a32f7db2e85cc65bd59910f8e7b40b3d9b224764", size = 1896245, upload-time = "2026-10-11T01:00:24.638Z" },
+    { url = "https://files.pythonhosted.org/packages/da/e4/aa172eb534b7f02f1f8ff1c3213347eaf3cf218db91a727c6863c22f1035/aiohttp-3.14.5-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:20f085697d7e911f1f73c43ed03fafbed1e7121797e2eb5428efa80398060584", size = 1919630, upload-time = "2026-10-11T01:00:32.548Z" },
+    { url = "https://files.pythonhosted.org/packages/b2/94/eee018537ba19da0ceb2ac79cab83faed4ac49568e08376e2799043f2538/aiohttp-3.14.5-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:19e2ba471507c34f8252402ab50f5ab512398b9ea8c8f1cb26beb3f75793ba30", size = 1858112, upload-time = "2026-10-11T01:00:36.581Z" },
+    { url = "https://files.pythonhosted.org/packages/03/ad/6ddfe0aacd931c17b53533336d97e9d11a98b96d6ae815a9da0b19f82ccf/aiohttp-3.14.5-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:3be7dd397d64ca3e1869626fa9318aaebb54b7bf93bc72d7a205448d83e4f748", size = 1899501, upload-time = "2026-10-11T01:00:46.629Z" },
+    { url = "https://files.pythonhosted.org/packages/68/30/173960c42b05a6c59f7558e4b12a4b0d9ba376cf6aa9bde7f9e08a30ca8d/aiohttp-3.14.5-py3-none-any.whl", hash = "sha256:efc21a454892828368b11c2c780de0ff8bc991f73f6b99c6b66e56205470929b", size = 279517, upload-time = "2026-10-11T01:05:08.523Z" },
 ]
 
 [[package]]
@@ -1048,7 +1048,7 @@ source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "loguru" },
     { name = "pydantic" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
     { name = "transformers" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/2c/9e/d7f18bd9a0354088abc11a0c1f2c7698f7c49e5a709faedf6a46e388f693/compressed_tensors-0.17.0.tar.gz", hash = "sha256:15c20d06bdbcf35b51fc99fd125e7b9be1e1855567c33b7a46dfac26ad6fb126", size = 257091, upload-time = "2026-06-03T16:49:17.208Z" }
@@ -1458,7 +1458,7 @@ version = "0.1.0"
 source = { git = "https://github.com/samsja/dion.git?rev=0c801b0#0c801b0ae7c3a98433255bff78994e43e859c738" }
 dependencies = [
     { name = "numpy" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
     { name = "triton" },
 ]
 
@@ -1796,17 +1796,17 @@ wheels = [
 
 [[package]]
 name = "flash-attn"
-version = "2.8.3+cu130torch2.13"
-source = { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.47/flash_attn-2.8.3+cu130torch2.13-cp312-cp312-linux_x86_64.whl" }
+version = "2.8.3+cu130torch2.14"
+source = { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.10.0/flash_attn-2.8.3+cu130torch2.14-cp312-cp312-linux_x86_64.whl" }
 resolution-markers = [
     "platform_machine == 'x86_64' and sys_platform == 'linux'",
 ]
 dependencies = [
     { name = "einops" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
 ]
 wheels = [
-    { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.47/flash_attn-2.8.3+cu130torch2.13-cp312-cp312-linux_x86_64.whl", hash = "sha256:c65d762d24b7dd1ea8dbd9f5a98708fa8abbfd8f33f46bfc6123e76f1e174d3f" },
+    { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.10.0/flash_attn-2.8.3+cu130torch2.14-cp312-cp312-linux_x86_64.whl", hash = "sha256:623308fbb31944b6c312b0cc0c500d73d41e5c41e9d8933c27d290e0ca45caac" },
 ]
 
 [package.metadata]
@@ -1817,17 +1817,17 @@ requires-dist = [
 
 [[package]]
 name = "flash-attn"
-version = "2.8.3+cu130torch2.13"
-source = { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.48/flash_attn-2.8.3+cu130torch2.13-cp312-cp312-linux_aarch64.whl" }
+version = "2.8.3+cu130torch2.14"
+source = { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.10.1/flash_attn-2.8.3+cu130torch2.14-cp312-cp312-linux_aarch64.whl" }
 resolution-markers = [
     "platform_machine == 'aarch64' and sys_platform == 'linux'",
 ]
 dependencies = [
     { name = "einops" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
 ]
 wheels = [
-    { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.48/flash_attn-2.8.3+cu130torch2.13-cp312-cp312-linux_aarch64.whl", hash = "sha256:7688c302fd99f3cb3187cff8336500d17d455f49591ed687932d4f785e3bcb3b" },
+    { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.10.1/flash_attn-2.8.3+cu130torch2.14-cp312-cp312-linux_aarch64.whl", hash = "sha256:8270ec4b84b48377bc4663fdf9a1101e663dc49c7c4aafc82f0e54009f72cbc8" },
 ]
 
 [package.metadata]
@@ -1844,7 +1844,7 @@ dependencies = [
     { name = "einops" },
     { name = "ninja" },
     { name = "packaging" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
 ]
 wheels = [
     { url = "https://download.pytorch.org/whl/test/cu130/flash_attn_3-3.0.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:44a019ca428d0866fc73ca49b8af02eb4b6a37ae8a1a6b1a07917c383a187fe6", upload-time = "2026-02-14T04:28:03Z" },
@@ -1860,7 +1860,7 @@ dependencies = [
     { name = "einops" },
     { name = "nvidia-cutlass-dsl" },
     { name = "quack-kernels" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
     { name = "torch-c-dlpack-ext" },
     { name = "typing-extensions" },
 ]
@@ -1874,26 +1874,18 @@ dependencies = [
     { name = "transformers" },
 ]
 
-[[package]]
-name = "flash-mla"
-version = "1.0.0+cu.13.0.torch.2.13.b7643bd"
-source = { url = "https://github.com/NovaSky-AI/skyrl-wheels/releases/download/cu13torch2.13-flashmla-b7643bd/flash_mla-1.0.0%2Bcu.13.0.torch.2.13.b7643bd-cp312-cp312-linux_x86_64.whl" }
-wheels = [
-    { url = "https://github.com/NovaSky-AI/skyrl-wheels/releases/download/cu13torch2.13-flashmla-b7643bd/flash_mla-1.0.0%2Bcu.13.0.torch.2.13.b7643bd-cp312-cp312-linux_x86_64.whl", hash = "sha256:0bdb32340979c739659b305537dbe0445ba9f3148a4ce2f97f4fd170887c01e4" },
-]
-
 [[package]]
 name = "flashinfer-cubin"
-version = "0.7.0.post1"
+version = "0.7.2.post1"
 source = { registry = "https://flashinfer.ai/whl" }
 wheels = [
-    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_cubin-0.7.0.post1-py3-none-any.whl", hash = "sha256:f5b5fbf6e914b358219b845ed3ee586ffc1ae3c9768c8b135115ac14afbceed3" },
+    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_cubin-0.7.2.post1-py3-none-any.whl", hash = "sha256:1998c7006bcc06ee5dd3208a682af247382c6710d6b370a961381262c04c11f0" },
 ]
 
 [[package]]
 name = "flashinfer-jit-cache"
-version = "0.7.0.post1+cu130"
-source = { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache-0.7.0.post1%2Bcu130-cp39-abi3-manylinux_2_28_aarch64.whl" }
+version = "0.7.2.post1+cu130"
+source = { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache-0.7.2.post1%2Bcu130-cp39-abi3-manylinux_2_28_aarch64.whl" }
 resolution-markers = [
     "platform_machine == 'aarch64' and sys_platform == 'linux'",
 ]
@@ -1908,25 +1900,25 @@ dependencies = [
     { name = "flashinfer-jit-cache-sm90a" },
 ]
 wheels = [
-    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache-0.7.0.post1%2Bcu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:dcad7b6519e85b9332d1a18c84046797ce26206d4f15b9e484e2abc66f50d6fb" },
+    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache-0.7.2.post1%2Bcu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:8449156043036d240f3581c25e6883810f7ea6d656f56e78183b39eee9066877" },
 ]
 
 [package.metadata]
 requires-dist = [
-    { name = "flashinfer-jit-cache-sm100a", specifier = "==0.7.0.post1+cu130" },
-    { name = "flashinfer-jit-cache-sm103a", specifier = "==0.7.0.post1+cu130" },
-    { name = "flashinfer-jit-cache-sm110a", specifier = "==0.7.0.post1+cu130" },
-    { name = "flashinfer-jit-cache-sm120f", specifier = "==0.7.0.post1+cu130" },
-    { name = "flashinfer-jit-cache-sm121a", specifier = "==0.7.0.post1+cu130" },
-    { name = "flashinfer-jit-cache-sm80", specifier = "==0.7.0.post1+cu130" },
-    { name = "flashinfer-jit-cache-sm89", specifier = "==0.7.0.post1+cu130" },
-    { name = "flashinfer-jit-cache-sm90a", specifier = "==0.7.0.post1+cu130" },
+    { name = "flashinfer-jit-cache-sm100a", specifier = "==0.7.2.post1+cu130" },
+    { name = "flashinfer-jit-cache-sm103a", specifier = "==0.7.2.post1+cu130" },
+    { name = "flashinfer-jit-cache-sm110a", specifier = "==0.7.2.post1+cu130" },
+    { name = "flashinfer-jit-cache-sm120f", specifier = "==0.7.2.post1+cu130" },
+    { name = "flashinfer-jit-cache-sm121a", specifier = "==0.7.2.post1+cu130" },
+    { name = "flashinfer-jit-cache-sm80", specifier = "==0.7.2.post1+cu130" },
+    { name = "flashinfer-jit-cache-sm89", specifier = "==0.7.2.post1+cu130" },
+    { name = "flashinfer-jit-cache-sm90a", specifier = "==0.7.2.post1+cu130" },
 ]
 
 [[package]]
 name = "flashinfer-jit-cache"
-version = "0.7.0.post1+cu130"
-source = { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache-0.7.0.post1%2Bcu130-cp39-abi3-manylinux_2_28_x86_64.whl" }
+version = "0.7.2.post1+cu130"
+source = { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache-0.7.2.post1%2Bcu130-cp39-abi3-manylinux_2_28_x86_64.whl" }
 resolution-markers = [
     "platform_machine == 'x86_64' and sys_platform == 'linux'",
 ]
@@ -1939,92 +1931,92 @@ dependencies = [
     { name = "flashinfer-jit-cache-sm90a" },
 ]
 wheels = [
-    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache-0.7.0.post1%2Bcu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:9630eeb32396adb89ce8ffa8465fedd14549eb2e2ad3aff182d9e69be4a12075" },
+    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache-0.7.2.post1%2Bcu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:5a0a92ce56e11c19cec31948a4828b245314377eb587fe17307730b096225500" },
 ]
 
 [package.metadata]
 requires-dist = [
-    { name = "flashinfer-jit-cache-sm100a", specifier = "==0.7.0.post1+cu130" },
-    { name = "flashinfer-jit-cache-sm103a", specifier = "==0.7.0.post1+cu130" },
-    { name = "flashinfer-jit-cache-sm120f", specifier = "==0.7.0.post1+cu130" },
-    { name = "flashinfer-jit-cache-sm80", specifier = "==0.7.0.post1+cu130" },
-    { name = "flashinfer-jit-cache-sm89", specifier = "==0.7.0.post1+cu130" },
-    { name = "flashinfer-jit-cache-sm90a", specifier = "==0.7.0.post1+cu130" },
+    { name = "flashinfer-jit-cache-sm100a", specifier = "==0.7.2.post1+cu130" },
+    { name = "flashinfer-jit-cache-sm103a", specifier = "==0.7.2.post1+cu130" },
+    { name = "flashinfer-jit-cache-sm120f", specifier = "==0.7.2.post1+cu130" },
+    { name = "flashinfer-jit-cache-sm80", specifier = "==0.7.2.post1+cu130" },
+    { name = "flashinfer-jit-cache-sm89", specifier = "==0.7.2.post1+cu130" },
+    { name = "flashinfer-jit-cache-sm90a", specifier = "==0.7.2.post1+cu130" },
 ]
 
 [[package]]
 name = "flashinfer-jit-cache-sm100a"
-version = "0.7.0.post1+cu130"
+version = "0.7.2.post1+cu130"
 source = { registry = "https://flashinfer.ai/whl/cu130" }
 wheels = [
-    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache_sm100a-0.7.0.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:09d34f7a7ac5ef0ffeceefcb02c3926b9812d79d1922b3eda09a28902c98ef96" },
-    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache_sm100a-0.7.0.post1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:3654aad2bbaebf7ede45590e5fc241035fec92219e505ea2d2d8b77b0d88ee91" },
+    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm100a-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:ed6c128c4033de6fb6edbd1743a008ee70d9a3d4c8d1b60b45d821bd7e4143a2" },
+    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm100a-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:2d118fdb9a4a403ae22c08192830885eb27aba6d363ffa9c88758c7a51d1c126" },
 ]
 
 [[package]]
 name = "flashinfer-jit-cache-sm103a"
-version = "0.7.0.post1+cu130"
+version = "0.7.2.post1+cu130"
 source = { registry = "https://flashinfer.ai/whl/cu130" }
 wheels = [
-    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache_sm103a-0.7.0.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:f68545c055e52d674df5ba590a7e35414ccdd0fc6817e697fa0c302129f75f17" },
-    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache_sm103a-0.7.0.post1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:cd1e327d17be3e9683595d9461e1cf594a069923c940a5c91e7ed809dd7f8f1f" },
+    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm103a-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:8883759aaf289b8d494f714bb28766c9bef2e84fe6115d439d10e2417a8c7ae3" },
+    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm103a-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:c9c8430cd8c49d80edc44af820d4817c1c1b99cbe738c40034b41a44403ec5c7" },
 ]
 
 [[package]]
 name = "flashinfer-jit-cache-sm110a"
-version = "0.7.0.post1+cu130"
+version = "0.7.2.post1+cu130"
 source = { registry = "https://flashinfer.ai/whl/cu130" }
 wheels = [
-    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache_sm110a-0.7.0.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:a77273cb80cb81a98b827f495ca0204e8db83a3aabbbacdc576ecd4808fbd92f" },
+    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm110a-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:ff08fecf03bec168d37c8518c679e2332ba77d79824e806cc3cdd67c4f6b7a75" },
 ]
 
 [[package]]
 name = "flashinfer-jit-cache-sm120f"
-version = "0.7.0.post1+cu130"
+version = "0.7.2.post1+cu130"
 source = { registry = "https://flashinfer.ai/whl/cu130" }
 wheels = [
-    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache_sm120f-0.7.0.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:bfa5e2025a7a09755bdaa102b3ddc03809081a5c0bf06bb5e910407a43bb3acb" },
-    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache_sm120f-0.7.0.post1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:1c8b5a3614fdc9ad4f0e26502c5438657e28d667d165b05ff0383ca8d208568f" },
+    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm120f-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:1d3582708053a70b383c6d4df842a18c86f19518e230eaff6c1f689ce279dda5" },
+    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm120f-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:ab05a6e987bf79e18594d43ccc89ccdc8dbeacb3241c031ede8bc054a2baa9a7" },
 ]
 
 [[package]]
 name = "flashinfer-jit-cache-sm121a"
-version = "0.7.0.post1+cu130"
+version = "0.7.2.post1+cu130"
 source = { registry = "https://flashinfer.ai/whl/cu130" }
 wheels = [
-    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache_sm121a-0.7.0.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:1bf568b1c7452a697ad313a82aed8102b15204d9843ad2247d9778ca84513f75" },
+    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm121a-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:bcbd9b7d51cefcaaac9629cec260d37e253f6a8493876d045b6869ddf42b49c7" },
 ]
 
 [[package]]
 name = "flashinfer-jit-cache-sm80"
-version = "0.7.0.post1+cu130"
+version = "0.7.2.post1+cu130"
 source = { registry = "https://flashinfer.ai/whl/cu130" }
 wheels = [
-    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache_sm80-0.7.0.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:aebb89a7781caa17e42e5b43782f88cc679f8bc9a5e218d9d4a3fe5f91800a65" },
-    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache_sm80-0.7.0.post1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:b142aef0b32682a80ff2fd55aa74630678820776ffe03ecc183f4a01ada1c02a" },
+    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm80-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:f84fd89aac415b9cf79494e5606a307ecff2e421bcff132b62c6dba88f447a3b" },
+    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm80-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:b5076a39be7b63d15df44ecdfb9ce5be7206d39629b6d317b06d6e314f0b4cd6" },
 ]
 
 [[package]]
 name = "flashinfer-jit-cache-sm89"
-version = "0.7.0.post1+cu130"
+version = "0.7.2.post1+cu130"
 source = { registry = "https://flashinfer.ai/whl/cu130" }
 wheels = [
-    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache_sm89-0.7.0.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:284f4be6e2ce311b2666903d49fd21f9e6318f1c99bbfdaccf1bbdce28512970" },
-    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache_sm89-0.7.0.post1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:056a86546f8970ad4f345eecc4cccb3f36e9b4491109a26cc0fcc3703578a0ca" },
+    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm89-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:18fa31a010e3942a00ed4b6ffb4def344b5e10fc1ff9f56a538ce646bc316a7c" },
+    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm89-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:61ae80c3a78f20ebd8ef350bfe16bde947242bb9ac93f1cafb63e26c04c009dc" },
 ]
 
 [[package]]
 name = "flashinfer-jit-cache-sm90a"
-version = "0.7.0.post1+cu130"
+version = "0.7.2.post1+cu130"
 source = { registry = "https://flashinfer.ai/whl/cu130" }
 wheels = [
-    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache_sm90a-0.7.0.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:a5b1eb7505895d799896ba4f206bdbead53ca3300c6659dd5a91d3be547df918" },
-    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache_sm90a-0.7.0.post1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:9f6398cf38d3d772ce9244f5c7882ee572ec3dbcb495313839d6f908b1ea3ea6" },
+    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm90a-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:3bd375e978255d3215ee69da3434ca4b9dcce7a2d655c5d3445cd048fcc62096" },
+    { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm90a-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:52e5cc064b52290dad2f8889ee0eb5cd59621dfac8351841e4b3987d42d034a3" },
 ]
 
 [[package]]
 name = "flashinfer-python"
-version = "0.7.0.post1"
+version = "0.7.2.post1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "apache-tvm-ffi" },
@@ -2042,12 +2034,12 @@ dependencies = [
     { name = "packaging" },
     { name = "requests" },
     { name = "tabulate" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
     { name = "tqdm" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/14/b0/2f27f2645056b0e671b3bcf8a934e73ed89ca5c5df9ebe651f21c0d218e2/flashinfer_python-0.7.0.post1.tar.gz", hash = "sha256:2e67b38a9af1d4eedff96f7f5d79c6a4d9f0f97cb4a30f71e4ffd717f1ad9af4", size = 26061698, upload-time = "2026-09-29T23:34:14.797Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/05/76/1a42eafe2a3cd26924e485fb85bc635a2a41fcb14f453a8c037214b32940/flashinfer_python-0.7.2.post1.tar.gz", hash = "sha256:af60319898dcfd552ba12a34165b9dd46c0bbd83d48ce665b1df6bd6ae77e9f1", size = 51431554, upload-time = "2026-10-23T18:49:35.84Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/00/7d/16237087a2af99086ef4ae9808dff9b24f0a4db4cfd673233836821a32e0/flashinfer_python-0.7.0.post1-py3-none-any.whl", hash = "sha256:c7adf826568d61fc1b7d3aadd4cae387a35a138bfc08e2deca2f34ecfa280716", size = 33659854, upload-time = "2026-09-29T23:34:11.137Z" },
+    { url = "https://files.pythonhosted.org/packages/b6/ad/da1fc7644b28ab82db922021955a0c85c8d6b3b1d6db224a13ef65cef85e/flashinfer_python-0.7.2.post1-py3-none-any.whl", hash = "sha256:bfdc18b9e1e4d8d0730fe4343dc0e284cd0c160086de80f99888bc4867922d49", size = 62052965, upload-time = "2026-10-23T18:49:31.516Z" },
 ]
 
 [[package]]
@@ -2668,7 +2660,7 @@ wheels = [
 
 [[package]]
 name = "humming-kernels"
-version = "0.1.17"
+version = "0.1.19"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "cuda-bindings" },
@@ -2676,14 +2668,14 @@ dependencies = [
     { name = "numpy" },
     { name = "safetensors" },
     { name = "tabulate" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
     { name = "tqdm" },
     { name = "triton" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/60/52/b3958f17315854ca781909f39392e6b40d45965820912cd610191597bd92/humming_kernels-0.1.17.tar.gz", hash = "sha256:a2565bf03b8b2ca07086db00a949c0c92727ea675e4d3037f9337192a840225a", size = 403057, upload-time = "2026-10-14T17:07:29.566Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/6e/c4/4ca7fd5042fed448b62874f1d2886c93c0ac78da11cd5548f926e0483f64/humming_kernels-0.1.19.tar.gz", hash = "sha256:3c6c7ef0e55ef6b97b8b0d4134f7f34cf988df22f7885bff27aa0b30fbbd8e65", size = 428158, upload-time = "2026-10-30T19:30:57.267Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/38/78/ee230d0165b47f08103fdb9a58a64b28c91e15ed7e2d3beb3920543a9f42/humming_kernels-0.1.17-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:83dda343f36a52efb617ab7b85747baabfa1c2d00dcf1078c45f92b8df7a01c6", size = 496273, upload-time = "2026-10-14T17:07:26.658Z" },
-    { url = "https://files.pythonhosted.org/packages/da/fa/c4e367ca51b2611faf88ac9ee1abe5e9ee6d4d174bbabdb54481c65f7e2e/humming_kernels-0.1.17-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:e418b5482dabd71f64622893bdbb8bf27eb30379c443b92d32af48825d9ad64c", size = 504747, upload-time = "2026-10-14T17:07:28.196Z" },
+    { url = "https://files.pythonhosted.org/packages/69/34/29856875ef58ca87eaa31acf3d37466e8523025c709ef52478492060dc73/humming_kernels-0.1.19-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:13038de14ed2910bcb46f0147d9935c6cac835045e07b3f38c2137d76b77bc4c", size = 547057, upload-time = "2026-10-30T19:30:51.069Z" },
+    { url = "https://files.pythonhosted.org/packages/af/a1/d6c92ed0290b8e1322d7b9ffb3fafe4f2af51ba9b564fa561c036a402266/humming_kernels-0.1.19-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:16ad0c8e9c2d5a8db739bdb97bcc510cedc9dc4c4c111098cbb2d7ac9c1dd973", size = 556040, upload-time = "2026-10-30T19:30:52.832Z" },
 ]
 
 [package.optional-dependencies]
@@ -2877,7 +2869,7 @@ name = "instanttensor"
 version = "0.2.2"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/ca/a1/95d0d4cee908ab0bd8e56b44cdb19e678081c54de52d280b0b3c9f7e4c43/instanttensor-0.2.2.tar.gz", hash = "sha256:aca07c1afe4fde23dcbfb53044e2bef50dad4c0af054317acc669a1c8dc2b325", size = 613106, upload-time = "2026-10-13T15:06:59.034Z" }
 wheels = [
@@ -3454,7 +3446,7 @@ dependencies = [
     { name = "quack-kernels" },
     { name = "setuptools", version = "80.10.2", source = { registry = "https://pypi.org/simple" } },
     { name = "tilelang" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
     { name = "transformers" },
     { name = "triton" },
 ]
@@ -3849,7 +3841,7 @@ dependencies = [
     { name = "runai-model-streamer", extra = ["azure", "gcs", "s3"] },
     { name = "s3transfer" },
     { name = "safetensors" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
     { name = "zstandard" },
 ]
 wheels = [
@@ -3992,17 +3984,17 @@ provides-extras = ["dev"]
 
 [[package]]
 name = "multidict"
-version = "6.7.1"
+version = "7.2.1"
 source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/1a/c2/c2d94cbe6ac1753f3fc980da97b3d930efe1da3af3c9f5125354436c073d/multidict-6.7.1.tar.gz", hash = "sha256:ec6652a1bee61c53a3e5776b6049172c53b6aaba34f18c9ad04f82712bac623d", size = 102010, upload-time = "2026-01-26T02:46:45.979Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/7b/6b/29ef8e752f5854dcc267ded0652c475bec99b86f255957a0d4fcad657b18/multidict-7.2.1.tar.gz", hash = "sha256:e8af3e6f1bdbc31fd762b094133b478b89988eaebc63cbe82b1d9e3bbd15cede", size = 369994, upload-time = "2026-10-24T20:41:35.457Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/8d/9c/f20e0e2cf80e4b2e4b1c365bf5fe104ee633c751a724246262db8f1a0b13/multidict-6.7.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:a90f75c956e32891a4eda3639ce6dd86e87105271f43d43442a3aedf3cddf172", size = 76893, upload-time = "2026-01-26T02:43:52.754Z" },
-    { url = "https://files.pythonhosted.org/packages/a9/65/1caac9d4cd32e8433908683446eebc953e82d22b03d10d41a5f0fefe991b/multidict-6.7.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b0fa96985700739c4c7853a43c0b3e169360d6855780021bfc6d0f1ce7c123e7", size = 43872, upload-time = "2026-01-26T02:43:55.041Z" },
-    { url = "https://files.pythonhosted.org/packages/fd/80/c959c5933adedb9ac15152e4067c702a808ea183a8b64cf8f31af8ad3155/multidict-6.7.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:eb0ce7b2a32d09892b3dd6cc44877a0d02a33241fafca5f25c8b6b62374f8b75", size = 258883, upload-time = "2026-01-26T02:43:57.499Z" },
-    { url = "https://files.pythonhosted.org/packages/f3/8d/5e5be3ced1d12966fefb5c4ea3b2a5b480afcea36406559442c6e31d4a48/multidict-6.7.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bfde23ef6ed9db7eaee6c37dcec08524cb43903c60b285b172b6c094711b3961", size = 256322, upload-time = "2026-01-26T02:44:03.56Z" },
-    { url = "https://files.pythonhosted.org/packages/31/6e/d8a26d81ac166a5592782d208dd90dfdc0a7a218adaa52b45a672b46c122/multidict-6.7.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3758692429e4e32f1ba0df23219cd0b4fc0a52f476726fff9337d1a57676a582", size = 253955, upload-time = "2026-01-26T02:44:04.845Z" },
-    { url = "https://files.pythonhosted.org/packages/2b/17/ebb9644da78c4ab36403739e0e6e0e30ebb135b9caf3440825001a0bddcb/multidict-6.7.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fc5907494fccf3e7d3f94f95c91d6336b092b5fc83811720fae5e2765890dfba", size = 251377, upload-time = "2026-01-26T02:44:12.042Z" },
-    { url = "https://files.pythonhosted.org/packages/81/08/7036c080d7117f28a4af526d794aab6a84463126db031b007717c1a6676e/multidict-6.7.1-py3-none-any.whl", hash = "sha256:55d97cc6dae627efa6a6e548885712d4864b81110ac76fa4e534c03819fa4a56", size = 12319, upload-time = "2026-01-26T02:46:44.004Z" },
+    { url = "https://files.pythonhosted.org/packages/27/21/53c299f1ea289f5fd438e70d4f1f053f59b4485f5a410bcc85eecb553394/multidict-7.2.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:cb43bb34346155d479acdc5423afccdff79ce1fc66adfcc9de995cbec8a13acf", size = 172948, upload-time = "2026-10-24T20:40:10.189Z" },
+    { url = "https://files.pythonhosted.org/packages/bf/e4/bd8838244fd1309165c79e74af9a6b9e9e42fd0f2c24c75bee665fc9ba69/multidict-7.2.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d37773029ba034a2cc970b6eeb67dad3bff9d894393c51f6bb1c6b4f563a032c", size = 103911, upload-time = "2026-10-24T20:40:13.517Z" },
+    { url = "https://files.pythonhosted.org/packages/72/13/c081a1af93258029471636ea1c741141ece10f1045578a5f73ba99b694a0/multidict-7.2.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fa0ecd064b0ae18d235283abb5adbc4e07304dba8f2aa7c0d61f28bebcc70d4c", size = 506972, upload-time = "2026-10-24T20:40:11.68Z" },
+    { url = "https://files.pythonhosted.org/packages/43/50/6bd8600b4843f0a9d939e0cc0f4160f4c727bd6c69e4f7e294a1caa93db8/multidict-7.2.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1fb8c1a3480c130c3a6422d5f16f3cad6552bf08c2089dd288ba7134649183b5", size = 516796, upload-time = "2026-10-24T20:40:13.481Z" },
+    { url = "https://files.pythonhosted.org/packages/e8/bf/30a0c0911e4b61accb9995165acfdc5107dd335ad3ce003007afaf168dec/multidict-7.2.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:b4821bd30b38144d8ab75df56af0b07e9d71327c9cf94c0c7a1965f865bb9615", size = 492916, upload-time = "2026-10-24T20:40:16.619Z" },
+    { url = "https://files.pythonhosted.org/packages/7b/9f/ae4d6c442e11aea5ff631f36e4e8d5ff4d1f4c330d2c2721cb729fa5eed6/multidict-7.2.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:d615377a8e7bfc4f1983ede3f411c2ac2b0f7c3053f262d0437b7d8bcc4485e3", size = 508675, upload-time = "2026-10-24T20:40:17.249Z" },
+    { url = "https://files.pythonhosted.org/packages/9d/96/1aaad0e5df33f1f5af21775c411349c0234a6edd581ea403bc81919a6eab/multidict-7.2.1-py3-none-any.whl", hash = "sha256:08ea94663d9e7c25d14d78222f7aee906c472d4837766bec62477c38f6e419dc", size = 28145, upload-time = "2026-10-24T20:40:59.159Z" },
 ]
 
 [[package]]
@@ -4119,26 +4111,26 @@ wheels = [
 
 [[package]]
 name = "nixl"
-version = "0.10.1"
+version = "1.5.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "nixl-cu13" },
 ]
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/bb/a4/10e80e623790d3fb070e966b36bced009419d483c92ae0df6645230f606f/nixl-0.10.1-py3-none-any.whl", hash = "sha256:616465673dae5180d296525a03237af4cd5f2c00c3228d185bc06dbe621509b7", size = 6680, upload-time = "2026-03-03T19:55:44.848Z" },
+    { url = "https://files.pythonhosted.org/packages/90/6d/ab1e2ef610d6242b864c4fd8a4adb8caaf8386d2bdd4b8740edbba3e3abd/nixl-1.5.0-py3-none-any.whl", hash = "sha256:1be21929dcf331b7637e4daafd538169c70467b0dfb9445621b4036e4831fb38", size = 11804, upload-time = "2026-09-30T07:29:24.029Z" },
 ]
 
 [[package]]
 name = "nixl-cu13"
-version = "0.10.1"
+version = "1.5.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "numpy" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
 ]
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/47/72/295cc323773ed95b755fb0d7ce5c0181c8941b3be69c514182b3b1ab89ca/nixl_cu13-0.10.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:909d00ffc1929ef45cd3cfa0cd3585999274c90a6bcf0799000677cf83a8f0e2", size = 49977656, upload-time = "2026-03-03T19:55:11.638Z" },
-    { url = "https://files.pythonhosted.org/packages/58/0c/b14d9c260c423551bcfaf1fdf3255a3eff471372fd990c2370e8ac21ab10/nixl_cu13-0.10.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:322e4702606ad498a493d99a065af46c0c16ce84a4bed6495f85efab75670f0a", size = 51283105, upload-time = "2026-03-03T19:55:12.655Z" },
+    { url = "https://files.pythonhosted.org/packages/1a/0a/8613ed87db02018596ac01f3f45a06a4db25eb571f9059a7ad5b04e5fb51/nixl_cu13-1.5.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:0a3f11e247122c6a3425d1b0f9da282bc1306329535a1d665a0f820b67dbc962", size = 57035701, upload-time = "2026-09-30T07:27:03.574Z" },
+    { url = "https://files.pythonhosted.org/packages/1d/6f/0d2920d5e138ce94ba9106f97e82555d43f4c89cef9ca77fec2e9ef08cb7/nixl_cu13-1.5.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:5d8d2d56439bdc91ed415098c05a74e4f577621104f940269bdcbf5411690d2b", size = 58770022, upload-time = "2026-09-30T07:25:03.262Z" },
 ]
 
 [[package]]
@@ -4306,15 +4298,15 @@ wheels = [
 
 [[package]]
 name = "nvidia-cudnn-frontend"
-version = "1.30.0"
+version = "1.31.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "apache-tvm-ffi" },
     { name = "nvidia-cutlass-dsl" },
 ]
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/e4/53/07836d9f74f1416a8768530d428c155f67d7728f283bce3a9b3096b7e18f/nvidia_cudnn_frontend-1.30.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:279d227a226429288ae9e4e88db744ce144ae69605f63a286dc84ae180869a6e", size = 8631629, upload-time = "2026-09-23T03:44:15.265Z" },
-    { url = "https://files.pythonhosted.org/packages/ee/95/c92076af7561e6e43da1ee3ac4155202dcd3876e5b999ff0dea5754eb3d2/nvidia_cudnn_frontend-1.30.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5dda72db38ea0c39f158ebffa9596eefab21ce7cc6a12824490ab8b6cb60e76d", size = 8797507, upload-time = "2026-09-23T03:44:35.051Z" },
+    { url = "https://files.pythonhosted.org/packages/bf/4b/ba084b26c9709b05446642413aa4bbb1056fe00af7f99e356a3cf24774d6/nvidia_cudnn_frontend-1.31.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:47d1a5432f59a2a7d0f903727c0d244e1413a20686b98d8e3890919c5d82a508", size = 9970357, upload-time = "2026-10-08T18:05:50.547Z" },
+    { url = "https://files.pythonhosted.org/packages/cb/24/8ec9de1fd9ea95e1c1cb9e6e2d3a5547aa9c5ece1136d650ffce1bb5d8c5/nvidia_cudnn_frontend-1.31.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3171f102254136418f24b4470dfa5419f34b03d9d792a4b773cf65c7a37e6bf8", size = 10159344, upload-time = "2026-10-08T18:06:15.766Z" },
 ]
 
 [[package]]
@@ -4448,11 +4440,11 @@ wheels = [
 
 [[package]]
 name = "nvidia-nccl-cu13"
-version = "2.29.7"
+version = "2.30.7"
 source = { registry = "https://pypi.org/simple" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/72/0d/daf50d44177ee0cbc7ff0a0c91eb5ff676c82be42f9a970bc7597f440c3a/nvidia_nccl_cu13-2.29.7-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:674a12383e3c38a1bcccae7d4f3633b37852230b6047883cb2f4c2d1b36d9bf5", size = 206014712, upload-time = "2026-03-03T05:34:20.843Z" },
-    { url = "https://files.pythonhosted.org/packages/67/f4/58e4e91b6919367c7aafb8e36fce9aad1a3047e536bf7e2fd560927d3a4c/nvidia_nccl_cu13-2.29.7-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d", size = 205976000, upload-time = "2026-03-03T05:36:24.472Z" },
+    { url = "https://files.pythonhosted.org/packages/d1/21/a73174c6157101bdf1ffc22b517f76ff0082613989dd9bc8f43e8034caac/nvidia_nccl_cu13-2.30.7-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:ca786ffa5a647c75d4d1f5cc72a6c4f537947e2ba8823d7c8aaf768e7a7b9f77", size = 215983881, upload-time = "2026-06-09T03:23:15.633Z" },
+    { url = "https://files.pythonhosted.org/packages/3f/34/c500f90c7ae641b8e0f98965b36b8a7ac79cc8b296e8d251fe3eb592ee54/nvidia_nccl_cu13-2.30.7-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:cefa7fdb9710efd0f39c5f1be1d61ff6fc9a996c451265bd7fbdcf9455ed4b50", size = 215965170, upload-time = "2026-06-09T03:23:39.73Z" },
 ]
 
 [[package]]
@@ -5007,15 +4999,18 @@ requires-dist = [{ name = "verifiers", extras = ["harbor"], specifier = ">=0.4.1
 
 [[package]]
 name = "pillow"
-version = "12.3.0"
+version = "13.0.0"
 source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/1c/3d/bb7fca845737cf9d7dbde16ed1843984665ff2e0a518f5db43e77ec540b9/pillow-12.3.0.tar.gz", hash = "sha256:3b8182a766685eaa002637e28b4ec8d6b18819a0c71f579bf0dbaa5830297cce", size = 47025035, upload-time = "2026-07-01T11:56:38.965Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/ce/a4/9cc30690a394a8305e8be4c9036aa16d0dfddadd09079860b582f562496d/pillow-13.0.0.tar.gz", hash = "sha256:cdcd1ba258178f7356d74170c9a81ba4eab6a5d55748bc402a5a2a6fe6be3510", size = 47058631, upload-time = "2026-10-16T22:18:09.718Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/d8/66/9a386a92561f402389a4fc70c18838bf6d35eb5eb5c6850b4b2dc64f5048/pillow-12.3.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ffd0c5368496f41b0944be820fcb7a838aa6e623d250b01acf2643939c3f99d7", size = 4780323, upload-time = "2026-07-01T11:54:09.351Z" },
-    { url = "https://files.pythonhosted.org/packages/25/27/ac8f99618ffd3dde21db0f4d4b1d2ab00c0880595bfd17df103f7f39fd0c/pillow-12.3.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d9c7f76c0673154f044e9d78c8655fb4213f6ca31a836df48b40fe5d187717b9", size = 6266838, upload-time = "2026-07-01T11:54:11.71Z" },
-    { url = "https://files.pythonhosted.org/packages/84/21/a35af28dcc61f37ed850a2d64c65c701321dfbf25085e469d5559360cbbf/pillow-12.3.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:78cb2c6865a35ab8ff8b75fd122f6033b92a62c82801110e48ddd6c936a45d91", size = 6940830, upload-time = "2026-07-01T11:54:13.732Z" },
-    { url = "https://files.pythonhosted.org/packages/eb/51/8b08617af3ad95e33ce6d7dd2c99ed6c8298f7fb131636303956be022e25/pillow-12.3.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e491916b378fba47242221bb9ead245211b70d504f495d105d17b14a24b4907c", size = 6344383, upload-time = "2026-07-01T11:54:15.756Z" },
-    { url = "https://files.pythonhosted.org/packages/1d/72/cf78ac9780bb93c28328f408973845a309d4d145041665f734572ced1b52/pillow-12.3.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:0dd2064cbc55aaec028ef5fbb60fa47bb6c3e7918e07ff17935284b227a9d2df", size = 7052934, upload-time = "2026-07-01T11:54:17.721Z" },
+    { url = "https://files.pythonhosted.org/packages/39/4b/f2ed53d83aa44ac19599ac0a54897f867fafb95f406980d1e24083f8cbe3/pillow-13.0.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:fc481f742ccda20e9ed760b8f978dff50f2a4e9bc253ac167fb0a05b0783ef52", size = 4852329, upload-time = "2026-10-16T22:15:46.284Z" },
+    { url = "https://files.pythonhosted.org/packages/3b/bb/9e30e8ed18c20800ee327a12107d5249ad8dc444867b3adfcc21654e3c71/pillow-13.0.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8bf295fe459343bc6093c8e5e0df60bfa189e0fcb775a22e2c2488fdeccc6c0d", size = 6519309, upload-time = "2026-10-16T22:15:48.14Z" },
+    { url = "https://files.pythonhosted.org/packages/e6/99/fdf65cc539e266ab93546c261bdf98411fe165c6e4e13f9dbf597097192b/pillow-13.0.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:42d63be3af4f9b951b199b0a9d23365d37ab14687e6d77b74940432c8515c4ca", size = 7144542, upload-time = "2026-10-16T22:15:50.093Z" },
+    { url = "https://files.pythonhosted.org/packages/26/9b/f38b7c4d8f599537e2d1a6d546cec9beef7536ff8f5c49a3f9c78179d8ac/pillow-13.0.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:4bd024d567eca931b668978d672dba914d4b8e2628e2c3073a2ad178de8a2620", size = 6618042, upload-time = "2026-10-16T22:15:52.054Z" },
+    { url = "https://files.pythonhosted.org/packages/a8/7c/45ee28531f89ee4582967d8a17f7a26d71ca64cd902b93c269590cd60a00/pillow-13.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:23edf4ec4b6f4f240e8df6e4985e9ee2ced86474632433e55fb1e2548744bf40", size = 7228137, upload-time = "2026-10-16T22:15:54.009Z" },
+    { url = "https://files.pythonhosted.org/packages/e4/b0/2594f202d743a35d4fe80682bfd96f277e5ca7aecb27cf64ef5bfc1d000b/pillow-13.0.0-pp312-pypy312_pp80-macosx_11_0_arm64.whl", hash = "sha256:7e91e97a929232cdb828889ee6ad0e48d7f40f4ec38818f4e4a8968b2b1af22e", size = 4806808, upload-time = "2026-10-16T22:18:01.188Z" },
+    { url = "https://files.pythonhosted.org/packages/8e/ed/d504165ef79af6fafe25c76726cf30f544cc236002f7ed054f1d51adc4ff/pillow-13.0.0-pp312-pypy312_pp80-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:93ff9c9c3663ce23aaa5d0b6b6abac55359e2a0c06c7e30f3771829af816be63", size = 5464358, upload-time = "2026-10-16T22:18:02.779Z" },
+    { url = "https://files.pythonhosted.org/packages/7b/fe/a10aba5503ca13e9e009b8bd0a1a4ce5584bdce246245fd6f3a83701369c/pillow-13.0.0-pp312-pypy312_pp80-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e65dd6e75de67b2362a4fa3e16aa627595f65054dc4ec73d64680c84d592fa61", size = 6043514, upload-time = "2026-10-16T22:18:04.538Z" },
 ]
 
 [[package]]
@@ -5145,7 +5140,7 @@ resolution-markers = [
     "platform_machine == 'aarch64' and sys_platform == 'linux'",
 ]
 dependencies = [
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
 ]
 wheels = [
     { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/prime_kernels-0.1.0+cu130torch2.13.0-cp312-cp312-linux_aarch64.whl", hash = "sha256:695d61bd1b9ba99969ca7ac1410f5a35c1d1c747dab13a92c6f9c3a8def9b41a" },
@@ -5162,7 +5157,7 @@ resolution-markers = [
     "platform_machine == 'x86_64' and sys_platform == 'linux'",
 ]
 dependencies = [
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
 ]
 wheels = [
     { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/prime_kernels-0.1.0+cu130torch2.13.0-cp312-cp312-linux_x86_64.whl", hash = "sha256:c181b28773116970f08c130d2576ea89d81174820b2cfa03006dcfed4c23d462" },
@@ -5247,15 +5242,14 @@ all = [
     { name = "deep-select", version = "1.0.0+bfa4507.20261026.121543", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121543-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
     { name = "deep-select", version = "1.0.0+bfa4507.20261026.121630", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121630-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "dion", marker = "sys_platform == 'linux'" },
-    { name = "flash-attn", version = "2.8.3+cu130torch2.13", source = { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.47/flash_attn-2.8.3+cu130torch2.13-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "flash-attn", version = "2.8.3+cu130torch2.13", source = { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.48/flash_attn-2.8.3+cu130torch2.13-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
+    { name = "flash-attn", version = "2.8.3+cu130torch2.14", source = { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.10.0/flash_attn-2.8.3+cu130torch2.14-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+    { name = "flash-attn", version = "2.8.3+cu130torch2.14", source = { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.10.1/flash_attn-2.8.3+cu130torch2.14-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
     { name = "flash-attn-3", marker = "sys_platform == 'linux'" },
     { name = "flash-attn-4", marker = "sys_platform == 'linux'" },
     { name = "flash-linear-attention", marker = "sys_platform == 'linux'" },
-    { name = "flash-mla", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "flashinfer-cubin", marker = "sys_platform == 'linux'" },
-    { name = "flashinfer-jit-cache", version = "0.7.0.post1+cu130", source = { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache-0.7.0.post1%2Bcu130-cp39-abi3-manylinux_2_28_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
-    { name = "flashinfer-jit-cache", version = "0.7.0.post1+cu130", source = { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache-0.7.0.post1%2Bcu130-cp39-abi3-manylinux_2_28_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+    { name = "flashinfer-jit-cache", version = "0.7.2.post1+cu130", source = { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache-0.7.2.post1%2Bcu130-cp39-abi3-manylinux_2_28_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
+    { name = "flashinfer-jit-cache", version = "0.7.2.post1+cu130", source = { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache-0.7.2.post1%2Bcu130-cp39-abi3-manylinux_2_28_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "mamba-ssm", marker = "sys_platform == 'linux'" },
     { name = "modelexpress", marker = "sys_platform == 'linux'" },
     { name = "mooncake-transfer-engine", marker = "sys_platform == 'linux'" },
@@ -5266,16 +5260,16 @@ all = [
     { name = "quack-kernels", marker = "sys_platform == 'linux'" },
     { name = "sonic-moe", marker = "sys_platform == 'linux'" },
     { name = "tilelang" },
-    { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux'" },
+    { name = "torch", version = "2.14.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux'" },
     { name = "torchao", version = "0.18.0+git5f2baf9d5", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/torchao-0.18.0+git5f2baf9d5-cp310-abi3-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
     { name = "torchao", version = "0.18.0+git5f2baf9d5", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/torchao-0.18.0+git5f2baf9d5-cp310-abi3-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "torchaudio", marker = "sys_platform == 'linux'" },
     { name = "torchdata", marker = "sys_platform == 'linux'" },
     { name = "torchtitan", marker = "sys_platform == 'linux'" },
     { name = "torchvision", marker = "sys_platform == 'linux'" },
-    { name = "vllm", version = "0.32.0", source = { url = "https://github.com/vllm-project/vllm/releases/download/v0.32.0/vllm-0.32.0-cp38-abi3-manylinux_2_28_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
-    { name = "vllm", version = "0.32.0", source = { url = "https://github.com/vllm-project/vllm/releases/download/v0.32.0/vllm-0.32.0-cp38-abi3-manylinux_2_28_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+    { name = "vllm", version = "0.33.0", source = { url = "https://github.com/vllm-project/vllm/releases/download/v0.33.0rc3/vllm-0.33.0-cp38-abi3-manylinux_2_28_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
+    { name = "vllm", version = "0.33.0", source = { url = "https://github.com/vllm-project/vllm/releases/download/v0.33.0rc3/vllm-0.33.0-cp38-abi3-manylinux_2_28_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "vllm-router", version = "0.2.3", source = { url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.2.3/vllm_router-0.2.3-cp38-abi3-manylinux_2_28_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
     { name = "vllm-router", version = "0.2.3", source = { url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.2.3/vllm_router-0.2.3-cp38-abi3-manylinux_2_28_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
 ]
@@ -5296,8 +5290,8 @@ dynamo = [
     { name = "ai-dynamo", marker = "sys_platform == 'linux'" },
 ]
 flash-attn = [
-    { name = "flash-attn", version = "2.8.3+cu130torch2.13", source = { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.47/flash_attn-2.8.3+cu130torch2.13-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "flash-attn", version = "2.8.3+cu130torch2.13", source = { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.48/flash_attn-2.8.3+cu130torch2.13-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
+    { name = "flash-attn", version = "2.8.3+cu130torch2.14", source = { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.10.0/flash_attn-2.8.3+cu130torch2.14-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+    { name = "flash-attn", version = "2.8.3+cu130torch2.14", source = { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.10.1/flash_attn-2.8.3+cu130torch2.14-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
 ]
 flash-attn-3 = [
     { name = "flash-attn-3", marker = "sys_platform == 'linux'" },
@@ -5310,10 +5304,9 @@ gpu = [
     { name = "deep-select", version = "1.0.0+bfa4507.20261026.121630", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121630-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "dion", marker = "sys_platform == 'linux'" },
     { name = "flash-linear-attention", marker = "sys_platform == 'linux'" },
-    { name = "flash-mla", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "flashinfer-cubin", marker = "sys_platform == 'linux'" },
-    { name = "flashinfer-jit-cache", version = "0.7.0.post1+cu130", source = { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache-0.7.0.post1%2Bcu130-cp39-abi3-manylinux_2_28_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
-    { name = "flashinfer-jit-cache", version = "0.7.0.post1+cu130", source = { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache-0.7.0.post1%2Bcu130-cp39-abi3-manylinux_2_28_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+    { name = "flashinfer-jit-cache", version = "0.7.2.post1+cu130", source = { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache-0.7.2.post1%2Bcu130-cp39-abi3-manylinux_2_28_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
+    { name = "flashinfer-jit-cache", version = "0.7.2.post1+cu130", source = { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache-0.7.2.post1%2Bcu130-cp39-abi3-manylinux_2_28_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "mamba-ssm", marker = "sys_platform == 'linux'" },
     { name = "modelexpress", marker = "sys_platform == 'linux'" },
     { name = "mooncake-transfer-engine", marker = "sys_platform == 'linux'" },
@@ -5321,16 +5314,16 @@ gpu = [
     { name = "nvidia-cudnn-frontend", marker = "sys_platform == 'linux'" },
     { name = "nvidia-ml-py", marker = "sys_platform == 'linux'" },
     { name = "tilelang" },
-    { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux'" },
+    { name = "torch", version = "2.14.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux'" },
     { name = "torchao", version = "0.18.0+git5f2baf9d5", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/torchao-0.18.0+git5f2baf9d5-cp310-abi3-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
     { name = "torchao", version = "0.18.0+git5f2baf9d5", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/torchao-0.18.0+git5f2baf9d5-cp310-abi3-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "torchaudio", marker = "sys_platform == 'linux'" },
     { name = "torchdata", marker = "sys_platform == 'linux'" },
     { name = "torchtitan", marker = "sys_platform == 'linux'" },
     { name = "torchvision", marker = "sys_platform == 'linux'" },
-    { name = "vllm", version = "0.32.0", source = { url = "https://github.com/vllm-project/vllm/releases/download/v0.32.0/vllm-0.32.0-cp38-abi3-manylinux_2_28_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
-    { name = "vllm", version = "0.32.0", source = { url = "https://github.com/vllm-project/vllm/releases/download/v0.32.0/vllm-0.32.0-cp38-abi3-manylinux_2_28_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+    { name = "vllm", version = "0.33.0", source = { url = "https://github.com/vllm-project/vllm/releases/download/v0.33.0rc3/vllm-0.33.0-cp38-abi3-manylinux_2_28_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
+    { name = "vllm", version = "0.33.0", source = { url = "https://github.com/vllm-project/vllm/releases/download/v0.33.0rc3/vllm-0.33.0-cp38-abi3-manylinux_2_28_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
 ]
 kernels = [
     { name = "prime-kernels", version = "0.1.0+cu130torch2.13.0", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/prime_kernels-0.1.0+cu130torch2.13.0-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
@@ -5372,26 +5365,25 @@ requires-dist = [
     { name = "dion", marker = "sys_platform == 'linux' and extra == 'gpu'", git = "https://github.com/samsja/dion.git?rev=0c801b0" },
     { name = "fastapi", specifier = ">=0.115" },
     { name = "fastapi", marker = "extra == 'dashboard'", specifier = ">=0.115" },
-    { name = "flash-attn", marker = "platform_machine == 'aarch64' and extra == 'flash-attn'", url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.48/flash_attn-2.8.3+cu130torch2.13-cp312-cp312-linux_aarch64.whl" },
-    { name = "flash-attn", marker = "platform_machine == 'x86_64' and extra == 'flash-attn'", url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.47/flash_attn-2.8.3+cu130torch2.13-cp312-cp312-linux_x86_64.whl" },
+    { name = "flash-attn", marker = "platform_machine == 'aarch64' and extra == 'flash-attn'", url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.10.1/flash_attn-2.8.3+cu130torch2.14-cp312-cp312-linux_aarch64.whl" },
+    { name = "flash-attn", marker = "platform_machine == 'x86_64' and extra == 'flash-attn'", url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.10.0/flash_attn-2.8.3+cu130torch2.14-cp312-cp312-linux_x86_64.whl" },
     { name = "flash-attn-3", marker = "sys_platform == 'linux' and extra == 'flash-attn-3'", index = "https://download.pytorch.org/whl/test/cu130" },
     { name = "flash-attn-4", marker = "sys_platform == 'linux' and extra == 'flash-attn-cute'", git = "https://github.com/Dao-AILab/flash-attention.git?subdirectory=flash_attn%2Fcute&rev=0251105" },
     { name = "flash-linear-attention", marker = "sys_platform == 'linux' and extra == 'gpu'", git = "https://github.com/fla-org/flash-linear-attention" },
-    { name = "flash-mla", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'gpu'", url = "https://github.com/NovaSky-AI/skyrl-wheels/releases/download/cu13torch2.13-flashmla-b7643bd/flash_mla-1.0.0%2Bcu.13.0.torch.2.13.b7643bd-cp312-cp312-linux_x86_64.whl" },
-    { name = "flashinfer-cubin", marker = "sys_platform == 'linux' and extra == 'gpu'", specifier = "==0.7.0.post1", index = "https://flashinfer.ai/whl" },
-    { name = "flashinfer-jit-cache", marker = "platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'gpu'", url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache-0.7.0.post1%2Bcu130-cp39-abi3-manylinux_2_28_aarch64.whl" },
-    { name = "flashinfer-jit-cache", marker = "platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux' and extra == 'gpu'", specifier = "==0.7.0.post1+cu130" },
-    { name = "flashinfer-jit-cache", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'gpu'", url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_jit_cache-0.7.0.post1%2Bcu130-cp39-abi3-manylinux_2_28_x86_64.whl" },
+    { name = "flashinfer-cubin", marker = "sys_platform == 'linux' and extra == 'gpu'", specifier = "==0.7.2.post1", index = "https://flashinfer.ai/whl" },
+    { name = "flashinfer-jit-cache", marker = "platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'gpu'", url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache-0.7.2.post1%2Bcu130-cp39-abi3-manylinux_2_28_aarch64.whl" },
+    { name = "flashinfer-jit-cache", marker = "platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux' and extra == 'gpu'", specifier = "==0.7.2.post1+cu130" },
+    { name = "flashinfer-jit-cache", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'gpu'", url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache-0.7.2.post1%2Bcu130-cp39-abi3-manylinux_2_28_x86_64.whl" },
     { name = "jaxtyping", specifier = ">=0.3.2" },
     { name = "loguru", specifier = ">=0.7.3" },
     { name = "mamba-ssm", marker = "sys_platform == 'linux' and extra == 'gpu'", git = "https://github.com/state-spaces/mamba?rev=f577286" },
     { name = "modelexpress", marker = "sys_platform == 'linux' and extra == 'gpu'", specifier = "==0.8.0.dev20261022" },
     { name = "mooncake-transfer-engine", marker = "sys_platform == 'linux' and extra == 'gpu'", specifier = ">=0.3.10.post2" },
     { name = "msgspec", specifier = ">=0.18" },
-    { name = "nixl", marker = "sys_platform == 'linux' and extra == 'disagg'" },
-    { name = "nixl-cu13", marker = "sys_platform == 'linux' and extra == 'gpu'", specifier = "==0.10.1" },
+    { name = "nixl", marker = "sys_platform == 'linux' and extra == 'disagg'", specifier = "==1.5.0" },
+    { name = "nixl-cu13", marker = "sys_platform == 'linux' and extra == 'gpu'", specifier = "==1.5.0" },
     { name = "numpy", specifier = ">=2.2.6" },
-    { name = "nvidia-cudnn-frontend", marker = "sys_platform == 'linux' and extra == 'gpu'", specifier = ">=1.30.0" },
+    { name = "nvidia-cudnn-frontend", marker = "sys_platform == 'linux' and extra == 'gpu'", specifier = ">=1.31.0" },
     { name = "nvidia-ml-py", marker = "sys_platform == 'linux' and extra == 'gpu'", specifier = ">=12.575.51" },
     { name = "openai", specifier = ">=1.106.1" },
     { name = "orjson", specifier = ">=3.11.0" },
@@ -5422,7 +5414,7 @@ requires-dist = [
     { name = "sympy", specifier = ">=1.13.3" },
     { name = "tenacity", specifier = ">=8.2.0" },
     { name = "tilelang", marker = "sys_platform == 'linux' and extra == 'gpu'", specifier = ">=0.1.8" },
-    { name = "torch", marker = "sys_platform == 'linux' and extra == 'gpu'", specifier = ">=2.13.0", index = "https://download.pytorch.org/whl/cu130" },
+    { name = "torch", marker = "sys_platform == 'linux' and extra == 'gpu'", specifier = ">=2.14.0", index = "https://download.pytorch.org/whl/cu130" },
     { name = "torchao", marker = "platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'gpu'", url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/torchao-0.18.0+git5f2baf9d5-cp310-abi3-linux_aarch64.whl" },
     { name = "torchao", marker = "platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux' and extra == 'gpu'", specifier = ">=0.18.0" },
     { name = "torchao", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'gpu'", url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/torchao-0.18.0+git5f2baf9d5-cp310-abi3-linux_x86_64.whl" },
@@ -5430,14 +5422,14 @@ requires-dist = [
     { name = "torchdata", marker = "sys_platform == 'linux' and extra == 'gpu'", specifier = ">=0.11.0" },
     { name = "torchtitan", marker = "sys_platform == 'linux' and extra == 'gpu'", git = "https://github.com/pytorch/torchtitan?rev=23e4dfc" },
     { name = "torchvision", marker = "sys_platform == 'linux' and extra == 'gpu'", index = "https://download.pytorch.org/whl/cu130" },
-    { name = "transformers", specifier = "==5.19.0" },
+    { name = "transformers", specifier = "==5.20.2" },
     { name = "uvicorn", specifier = ">=0.30" },
     { name = "uvicorn", marker = "extra == 'dashboard'", specifier = ">=0.30" },
     { name = "uvloop", specifier = ">=0.21.0" },
     { name = "verifiers", extras = ["harbor"], editable = "deps/verifiers" },
-    { name = "vllm", marker = "platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'gpu'", url = "https://github.com/vllm-project/vllm/releases/download/v0.32.0/vllm-0.32.0-cp38-abi3-manylinux_2_28_aarch64.whl" },
-    { name = "vllm", marker = "platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux' and extra == 'gpu'", specifier = ">=0.32.0" },
-    { name = "vllm", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'gpu'", url = "https://github.com/vllm-project/vllm/releases/download/v0.32.0/vllm-0.32.0-cp38-abi3-manylinux_2_28_x86_64.whl" },
+    { name = "vllm", marker = "platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'gpu'", url = "https://github.com/vllm-project/vllm/releases/download/v0.33.0rc3/vllm-0.33.0-cp38-abi3-manylinux_2_28_aarch64.whl" },
+    { name = "vllm", marker = "platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux' and extra == 'gpu'", specifier = ">=0.33.0" },
+    { name = "vllm", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'gpu'", url = "https://github.com/vllm-project/vllm/releases/download/v0.33.0rc3/vllm-0.33.0-cp38-abi3-manylinux_2_28_x86_64.whl" },
     { name = "vllm-router", marker = "platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.2.3/vllm_router-0.2.3-cp38-abi3-manylinux_2_28_aarch64.whl" },
     { name = "vllm-router", marker = "platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux' and extra == 'disagg'" },
     { name = "vllm-router", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.2.3/vllm_router-0.2.3-cp38-abi3-manylinux_2_28_x86_64.whl" },
@@ -6204,7 +6196,7 @@ dependencies = [
     { name = "apache-tvm-ffi" },
     { name = "einops" },
     { name = "nvidia-cutlass-dsl" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
     { name = "torch-c-dlpack-ext" },
 ]
 
@@ -6503,7 +6495,7 @@ source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "humanize" },
     { name = "numpy" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/b9/12/7fe4ed110506c992a1194983a816f16e0f2410416952207f89c895af7d04/runai_model_streamer-0.17.0-py3-none-manylinux_2_30_aarch64.whl", hash = "sha256:e722bb0338ed9b047c140745d3f5431af01fd2d5d6d809be01423bc6c62d7eb2", size = 745324, upload-time = "2026-10-13T16:09:10.62Z" },
@@ -6706,8 +6698,8 @@ dependencies = [
     { name = "numpy" },
     { name = "scikit-learn" },
     { name = "scipy" },
-    { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux'" },
+    { name = "torch", version = "2.14.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux'" },
     { name = "tqdm" },
     { name = "transformers" },
     { name = "typing-extensions" },
@@ -6860,7 +6852,7 @@ version = "0.1.2.post1"
 source = { git = "https://github.com/S1ro1/sonic-moe?rev=0bf06ac#0bf06ac7bf5b0800a7202743c1fb5722ed004350" }
 dependencies = [
     { name = "quack-kernels" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
 ]
 
 [[package]]
@@ -7335,8 +7327,8 @@ dependencies = [
     { name = "numpy" },
     { name = "psutil" },
     { name = "setuptools", version = "77.0.3", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" },
-    { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux'" },
+    { name = "torch", version = "2.14.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux'" },
     { name = "torch-c-dlpack-ext" },
     { name = "tqdm" },
     { name = "typing-extensions" },
@@ -7378,16 +7370,16 @@ wheels = [
 
 [[package]]
 name = "tokenspeed-mla"
-version = "0.2.15"
+version = "2.0.3"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "apache-tvm-ffi" },
     { name = "nvidia-cutlass-dsl" },
     { name = "tokenspeed-triton" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
 ]
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/14/56/495ed0feea7c332e02bbf23fa565aee3a176fe463c81ff7cc85afc01817e/tokenspeed_mla-0.2.15-py3-none-any.whl", hash = "sha256:7b4d8bfe92e1ad984e94db262671ef33ce2ecfe0a44f79c5ea1c2816d8397a1f", size = 151477, upload-time = "2026-09-28T19:49:06.103Z" },
+    { url = "https://files.pythonhosted.org/packages/39/74/cce1840ddb1736baf8c3115ced5f657f98151bc42d5e84f8f9c2f0f3208d/tokenspeed_mla-2.0.3-py3-none-any.whl", hash = "sha256:576ce91e861598347f09c429fe41269863d351bf3de8410a845dae39fa172e4a", size = 238130, upload-time = "2026-10-14T15:18:03.665Z" },
 ]
 
 [[package]]
@@ -7433,7 +7425,7 @@ wheels = [
 
 [[package]]
 name = "torch"
-version = "2.13.0"
+version = "2.14.0"
 source = { registry = "https://pypi.org/simple" }
 resolution-markers = [
     "platform_machine == 'arm64' and sys_platform == 'darwin'",
@@ -7448,12 +7440,12 @@ dependencies = [
     { name = "typing-extensions" },
 ]
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/c4/3a/ed0f4d4d1dcde03bced7aac9a28e800abcdc0cbd06b6775044c9fbd877b7/torch-2.13.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:2fe228aba290d14b9f31b049be550dbd469c3fd3013d7a19705b30454da97027", size = 111213045, upload-time = "2026-07-08T16:05:22.997Z" },
+    { url = "https://files.pythonhosted.org/packages/17/76/bb4770f56cf6d8971671dbcbb7493e5a6a15ad2825f4e359b02c27c38297/torch-2.14.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:c1f844f1c750e87df4b68bc3afbc0e2b0c7ef19d7b8f666e48bdcf6a0c4f0056", size = 127303200, upload-time = "2026-09-02T13:43:20.311Z" },
 ]
 
 [[package]]
 name = "torch"
-version = "2.13.0+cu130"
+version = "2.14.0+cu130"
 source = { registry = "https://download.pytorch.org/whl/cu130" }
 resolution-markers = [
     "platform_machine == 'x86_64' and sys_platform == 'linux'",
@@ -7476,8 +7468,8 @@ dependencies = [
     { name = "typing-extensions" },
 ]
 wheels = [
-    { url = "https://download-r2.pytorch.org/whl/cu130/torch-2.13.0%2Bcu130-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:5ebd552c887e707c8e64927aceb8377ca2e81588c4e7494bcd23cb8ac0aca14d", upload-time = "2026-07-08T20:24:19Z" },
-    { url = "https://download-r2.pytorch.org/whl/cu130/torch-2.13.0%2Bcu130-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:8db7338e6895c3d4bd89a02ff4209507d1f0cf2ffeb3b898538b5a07d1ea8c1e", upload-time = "2026-07-08T20:24:52Z" },
+    { url = "https://download-r2.pytorch.org/whl/cu130/torch-2.14.0%2Bcu130-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:e729d57eb3e0c859fa7a9c27f9f274121c2c09ae70a376ac227a650807b0457c", upload-time = "2026-09-02T18:51:18Z" },
+    { url = "https://download-r2.pytorch.org/whl/cu130/torch-2.14.0%2Bcu130-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:9c5693d3ca91743202034094ad3b3df20b5eb1373d5e0bc8a5d976be12c35302", upload-time = "2026-09-02T18:51:32Z" },
 ]
 
 [[package]]
@@ -7485,8 +7477,8 @@ name = "torch-c-dlpack-ext"
 version = "0.1.5"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux'" },
+    { name = "torch", version = "2.14.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/37/de/921b6491efce5c389a5ef9bbed3d2d6660005840dae488124173180859ab/torch_c_dlpack_ext-0.1.5.tar.gz", hash = "sha256:d06f0357d575d22a168cc77acb9020fc4bae30968ceb6718a055dcbe92bacabe", size = 12913, upload-time = "2026-01-12T11:25:08.484Z" }
 wheels = [
@@ -7601,7 +7593,7 @@ version = "0.11.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "requests" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
     { name = "urllib3" },
 ]
 wheels = [
@@ -7624,16 +7616,16 @@ dependencies = [
 
 [[package]]
 name = "torchvision"
-version = "0.28.0+cu130"
+version = "0.29.0+cu130"
 source = { registry = "https://download.pytorch.org/whl/cu130" }
 dependencies = [
     { name = "numpy" },
     { name = "pillow" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
 ]
 wheels = [
-    { url = "https://download-r2.pytorch.org/whl/cu130/torchvision-0.28.0%2Bcu130-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:ecf72161734b0cf75aabeb2a83101ee021ba8a9cfea57b40050deed7accd4615", upload-time = "2026-07-08T12:26:52Z" },
-    { url = "https://download-r2.pytorch.org/whl/cu130/torchvision-0.28.0%2Bcu130-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:8a0008d34ccc4e81066b97ff0ae5a34c676bfdf3464baf40c01b320dc9a45ce0", upload-time = "2026-07-08T12:26:52Z" },
+    { url = "https://download-r2.pytorch.org/whl/cu130/torchvision-0.29.0%2Bcu130-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:7678a0ed87d39580ab0369b4183d815ba92b84aa6a2c5bd60c67ab80274acef9", upload-time = "2026-09-02T00:28:04Z" },
+    { url = "https://download-r2.pytorch.org/whl/cu130/torchvision-0.29.0%2Bcu130-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:fb681a367becdbf40cd559f1b553be79832885a58414980e77bed192fd960673", upload-time = "2026-09-02T00:28:04Z" },
 ]
 
 [[package]]
@@ -7686,7 +7678,7 @@ wheels = [
 
 [[package]]
 name = "transformers"
-version = "5.19.0"
+version = "5.20.2"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "huggingface-hub" },
@@ -7699,9 +7691,9 @@ dependencies = [
     { name = "tqdm" },
     { name = "typer" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/65/4c/70012ed0795235fb6bd7806fafd4a1842f03e49ebbfa3bb9580e451948ce/transformers-5.19.0.tar.gz", hash = "sha256:87f38dd25e4521151b97e94520ac457f44a0ae8a8358a5b112daff6c64a822d6", size = 10032022, upload-time = "2026-10-06T16:39:00.229Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/8c/4e/3cd6f08cc84ee612d40f3cdb1649475e7c997f2d5cbd95ade8780708af62/transformers-5.20.2.tar.gz", hash = "sha256:080871e061546d720be57f5846fce426f46b44d61574907cc7404c2cba49d180", size = 10225379, upload-time = "2026-10-19T17:46:03.092Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/b9/59/1e7f1212c215c73690ae6cab4678efe1555ff985953e145d88cc212859b0/transformers-5.19.0-py3-none-any.whl", hash = "sha256:afcd2dd5f603ed28c1e1fcb00a338ccbb4ef5f878ed289635df8b58187afb518", size = 12599123, upload-time = "2026-10-06T16:38:56.794Z" },
+    { url = "https://files.pythonhosted.org/packages/96/b6/724faefed06913dd3e42e0476fe2c42448675c42f687a12bf810aa098005/transformers-5.20.2-py3-none-any.whl", hash = "sha256:7ae75dbc1fb3922aff1b6c97b9e5479d6cf26fd86a8866bfb71324190479f807", size = 12844921, upload-time = "2026-10-19T17:45:59.293Z" },
 ]
 
 [[package]]
@@ -7743,11 +7735,11 @@ wheels = [
 
 [[package]]
 name = "triton"
-version = "3.7.1"
+version = "3.8.0"
 source = { registry = "https://pypi.org/simple" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/94/fa/f856e24deb462d5f18bd4b5a746957862ab9b6ee5834bda60605ec348366/triton-3.7.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9497f2e696ee368862a181a90b2dcc03ca978cc4f602abd67c7d81022a6988e1", size = 184692359, upload-time = "2026-06-17T20:03:48.288Z" },
-    { url = "https://files.pythonhosted.org/packages/c4/6f/fb96d15db6f36d6eae4cafb998c2e0353bf59d7c4ea1662d7497f269134a/triton-3.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7e40869937a68206ec70d7f25bb7ec6433cb083f9135e1f36dbd318dc449a728", size = 197719725, upload-time = "2026-06-17T19:53:20.419Z" },
+    { url = "https://files.pythonhosted.org/packages/87/07/0f8cd8e8db0472334253efdaaab3d0819fea27aa99bf0e7f1aeea4ceb5ae/triton-3.8.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a9c404c69ed4a39e8ec632eaf6b9fe058a060bf98979c177f6ef666f06bb8d50", size = 226474486, upload-time = "2026-08-28T16:08:18.29Z" },
+    { url = "https://files.pythonhosted.org/packages/c1/09/b7012e5bfae67640f268aa584caa80fe1674f6b0da949046b679972c33e3/triton-3.8.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e91ffa46d095b252248297292dd22bcbacd53a125a0c2eefbbbf74925a320bc3", size = 247972921, upload-time = "2026-08-28T15:55:53.157Z" },
 ]
 
 [[package]]
@@ -8112,8 +8104,8 @@ wheels = [
 
 [[package]]
 name = "vllm"
-version = "0.32.0"
-source = { url = "https://github.com/vllm-project/vllm/releases/download/v0.32.0/vllm-0.32.0-cp38-abi3-manylinux_2_28_aarch64.whl" }
+version = "0.33.0"
+source = { url = "https://github.com/vllm-project/vllm/releases/download/v0.33.0rc3/vllm-0.33.0-cp38-abi3-manylinux_2_28_aarch64.whl" }
 resolution-markers = [
     "platform_machine == 'aarch64' and sys_platform == 'linux'",
 ]
@@ -8124,6 +8116,7 @@ dependencies = [
     { name = "blake3" },
     { name = "cachetools" },
     { name = "cbor2" },
+    { name = "certifi" },
     { name = "cloudpickle" },
     { name = "compressed-tensors" },
     { name = "depyf" },
@@ -8144,6 +8137,7 @@ dependencies = [
     { name = "mistral-common", extra = ["image"] },
     { name = "model-hosting-container-standards" },
     { name = "msgspec" },
+    { name = "multidict" },
     { name = "ninja" },
     { name = "numba" },
     { name = "numpy" },
@@ -8156,6 +8150,7 @@ dependencies = [
     { name = "opentelemetry-exporter-otlp" },
     { name = "opentelemetry-sdk" },
     { name = "opentelemetry-semantic-conventions-ai" },
+    { name = "orjson" },
     { name = "oss-harmony" },
     { name = "outlines-core" },
     { name = "partial-json-parser" },
@@ -8184,7 +8179,7 @@ dependencies = [
     { name = "tilelang" },
     { name = "tokenizers" },
     { name = "tokenspeed-mla" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
     { name = "torchaudio" },
     { name = "torchcodec" },
     { name = "torchvision" },
@@ -8195,12 +8190,12 @@ dependencies = [
     { name = "xgrammar" },
 ]
 wheels = [
-    { url = "https://github.com/vllm-project/vllm/releases/download/v0.32.0/vllm-0.32.0-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:d0620dbdf99528bcfebe551a8906347a229921fb1602f887bd1286a920da2d44" },
+    { url = "https://github.com/vllm-project/vllm/releases/download/v0.33.0rc3/vllm-0.33.0-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:e9e404f52c8494b69bd83fe99fe7ecfa9f6b89fbd156d2a9639c24ea41d6d606" },
 ]
 
 [package.metadata]
 requires-dist = [
-    { name = "aiohttp", specifier = ">=3.14.3" },
+    { name = "aiohttp", specifier = ">=3.14.4" },
     { name = "anthropic", specifier = ">=0.71.0" },
     { name = "apache-tvm-ffi", specifier = "==0.1.11" },
     { name = "av", marker = "extra == 'audio'" },
@@ -8208,6 +8203,7 @@ requires-dist = [
     { name = "blake3" },
     { name = "cachetools" },
     { name = "cbor2" },
+    { name = "certifi", specifier = ">=2026.10.19" },
     { name = "cloudpickle" },
     { name = "compressed-tensors", specifier = "==0.17.0" },
     { name = "datasets", marker = "extra == 'bench'" },
@@ -8217,10 +8213,10 @@ requires-dist = [
     { name = "fastsafetensors", specifier = ">=0.3.3" },
     { name = "fastsafetensors", marker = "extra == 'fastsafetensors'", specifier = ">=0.3.3" },
     { name = "filelock", specifier = ">=3.20.3" },
-    { name = "flashinfer-python", specifier = "==0.7.0.post1" },
+    { name = "flashinfer-python", specifier = "==0.7.2.post1" },
     { name = "helion", marker = "extra == 'helion'", specifier = "==1.4.0" },
     { name = "huggingface-hub", specifier = ">=1.31.0" },
-    { name = "humming-kernels", extras = ["cu13"], specifier = "==0.1.17" },
+    { name = "humming-kernels", extras = ["cu13"], specifier = "==0.1.19" },
     { name = "ijson" },
     { name = "instanttensor", specifier = ">=0.2.2" },
     { name = "instanttensor", marker = "extra == 'instanttensor'", specifier = ">=0.2.2" },
@@ -8229,15 +8225,16 @@ requires-dist = [
     { name = "llguidance", marker = "platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'ppc64le' or platform_machine == 's390x' or platform_machine == 'x86_64'", specifier = ">=1.7.0,<1.8.0" },
     { name = "lm-format-enforcer", specifier = "==0.11.3" },
     { name = "matplotlib", marker = "extra == 'bench'" },
-    { name = "mcp", specifier = ">=2.0.0,<3.0.0" },
+    { name = "mcp", specifier = ">=2.2.0,<3.0.0" },
     { name = "mistral-common", extras = ["audio"], marker = "extra == 'audio'" },
     { name = "mistral-common", extras = ["image"], specifier = ">=1.11.6" },
     { name = "model-hosting-container-standards", specifier = ">=0.1.14,<1.0.0" },
     { name = "msgspec" },
+    { name = "multidict", specifier = ">=6.9.1" },
     { name = "ninja" },
     { name = "numba", specifier = "==0.65.0" },
     { name = "numpy" },
-    { name = "nvidia-cudnn-frontend", specifier = ">=1.19.1" },
+    { name = "nvidia-cudnn-frontend", specifier = ">=1.31.0" },
     { name = "nvidia-cutlass-dsl", extras = ["cu13"], specifier = "==4.8.0" },
     { name = "nvidia-deepstream-videodecode-cu13", marker = "extra == 'deepstream'", specifier = ">=9.0.2" },
     { name = "nvtx", specifier = "==0.2.15" },
@@ -8251,11 +8248,12 @@ requires-dist = [
     { name = "opentelemetry-sdk", marker = "extra == 'otel'", specifier = ">=1.26.0" },
     { name = "opentelemetry-semantic-conventions-ai", specifier = ">=0.4.1" },
     { name = "opentelemetry-semantic-conventions-ai", marker = "extra == 'otel'", specifier = ">=0.4.1" },
+    { name = "orjson", specifier = ">=3.11.6" },
     { name = "oss-harmony", specifier = ">=0.0.11" },
     { name = "outlines-core", specifier = "==0.2.14" },
     { name = "pandas", marker = "extra == 'bench'" },
     { name = "partial-json-parser" },
-    { name = "pillow", specifier = ">=12.3.0" },
+    { name = "pillow", specifier = ">=13.0.0" },
     { name = "plotly", marker = "extra == 'bench'" },
     { name = "prometheus-client", specifier = ">=0.18.0" },
     { name = "prometheus-fastapi-instrumentator", specifier = ">=8.0.0" },
@@ -8289,25 +8287,25 @@ requires-dist = [
     { name = "tiktoken", specifier = ">=0.6.0" },
     { name = "tilelang", specifier = "==0.1.12" },
     { name = "tokenizers", specifier = ">=0.23.2" },
-    { name = "tokenspeed-mla", marker = "sys_platform == 'linux'", specifier = "==0.2.15" },
-    { name = "torch", specifier = "==2.13.0" },
+    { name = "tokenspeed-mla", marker = "sys_platform == 'linux'", specifier = "==2.0.3" },
+    { name = "torch", specifier = "==2.14.0" },
     { name = "torchaudio", specifier = "==2.11.0" },
     { name = "torchcodec", specifier = ">=0.14" },
-    { name = "torchvision", specifier = "==0.28.0" },
+    { name = "torchvision", specifier = "==0.29.0" },
     { name = "tqdm" },
-    { name = "transformers", specifier = ">=5.16.1,<5.20.0" },
+    { name = "transformers", specifier = ">=5.16.1,<5.21.0" },
     { name = "typing-extensions", specifier = ">=4.10" },
     { name = "vllm-gguf-plugin", marker = "extra == 'extra-quant'", specifier = ">=0.0.2" },
     { name = "watchfiles" },
     { name = "xgrammar", marker = "platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'ppc64le' or platform_machine == 's390x' or platform_machine == 'x86_64'", specifier = "==0.2.8" },
-    { name = "zentorch", marker = "extra == 'zen'", specifier = "==2.13.0.1" },
+    { name = "zentorch", marker = "extra == 'zen'", specifier = "==2.14.0.1" },
 ]
 provides-extras = ["zen", "bench", "tensorizer", "fastsafetensors", "instanttensor", "runai", "audio", "video", "deepstream", "flashinfer", "b12x", "helion", "grpc", "otel", "extra-quant"]
 
 [[package]]
 name = "vllm"
-version = "0.32.0"
-source = { url = "https://github.com/vllm-project/vllm/releases/download/v0.32.0/vllm-0.32.0-cp38-abi3-manylinux_2_28_x86_64.whl" }
+version = "0.33.0"
+source = { url = "https://github.com/vllm-project/vllm/releases/download/v0.33.0rc3/vllm-0.33.0-cp38-abi3-manylinux_2_28_x86_64.whl" }
 resolution-markers = [
     "platform_machine == 'x86_64' and sys_platform == 'linux'",
 ]
@@ -8318,6 +8316,7 @@ dependencies = [
     { name = "blake3" },
     { name = "cachetools" },
     { name = "cbor2" },
+    { name = "certifi" },
     { name = "cloudpickle" },
     { name = "compressed-tensors" },
     { name = "depyf" },
@@ -8338,6 +8337,7 @@ dependencies = [
     { name = "mistral-common", extra = ["image"] },
     { name = "model-hosting-container-standards" },
     { name = "msgspec" },
+    { name = "multidict" },
     { name = "ninja" },
     { name = "numba" },
     { name = "numpy" },
@@ -8350,6 +8350,7 @@ dependencies = [
     { name = "opentelemetry-exporter-otlp" },
     { name = "opentelemetry-sdk" },
     { name = "opentelemetry-semantic-conventions-ai" },
+    { name = "orjson" },
     { name = "oss-harmony" },
     { name = "outlines-core" },
     { name = "partial-json-parser" },
@@ -8378,7 +8379,7 @@ dependencies = [
     { name = "tilelang" },
     { name = "tokenizers" },
     { name = "tokenspeed-mla" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
     { name = "torchaudio" },
     { name = "torchcodec" },
     { name = "torchvision" },
@@ -8389,12 +8390,12 @@ dependencies = [
     { name = "xgrammar" },
 ]
 wheels = [
-    { url = "https://github.com/vllm-project/vllm/releases/download/v0.32.0/vllm-0.32.0-cp38-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:5e2ecdc5b0baf72921c288486417ca82778456a448e4fe60b62ee43c0dfc8611" },
+    { url = "https://github.com/vllm-project/vllm/releases/download/v0.33.0rc3/vllm-0.33.0-cp38-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:4b6bdec4ab269990bf70d45a3c1e94c3760a27064513760193b26d6f0da883aa" },
 ]
 
 [package.metadata]
 requires-dist = [
-    { name = "aiohttp", specifier = ">=3.14.3" },
+    { name = "aiohttp", specifier = ">=3.14.4" },
     { name = "anthropic", specifier = ">=0.71.0" },
     { name = "apache-tvm-ffi", specifier = "==0.1.11" },
     { name = "av", marker = "extra == 'audio'" },
@@ -8402,6 +8403,7 @@ requires-dist = [
     { name = "blake3" },
     { name = "cachetools" },
     { name = "cbor2" },
+    { name = "certifi", specifier = ">=2026.10.19" },
     { name = "cloudpickle" },
     { name = "compressed-tensors", specifier = "==0.17.0" },
     { name = "datasets", marker = "extra == 'bench'" },
@@ -8411,10 +8413,10 @@ requires-dist = [
     { name = "fastsafetensors", specifier = ">=0.3.3" },
     { name = "fastsafetensors", marker = "extra == 'fastsafetensors'", specifier = ">=0.3.3" },
     { name = "filelock", specifier = ">=3.20.3" },
-    { name = "flashinfer-python", specifier = "==0.7.0.post1" },
+    { name = "flashinfer-python", specifier = "==0.7.2.post1" },
     { name = "helion", marker = "extra == 'helion'", specifier = "==1.4.0" },
     { name = "huggingface-hub", specifier = ">=1.31.0" },
-    { name = "humming-kernels", extras = ["cu13"], specifier = "==0.1.17" },
+    { name = "humming-kernels", extras = ["cu13"], specifier = "==0.1.19" },
     { name = "ijson" },
     { name = "instanttensor", specifier = ">=0.2.2" },
     { name = "instanttensor", marker = "extra == 'instanttensor'", specifier = ">=0.2.2" },
@@ -8423,15 +8425,16 @@ requires-dist = [
     { name = "llguidance", marker = "platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'ppc64le' or platform_machine == 's390x' or platform_machine == 'x86_64'", specifier = ">=1.7.0,<1.8.0" },
     { name = "lm-format-enforcer", specifier = "==0.11.3" },
     { name = "matplotlib", marker = "extra == 'bench'" },
-    { name = "mcp", specifier = ">=2.0.0,<3.0.0" },
+    { name = "mcp", specifier = ">=2.2.0,<3.0.0" },
     { name = "mistral-common", extras = ["audio"], marker = "extra == 'audio'" },
     { name = "mistral-common", extras = ["image"], specifier = ">=1.11.6" },
     { name = "model-hosting-container-standards", specifier = ">=0.1.14,<1.0.0" },
     { name = "msgspec" },
+    { name = "multidict", specifier = ">=6.9.1" },
     { name = "ninja" },
     { name = "numba", specifier = "==0.65.0" },
     { name = "numpy" },
-    { name = "nvidia-cudnn-frontend", specifier = ">=1.19.1" },
+    { name = "nvidia-cudnn-frontend", specifier = ">=1.31.0" },
     { name = "nvidia-cutlass-dsl", extras = ["cu13"], specifier = "==4.8.0" },
     { name = "nvidia-deepstream-videodecode-cu13", marker = "extra == 'deepstream'", specifier = ">=9.0.2" },
     { name = "nvtx", specifier = "==0.2.15" },
@@ -8445,11 +8448,12 @@ requires-dist = [
     { name = "opentelemetry-sdk", marker = "extra == 'otel'", specifier = ">=1.26.0" },
     { name = "opentelemetry-semantic-conventions-ai", specifier = ">=0.4.1" },
     { name = "opentelemetry-semantic-conventions-ai", marker = "extra == 'otel'", specifier = ">=0.4.1" },
+    { name = "orjson", specifier = ">=3.11.6" },
     { name = "oss-harmony", specifier = ">=0.0.11" },
     { name = "outlines-core", specifier = "==0.2.14" },
     { name = "pandas", marker = "extra == 'bench'" },
     { name = "partial-json-parser" },
-    { name = "pillow", specifier = ">=12.3.0" },
+    { name = "pillow", specifier = ">=13.0.0" },
     { name = "plotly", marker = "extra == 'bench'" },
     { name = "prometheus-client", specifier = ">=0.18.0" },
     { name = "prometheus-fastapi-instrumentator", specifier = ">=8.0.0" },
@@ -8483,18 +8487,18 @@ requires-dist = [
     { name = "tiktoken", specifier = ">=0.6.0" },
     { name = "tilelang", specifier = "==0.1.12" },
     { name = "tokenizers", specifier = ">=0.23.2" },
-    { name = "tokenspeed-mla", marker = "sys_platform == 'linux'", specifier = "==0.2.15" },
-    { name = "torch", specifier = "==2.13.0" },
+    { name = "tokenspeed-mla", marker = "sys_platform == 'linux'", specifier = "==2.0.3" },
+    { name = "torch", specifier = "==2.14.0" },
     { name = "torchaudio", specifier = "==2.11.0" },
     { name = "torchcodec", specifier = ">=0.14" },
-    { name = "torchvision", specifier = "==0.28.0" },
+    { name = "torchvision", specifier = "==0.29.0" },
     { name = "tqdm" },
-    { name = "transformers", specifier = ">=5.16.1,<5.20.0" },
+    { name = "transformers", specifier = ">=5.16.1,<5.21.0" },
     { name = "typing-extensions", specifier = ">=4.10" },
     { name = "vllm-gguf-plugin", marker = "extra == 'extra-quant'", specifier = ">=0.0.2" },
     { name = "watchfiles" },
     { name = "xgrammar", marker = "platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'ppc64le' or platform_machine == 's390x' or platform_machine == 'x86_64'", specifier = "==0.2.8" },
-    { name = "zentorch", marker = "extra == 'zen'", specifier = "==2.13.0.1" },
+    { name = "zentorch", marker = "extra == 'zen'", specifier = "==2.14.0.1" },
 ]
 provides-extras = ["zen", "bench", "tensorizer", "fastsafetensors", "instanttensor", "runai", "audio", "video", "deepstream", "flashinfer", "b12x", "helion", "grpc", "otel", "extra-quant"]
 
@@ -8761,7 +8765,7 @@ dependencies = [
     { name = "apache-tvm-ffi" },
     { name = "numpy" },
     { name = "pydantic" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
+    { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
     { name = "transformers" },
     { name = "triton", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "typing-extensions" },
@@ -8787,22 +8791,22 @@ wheels = [
 
 [[package]]
 name = "yarl"
-version = "1.24.2"
+version = "1.25.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "idna" },
     { name = "multidict" },
     { name = "propcache" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/79/12/1e8f37460ea0f7eb59c221fdaf0ed75e7ac43e97f8093b9c6f411df50a78/yarl-1.24.2.tar.gz", hash = "sha256:9ac374123c6fd7abf64d1fec93962b0bd4ee2c19751755a762a72dd96c0378f8", size = 210798, upload-time = "2026-05-19T21:31:05.599Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/75/16/e8be8e2fb175bbf41a0680381a319f1199fae256588241a2ac8677eafb49/yarl-1.25.1.tar.gz", hash = "sha256:03dd38de09bc213e9a8b29761eec33ee1d5318dac0e49d8af36e4d27830e23a7", size = 246245, upload-time = "2026-09-15T19:35:02.264Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/f0/da/866bcb01076ba49d2b42b309867bed3826421f1c479655eb7a607b44f20b/yarl-1.24.2-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:b975866c184564c827e0877380f0dae57dcca7e52782128381b72feff6dfceb8", size = 129957, upload-time = "2026-05-19T21:28:51.695Z" },
-    { url = "https://files.pythonhosted.org/packages/29/b6/170e2b8d4e3bc30e6bfdcca53556537f5bf595e938632dfcb059311f3ff6/yarl-1.24.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8ae44649b00947634ab0dab2a374a638f52923a6e67083f2c156cd5cbd1a881d", size = 91688, upload-time = "2026-05-19T21:28:54.865Z" },
-    { url = "https://files.pythonhosted.org/packages/fe/a5/c9f655d5553ea0b99fdac9d6a99ad3f9b3e73b8e5758bb46f58c9831f74c/yarl-1.24.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:507cc19f0b45454e2d6dcd62ff7d062b9f77a2812404e62dbdaec05b50faa035", size = 102902, upload-time = "2026-05-19T21:28:56.963Z" },
-    { url = "https://files.pythonhosted.org/packages/7a/da/323a01c349bd5fb01bb6652e314d9bb218cee630a736bdb810ad50e4013f/yarl-1.24.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7e7ebcdef69dec6c6451e616f32b622a6d4a2e92b445c992f7c8e5274a6bbc4c", size = 105612, upload-time = "2026-05-19T21:29:04.247Z" },
-    { url = "https://files.pythonhosted.org/packages/41/07/efabe5df87e96d7ad5959760b888344be48cd6884db127b407c6b5503adc/yarl-1.24.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3065657c80a2321225e804048597ad55658a7e76b32d6f5ee4074d04c50401db", size = 102333, upload-time = "2026-05-19T21:29:08.267Z" },
-    { url = "https://files.pythonhosted.org/packages/48/2d/1c8d89c7c5f9cad9fb2902445d94e2ab1d7aa35de029afbb8ae95c42d00f/yarl-1.24.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:e30dd55825dc554ec5b66a94953b8eda8745926514c5089dfcacecb9c99b5bd1", size = 105719, upload-time = "2026-05-19T21:29:18.367Z" },
-    { url = "https://files.pythonhosted.org/packages/fd/4d/4b880086bd0d3e034d25647be1d830afc3e3f610e98c4ab3490af6b1b6d5/yarl-1.24.2-py3-none-any.whl", hash = "sha256:2783d9226db8797636cd6896e4de81feed252d1db72265686c9558d97a4d94b9", size = 53576, upload-time = "2026-05-19T21:31:03.909Z" },
+    { url = "https://files.pythonhosted.org/packages/75/b3/cd32ac66ae622b854c2df0ac52106dda220d361b65a64fde7d5b3684aa3f/yarl-1.25.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:94d7aa6debf92a1dd14cb5280b083a764169a13cfb23a452111160274ed989f4", size = 144798, upload-time = "2026-09-15T19:31:01.821Z" },
+    { url = "https://files.pythonhosted.org/packages/be/dd/ee38aec8e09fdf957e50d4085453fbe202f56c6c3b4cf07b81cdb4f09ee9/yarl-1.25.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:e029648f9c951db30e98a7d7ec90835db88ec4b32820efe2a9bdc2287e032eb6", size = 104325, upload-time = "2026-09-15T19:31:06.338Z" },
+    { url = "https://files.pythonhosted.org/packages/1e/b3/058dbfb1857b484c9cf9cc135659f50b85ce66e03c99e44dc2f7b6161f55/yarl-1.25.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4d781294bb815ecb5ea57ff6bbf8038e0a31a95fdf3e1788f66e0dc100d64b58", size = 115358, upload-time = "2026-09-15T19:31:08.593Z" },
+    { url = "https://files.pythonhosted.org/packages/bc/7b/ca212cbe170ac8b96e45317ecbcf9c3c3ecf0cdec98d5b088a9c4088929b/yarl-1.25.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c6f117789d22dce188e5754e8bc65b7e6ebf8cb73963b9fa761f672a5883769d", size = 117050, upload-time = "2026-09-15T19:31:17.241Z" },
+    { url = "https://files.pythonhosted.org/packages/e8/43/268717870f9ba0cc9701a95181587f6dc8c5f387aab4aeecc83158f38a79/yarl-1.25.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e80f557716fd765439577131e526b8942ffc2c07bdbc5e39fa62f660ba1e963f", size = 114944, upload-time = "2026-09-15T19:31:21.414Z" },
+    { url = "https://files.pythonhosted.org/packages/05/79/ad94f93ca731bc9e44d321833ab96b82a4f9f5f63cf773f81a4aeea5ecc1/yarl-1.25.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:68782fdb4027b8d1eee25ec35e9a6db05e863b899eb0310b3a33b6c3fef55707", size = 117060, upload-time = "2026-09-15T19:31:31.367Z" },
+    { url = "https://files.pythonhosted.org/packages/54/22/318c7980066769c6bcd9221ed2248294f5698811da099013098c670565ed/yarl-1.25.1-py3-none-any.whl", hash = "sha256:681c758b0490f9e96b78e5fa8e8dc6e648e9185bb6eaebe73183c33ea0c445f3", size = 63617, upload-time = "2026-09-15T19:34:59.616Z" },
 ]
 
 [[package]]

From 3307dc2b5f93bd33781cc26d7dc2a5f7beed4a77 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Mon, 2 Nov 2026 19:42:24 +0000
Subject: [PATCH 76/88] refactor(inference): drop the LoRA load_inplace reset
 vLLM 0.33 does itself

vllm-project/vllm#41482 resets load_inplace before caching the adapter.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/prime_rl/inference/vllm/server.py | 16 +++-------------
 1 file changed, 3 insertions(+), 13 deletions(-)

diff --git a/src/prime_rl/inference/vllm/server.py b/src/prime_rl/inference/vllm/server.py
index 00a8e0df6dd..e3d12b6fad3 100644
--- a/src/prime_rl/inference/vllm/server.py
+++ b/src/prime_rl/inference/vllm/server.py
@@ -83,23 +83,13 @@ async def load_lora_adapter(lora_request: LoadLoRAAdapterRequest, raw_request: R
     the base model is served as ``<model>-base``; see ``InferenceConfig.to_namespace``.)
     vLLM's native loader rejects a same-name reload unless ``load_inplace=True``, so
     we force it here — that makes the worker re-read the new weights during
-    ``add_lora``, reusing the existing ``lora_int_id``.
-
-    We then reset the stored request's flag back to ``False``. ``load_inplace`` is
-    a sticky field on the ``LoRARequest`` that ``_maybe_get_adapters`` hands to
-    every generation request; left ``True`` it would force a disk reload on each
-    scheduler step. The orchestrator awaits this endpoint before dispatching
-    rollouts for the new version, so the reset always lands before generation.
-    The reset runs regardless of success/error: vLLM only stores the adapter on
-    success today, but resetting whatever is stored keeps us correct even if a
-    future version were to leave a ``load_inplace=True`` request behind on error.
+    ``add_lora``, reusing the existing ``lora_int_id``. vLLM resets the stored
+    request's flag to ``False`` before caching it (vllm-project/vllm#41482), so it
+    does not force a disk reload on every later generation request.
     """
     handler = models(raw_request)
     lora_request.load_inplace = True
     response = await handler.load_lora_adapter(lora_request)
-    stored = handler.lora_requests.get(lora_request.lora_name)
-    if stored is not None:
-        stored.load_inplace = False
     if isinstance(response, ErrorResponse):
         return JSONResponse(content=response.model_dump(), status_code=response.error.code)
     await engine_client(raw_request).update_weight_version(weight_version)

From 95e9768806205182c70c5d33a41eb93bc6847363 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Mon, 2 Nov 2026 20:15:40 +0000
Subject: [PATCH 77/88] chore: pin prime-kernels, deep-ep, deep-gemm,
 deep-select built for torch 2.14

From PrimeIntellect-ai/prime-kernels release v0.1.0-b050074 (PrimeIntellect-ai/prime-kernels#24).
torchao stays on v0.1.1: it only uses the stable libtorch ABI.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 pyproject.toml | 16 ++++-----
 uv.lock        | 88 +++++++++++++++++++++++++-------------------------
 2 files changed, 52 insertions(+), 52 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index f56e284f0e3..1c010427f60 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -287,20 +287,20 @@ vllm = [
     { url = "https://github.com/vllm-project/vllm/releases/download/v0.33.0rc3/vllm-0.33.0-cp38-abi3-manylinux_2_28_aarch64.whl", marker = "platform_machine == 'aarch64'" },
 ]
 deep-ep = [
-    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl", marker = "platform_machine == 'x86_64'" },
-    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_aarch64.whl", marker = "platform_machine == 'aarch64'" },
+    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl", marker = "platform_machine == 'x86_64'" },
+    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_aarch64.whl", marker = "platform_machine == 'aarch64'" },
 ]
 deep-gemm = [
-    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_x86_64.whl", marker = "platform_machine == 'x86_64'" },
-    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_aarch64.whl", marker = "platform_machine == 'aarch64'" },
+    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_gemm-2.5.0+local-cp312-cp312-linux_x86_64.whl", marker = "platform_machine == 'x86_64'" },
+    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_gemm-2.5.0+local-cp312-cp312-linux_aarch64.whl", marker = "platform_machine == 'aarch64'" },
 ]
 prime-kernels = [
-    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/prime_kernels-0.1.0+cu130torch2.13.0-cp312-cp312-linux_x86_64.whl", marker = "platform_machine == 'x86_64'" },
-    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/prime_kernels-0.1.0+cu130torch2.13.0-cp312-cp312-linux_aarch64.whl", marker = "platform_machine == 'aarch64'" },
+    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/prime_kernels-0.1.0+cu130torch2.14.0-cp312-cp312-linux_x86_64.whl", marker = "platform_machine == 'x86_64'" },
+    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/prime_kernels-0.1.0+cu130torch2.14.0-cp312-cp312-linux_aarch64.whl", marker = "platform_machine == 'aarch64'" },
 ]
 deep-select = [
-    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121630-cp312-cp312-linux_x86_64.whl", marker = "platform_machine == 'x86_64'" },
-    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121543-cp312-cp312-linux_aarch64.whl", marker = "platform_machine == 'aarch64'" },
+    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_select-1.0.0+bfa4507.20261102.200747-cp312-cp312-linux_x86_64.whl", marker = "platform_machine == 'x86_64'" },
+    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_select-1.0.0+bfa4507.20261102.200624-cp312-cp312-linux_aarch64.whl", marker = "platform_machine == 'aarch64'" },
 ]
 flash-linear-attention = { git = "https://github.com/fla-org/flash-linear-attention" }
 mamba-ssm = { git = "https://github.com/state-spaces/mamba", rev = "f577286" }
diff --git a/uv.lock b/uv.lock
index 62221e7a921..10c27f2bb4d 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1305,67 +1305,67 @@ wheels = [
 [[package]]
 name = "deep-ep"
 version = "1.2.1+29d31c0"
-source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_aarch64.whl" }
+source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_aarch64.whl" }
 resolution-markers = [
     "platform_machine == 'aarch64' and sys_platform == 'linux'",
 ]
 wheels = [
-    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_aarch64.whl", hash = "sha256:ead61bea6a830e12a743510ab9633a71edf2127ec51a2b4d783b34b253cf6737" },
+    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_aarch64.whl", hash = "sha256:f369f166368e79ef51fa9aa862c31b59e321d3d4775726c0cac4d16943e2a18f" },
 ]
 
 [[package]]
 name = "deep-ep"
 version = "1.2.1+29d31c0"
-source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl" }
+source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl" }
 resolution-markers = [
     "platform_machine == 'x86_64' and sys_platform == 'linux'",
 ]
 wheels = [
-    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl", hash = "sha256:808f93c0863d942fef8a098891b09e7a733ba83e56c5df301b7053dd8e0eb631" },
+    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl", hash = "sha256:6559aaf2f00177c894664e15f4989c031d3d6eaa692db0ecd657a87b90385253" },
 ]
 
 [[package]]
 name = "deep-gemm"
-version = "2.5.0+891d57b"
-source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_aarch64.whl" }
+version = "2.5.0+local"
+source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_gemm-2.5.0+local-cp312-cp312-linux_aarch64.whl" }
 resolution-markers = [
     "platform_machine == 'aarch64' and sys_platform == 'linux'",
 ]
 wheels = [
-    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_aarch64.whl", hash = "sha256:657b4ef2ba3e7449216b0ea8360f6fd0aab400e1503ccd609f20c2fd5cdcee66" },
+    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_gemm-2.5.0+local-cp312-cp312-linux_aarch64.whl", hash = "sha256:e1da2fda91d184eb342d94fb8c6c4768869ec3f2fa508ebbfee7333e4e554210" },
 ]
 
 [[package]]
 name = "deep-gemm"
-version = "2.5.0+891d57b"
-source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_x86_64.whl" }
+version = "2.5.0+local"
+source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_gemm-2.5.0+local-cp312-cp312-linux_x86_64.whl" }
 resolution-markers = [
     "platform_machine == 'x86_64' and sys_platform == 'linux'",
 ]
 wheels = [
-    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_x86_64.whl", hash = "sha256:77b8f9ff3a5c2ef3b6621cd310aadc161115d4cbaf4aeace722dfe7fe40a2252" },
+    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_gemm-2.5.0+local-cp312-cp312-linux_x86_64.whl", hash = "sha256:f4e3a901c0eb22767d100f0e918710552dbe84300067a48d16145bae7dd44ee7" },
 ]
 
 [[package]]
 name = "deep-select"
-version = "1.0.0+bfa4507.20261026.121543"
-source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121543-cp312-cp312-linux_aarch64.whl" }
+version = "1.0.0+bfa4507.20261102.200624"
+source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_select-1.0.0+bfa4507.20261102.200624-cp312-cp312-linux_aarch64.whl" }
 resolution-markers = [
     "platform_machine == 'aarch64' and sys_platform == 'linux'",
 ]
 wheels = [
-    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121543-cp312-cp312-linux_aarch64.whl", hash = "sha256:d16b315567a2d9ec42cc822e64c5d72ca9786085943b1d8cd288fb601d030c7b" },
+    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_select-1.0.0+bfa4507.20261102.200624-cp312-cp312-linux_aarch64.whl", hash = "sha256:e3af71e86198341baea7e88e2f3078327601b37e4380957649e88be9cff2dd8e" },
 ]
 
 [[package]]
 name = "deep-select"
-version = "1.0.0+bfa4507.20261026.121630"
-source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121630-cp312-cp312-linux_x86_64.whl" }
+version = "1.0.0+bfa4507.20261102.200747"
+source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_select-1.0.0+bfa4507.20261102.200747-cp312-cp312-linux_x86_64.whl" }
 resolution-markers = [
     "platform_machine == 'x86_64' and sys_platform == 'linux'",
 ]
 wheels = [
-    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121630-cp312-cp312-linux_x86_64.whl", hash = "sha256:bd4a577996695b6e78dbb010339951a9e980faecad3075b854debec448da7267" },
+    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_select-1.0.0+bfa4507.20261102.200747-cp312-cp312-linux_x86_64.whl", hash = "sha256:addffb35be12eda107605c202dc62fb1174ed577083c6eb7ba6c20b4cc3de342" },
 ]
 
 [[package]]
@@ -5134,8 +5134,8 @@ wheels = [
 
 [[package]]
 name = "prime-kernels"
-version = "0.1.0+cu130torch2.13.0"
-source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/prime_kernels-0.1.0+cu130torch2.13.0-cp312-cp312-linux_aarch64.whl" }
+version = "0.1.0+cu130torch2.14.0"
+source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/prime_kernels-0.1.0+cu130torch2.14.0-cp312-cp312-linux_aarch64.whl" }
 resolution-markers = [
     "platform_machine == 'aarch64' and sys_platform == 'linux'",
 ]
@@ -5143,7 +5143,7 @@ dependencies = [
     { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
 ]
 wheels = [
-    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/prime_kernels-0.1.0+cu130torch2.13.0-cp312-cp312-linux_aarch64.whl", hash = "sha256:695d61bd1b9ba99969ca7ac1410f5a35c1d1c747dab13a92c6f9c3a8def9b41a" },
+    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/prime_kernels-0.1.0+cu130torch2.14.0-cp312-cp312-linux_aarch64.whl", hash = "sha256:5d356d2a544de52fa50411e2a32b276c11f102440b29a759ac333e12559f7503" },
 ]
 
 [package.metadata]
@@ -5151,8 +5151,8 @@ requires-dist = [{ name = "torch", specifier = ">=2.9.0" }]
 
 [[package]]
 name = "prime-kernels"
-version = "0.1.0+cu130torch2.13.0"
-source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/prime_kernels-0.1.0+cu130torch2.13.0-cp312-cp312-linux_x86_64.whl" }
+version = "0.1.0+cu130torch2.14.0"
+source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/prime_kernels-0.1.0+cu130torch2.14.0-cp312-cp312-linux_x86_64.whl" }
 resolution-markers = [
     "platform_machine == 'x86_64' and sys_platform == 'linux'",
 ]
@@ -5160,7 +5160,7 @@ dependencies = [
     { name = "torch", version = "2.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
 ]
 wheels = [
-    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/prime_kernels-0.1.0+cu130torch2.13.0-cp312-cp312-linux_x86_64.whl", hash = "sha256:c181b28773116970f08c130d2576ea89d81174820b2cfa03006dcfed4c23d462" },
+    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/prime_kernels-0.1.0+cu130torch2.14.0-cp312-cp312-linux_x86_64.whl", hash = "sha256:7bb66fb521142bc04ea674c3fb81a0c7130027fcad728c287ba22972ca9c5230" },
 ]
 
 [package.metadata]
@@ -5235,12 +5235,12 @@ dependencies = [
 
 [package.optional-dependencies]
 all = [
-    { name = "deep-ep", version = "1.2.1+29d31c0", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
-    { name = "deep-ep", version = "1.2.1+29d31c0", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "deep-gemm", version = "2.5.0+891d57b", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
-    { name = "deep-gemm", version = "2.5.0+891d57b", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "deep-select", version = "1.0.0+bfa4507.20261026.121543", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121543-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
-    { name = "deep-select", version = "1.0.0+bfa4507.20261026.121630", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121630-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+    { name = "deep-ep", version = "1.2.1+29d31c0", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
+    { name = "deep-ep", version = "1.2.1+29d31c0", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+    { name = "deep-gemm", version = "2.5.0+local", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_gemm-2.5.0+local-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
+    { name = "deep-gemm", version = "2.5.0+local", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_gemm-2.5.0+local-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+    { name = "deep-select", version = "1.0.0+bfa4507.20261102.200624", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_select-1.0.0+bfa4507.20261102.200624-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
+    { name = "deep-select", version = "1.0.0+bfa4507.20261102.200747", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_select-1.0.0+bfa4507.20261102.200747-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "dion", marker = "sys_platform == 'linux'" },
     { name = "flash-attn", version = "2.8.3+cu130torch2.14", source = { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.10.0/flash_attn-2.8.3+cu130torch2.14-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "flash-attn", version = "2.8.3+cu130torch2.14", source = { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.10.1/flash_attn-2.8.3+cu130torch2.14-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
@@ -5278,10 +5278,10 @@ dashboard = [
     { name = "uvicorn" },
 ]
 disagg = [
-    { name = "deep-ep", version = "1.2.1+29d31c0", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
-    { name = "deep-ep", version = "1.2.1+29d31c0", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "deep-gemm", version = "2.5.0+891d57b", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
-    { name = "deep-gemm", version = "2.5.0+891d57b", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+    { name = "deep-ep", version = "1.2.1+29d31c0", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
+    { name = "deep-ep", version = "1.2.1+29d31c0", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+    { name = "deep-gemm", version = "2.5.0+local", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_gemm-2.5.0+local-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
+    { name = "deep-gemm", version = "2.5.0+local", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_gemm-2.5.0+local-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "nixl", marker = "sys_platform == 'linux'" },
     { name = "vllm-router", version = "0.2.3", source = { url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.2.3/vllm_router-0.2.3-cp38-abi3-manylinux_2_28_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
     { name = "vllm-router", version = "0.2.3", source = { url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.2.3/vllm_router-0.2.3-cp38-abi3-manylinux_2_28_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
@@ -5300,8 +5300,8 @@ flash-attn-cute = [
     { name = "flash-attn-4", marker = "sys_platform == 'linux'" },
 ]
 gpu = [
-    { name = "deep-select", version = "1.0.0+bfa4507.20261026.121543", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121543-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
-    { name = "deep-select", version = "1.0.0+bfa4507.20261026.121630", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121630-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+    { name = "deep-select", version = "1.0.0+bfa4507.20261102.200624", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_select-1.0.0+bfa4507.20261102.200624-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
+    { name = "deep-select", version = "1.0.0+bfa4507.20261102.200747", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_select-1.0.0+bfa4507.20261102.200747-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "dion", marker = "sys_platform == 'linux'" },
     { name = "flash-linear-attention", marker = "sys_platform == 'linux'" },
     { name = "flashinfer-cubin", marker = "sys_platform == 'linux'" },
@@ -5326,8 +5326,8 @@ gpu = [
     { name = "vllm", version = "0.33.0", source = { url = "https://github.com/vllm-project/vllm/releases/download/v0.33.0rc3/vllm-0.33.0-cp38-abi3-manylinux_2_28_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
 ]
 kernels = [
-    { name = "prime-kernels", version = "0.1.0+cu130torch2.13.0", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/prime_kernels-0.1.0+cu130torch2.13.0-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
-    { name = "prime-kernels", version = "0.1.0+cu130torch2.13.0", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/prime_kernels-0.1.0+cu130torch2.13.0-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+    { name = "prime-kernels", version = "0.1.0+cu130torch2.14.0", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/prime_kernels-0.1.0+cu130torch2.14.0-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
+    { name = "prime-kernels", version = "0.1.0+cu130torch2.14.0", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/prime_kernels-0.1.0+cu130torch2.14.0-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
 ]
 quack = [
     { name = "quack-kernels", marker = "sys_platform == 'linux'" },
@@ -5353,15 +5353,15 @@ requires-dist = [
     { name = "aiolimiter", specifier = ">=1.2.1" },
     { name = "beartype", specifier = ">=0.21.0" },
     { name = "datasets", specifier = ">=4.0.0" },
-    { name = "deep-ep", marker = "platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_aarch64.whl" },
+    { name = "deep-ep", marker = "platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_aarch64.whl" },
     { name = "deep-ep", marker = "platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux' and extra == 'disagg'" },
-    { name = "deep-ep", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl" },
-    { name = "deep-gemm", marker = "platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_aarch64.whl" },
+    { name = "deep-ep", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl" },
+    { name = "deep-gemm", marker = "platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_gemm-2.5.0+local-cp312-cp312-linux_aarch64.whl" },
     { name = "deep-gemm", marker = "platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux' and extra == 'disagg'" },
-    { name = "deep-gemm", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_x86_64.whl" },
-    { name = "deep-select", marker = "platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'gpu'", url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121543-cp312-cp312-linux_aarch64.whl" },
+    { name = "deep-gemm", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_gemm-2.5.0+local-cp312-cp312-linux_x86_64.whl" },
+    { name = "deep-select", marker = "platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'gpu'", url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_select-1.0.0+bfa4507.20261102.200624-cp312-cp312-linux_aarch64.whl" },
     { name = "deep-select", marker = "platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux' and extra == 'gpu'" },
-    { name = "deep-select", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'gpu'", url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121630-cp312-cp312-linux_x86_64.whl" },
+    { name = "deep-select", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'gpu'", url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/deep_select-1.0.0+bfa4507.20261102.200747-cp312-cp312-linux_x86_64.whl" },
     { name = "dion", marker = "sys_platform == 'linux' and extra == 'gpu'", git = "https://github.com/samsja/dion.git?rev=0c801b0" },
     { name = "fastapi", specifier = ">=0.115" },
     { name = "fastapi", marker = "extra == 'dashboard'", specifier = ">=0.115" },
@@ -5388,9 +5388,9 @@ requires-dist = [
     { name = "openai", specifier = ">=1.106.1" },
     { name = "orjson", specifier = ">=3.11.0" },
     { name = "prime", specifier = ">=0.18.0" },
-    { name = "prime-kernels", marker = "platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'kernels'", url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/prime_kernels-0.1.0+cu130torch2.13.0-cp312-cp312-linux_aarch64.whl" },
+    { name = "prime-kernels", marker = "platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'kernels'", url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/prime_kernels-0.1.0+cu130torch2.14.0-cp312-cp312-linux_aarch64.whl" },
     { name = "prime-kernels", marker = "platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux' and extra == 'kernels'" },
-    { name = "prime-kernels", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'kernels'", url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/prime_kernels-0.1.0+cu130torch2.13.0-cp312-cp312-linux_x86_64.whl" },
+    { name = "prime-kernels", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'kernels'", url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b050074/prime_kernels-0.1.0+cu130torch2.14.0-cp312-cp312-linux_x86_64.whl" },
     { name = "prime-pydantic-config", editable = "deps/pydantic-config" },
     { name = "prime-rl", extras = ["disagg"], marker = "extra == 'all'" },
     { name = "prime-rl", extras = ["flash-attn"], marker = "extra == 'all'" },

From 27ce941cf09623b5982f213a57bd6f5eff35006e Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Mon, 2 Nov 2026 20:22:58 +0000
Subject: [PATCH 78/88] fix(inference): keep KV offload reuse with routed
 experts on vLLM 0.33

vllm-project/vllm#62591 admits an external hit only when every block's routing
rows are still retained, which it does for at most 2x the GPU KV cache. With a
CPU offload tier every multi-turn reload was refused: 0 external hit tokens per
turn on a 48-conversation x 4-turn OLMoE run (RTX PRO 6000), vs 13-25k with
routed experts off. A load within the routed-experts clamp needs none of its rows,
so admit it whole; with the fix the external hits match the routed-experts-off
run exactly, every row count is right and generated tokens are identical.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/prime_rl/inference/patches.py             | 17 +++++++++++
 .../test_routed_experts_kv_offload.py         | 29 +++++++++++++++++++
 .../unit/inference/test_routed_experts_pd.py  | 10 ++++++-
 3 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/src/prime_rl/inference/patches.py b/src/prime_rl/inference/patches.py
index 31157d5fed7..6f49466074a 100644
--- a/src/prime_rl/inference/patches.py
+++ b/src/prime_rl/inference/patches.py
@@ -493,7 +493,15 @@ def monkey_patch_clamp_kv_offload_loads_for_routed_experts():
 
     Multi-turn reloads stay below the next turn's ``routed_experts_prompt_start`` and keep their
     hits; offloaded prefixes past it (first turns, shared prompts) are recomputed.
+
+    vLLM 0.33 (vllm-project/vllm#62591) also admits an external hit only when this instance still retains the
+    routing rows of every block in it, and it retains rows for at most twice the GPU KV cache's blocks. A CPU
+    offload tier holds far more than that, so every multi-turn reload of an evicted prefix was refused and
+    recomputed. A load within the clamp needs none of its rows, so ``cap_external_hit`` admits it whole; the
+    rest still goes through upstream's check. Blocks loaded this way carry no rows, which the local prefix-hit
+    cut (`monkey_patch_routed_experts_cut_loaded_prefix_hits`) keeps later requests from reading.
     """
+    from vllm.distributed.aux_output_connector.connector import AuxOutputSchedulerConnector
     from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.connector import MooncakeStoreConnector
     from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import OffloadingConnector
 
@@ -530,9 +538,18 @@ def _get_num_new_matched_tokens(self, request, num_computed_tokens):
         load_spec.tail_key_boundaries = ()
         return limit, load_async
 
+    original_cap_external_hit = AuxOutputSchedulerConnector.cap_external_hit
+
+    def cap_external_hit(self, request, num_local, num_external):
+        limit = _routed_experts_load_limit(request, num_local, self._block_size)
+        if limit is not None and num_external <= limit:
+            return num_external
+        return original_cap_external_hit(self, request, num_local, num_external)
+
     _max_loadable_tokens._prime_rl_routed_experts_clamp = True
     OffloadingConnector._max_loadable_tokens = _max_loadable_tokens
     MooncakeStoreConnector.get_num_new_matched_tokens = _get_num_new_matched_tokens
+    AuxOutputSchedulerConnector.cap_external_hit = cap_external_hit
 
 
 def _track_externally_loaded_blocks(kv_cache_manager, gid: int, block_size: int) -> set[int]:
diff --git a/tests/unit/inference/test_routed_experts_kv_offload.py b/tests/unit/inference/test_routed_experts_kv_offload.py
index 5e93b085033..051167062b4 100644
--- a/tests/unit/inference/test_routed_experts_kv_offload.py
+++ b/tests/unit/inference/test_routed_experts_kv_offload.py
@@ -99,3 +99,32 @@ def test_cut_patch_takes_upstream_max_cache_hit_length():
     assert num_local == 32
     _, num_local, _, _ = Scheduler._get_local_prefix_cache_hit(scheduler, b)
     assert num_local == 64
+
+
+def test_offload_loads_below_prompt_start_need_no_retained_rows():
+    import torch
+    from vllm.distributed.aux_output_connector.connector import AuxOutputSchedulerConnector
+    from vllm.v1.core.kv_cache_manager import KVCacheManager
+    from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
+
+    from prime_rl.inference.patches import monkey_patch_clamp_kv_offload_loads_for_routed_experts
+
+    monkey_patch_clamp_kv_offload_loads_for_routed_experts()
+    spec = FullAttentionSpec(block_size=16, num_kv_heads=1, head_size=8, dtype=torch.float16)
+    config = KVCacheConfig(num_blocks=12, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["l0"], spec)])
+    manager = KVCacheManager(config, max_model_len=1024, scheduler_block_size=16, hash_block_size=16)
+    connector = AuxOutputSchedulerConnector(manager, kv_connector=object())
+    request = SimpleNamespace(
+        request_id="r",
+        block_hashes=[],
+        num_tokens=100,
+        num_output_tokens=0,
+        kv_transfer_params=None,
+        sampling_params=SimpleNamespace(routed_experts_prompt_start=70),
+    )
+    # The next turn needs rows from token 70 on, so a load of the first 64 tokens needs none: admitted
+    # although this instance retains no rows for those blocks.
+    assert connector.cap_external_hit(request, 0, 64) == 64
+    assert connector.cap_external_hit(request, 16, 48) == 48
+    # Past the prompt-start block, upstream's retained-rows check still decides (no rows -> no hit).
+    assert connector.cap_external_hit(request, 0, 80) == 0
diff --git a/tests/unit/inference/test_routed_experts_pd.py b/tests/unit/inference/test_routed_experts_pd.py
index 866e948d4a8..dbae0405b8c 100644
--- a/tests/unit/inference/test_routed_experts_pd.py
+++ b/tests/unit/inference/test_routed_experts_pd.py
@@ -57,7 +57,15 @@ def test_pd_decode_requests_skip_retained_row_caps():
     connector = AuxOutputSchedulerConnector(_kv_cache_manager(), kv_connector=object())
 
     def request(kv_transfer_params):
-        return SimpleNamespace(request_id="r", block_hashes=[], num_tokens=10, kv_transfer_params=kv_transfer_params)
+        return SimpleNamespace(
+            request_id="r",
+            block_hashes=[],
+            num_tokens=10,
+            num_output_tokens=0,
+            kv_transfer_params=kv_transfer_params,
+            # Rows from token 0 on: no load is free of the retained-rows check.
+            sampling_params=SimpleNamespace(routed_experts_prompt_start=0),
+        )
 
     # A decode instance has no rows for a remotely prefilled prompt: the NIXL load must not be zeroed.
     decode = request({"do_remote_prefill": True, "remote_engine_id": "p0"})

From 0ea767d4fb56edb07ab0f9886ab519820139d9be Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Tue, 3 Nov 2026 00:00:51 +0000
Subject: [PATCH 79/88] fix(weights): one update plan per TP group;
 replica-aware, incremental checksums

- TP rank 0 plans each update (checkpoint to load, deltas to apply, full or
  incremental verify) and the others follow it, so workers can no longer split
  between apply, replay and anchor load (an anchor completing mid-update, or a
  failure on one worker) and issue different collectives.
- A failure on any worker fails the update on all of them (gloo flag), so they
  all reload and replay the same way next time.
- One fixed-size check over every checkpoint tensor of the run: every worker
  runs the same reductions whatever it holds. A tensor is checked against
  multiplicity x expected when its holders cover it evenly, so KV heads
  replicated over TP workers (TP > num KV heads) no longer fail every load.
- Checksums move with the changes in apply_delta (O(changes) instead of a pass
  over every param); steps that carry an anchor recompute them in full.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../inference/vllm/worker/filesystem.py       | 234 +++++++++++-------
 .../inference/vllm/worker/weight_transfer.py  |  73 ++++--
 .../inference/test_checkpoint_index_map.py    |  21 ++
 3 files changed, 221 insertions(+), 107 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index a272988ee63..85f3d9d9cad 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -164,30 +164,99 @@ def prefetch_weights_done(self) -> bool:
         return True
 
     def update_weights_from_path(self, weight_path: str) -> None:
-        """Move to the broadcast step in ``weight_path`` (a full HF checkpoint and/or a sparse delta)."""
+        """Move to the broadcast step in ``weight_path`` (a full HF checkpoint and/or a sparse delta).
+
+        The TP group's workers run collectives while they verify, so they must take the same steps: TP rank 0
+        plans the update (which checkpoint to load, which deltas to apply) and the others follow its plan. A
+        failure on any worker fails the update on all of them, so they all replay the same way next time."""
         step_dir = Path(weight_path)
         step = int(step_dir.name.removeprefix("step_"))
+        tp_group = get_tp_group()
+        plan, error = None, None
+        if tp_group.rank_in_group == 0:
+            try:
+                plan = self._plan_update(step_dir)
+            except Exception as e:  # re-raised below, after the others got the failure instead of a plan
+                error = e
+        if tp_group.world_size > 1:
+            plan = tp_group.broadcast_object(plan if error is None else {"failed": repr(error)}, src=0)
+        if error is not None:
+            raise error
+        if "failed" in plan:
+            raise RuntimeError(f"TP rank 0 could not plan the update to {step_dir.name}: {plan['failed']}")
+        # Unknown until this update succeeds: after a failure, the next update reloads and replays
+        self.weight_step = -1
+        check = None
+        try:
+            if plan["index_map"] is not None:
+                self._build_index_map(Path(plan["index_map"]))
+                if plan["write_layout"]:
+                    # A full broadcast: the trainer takes the layout at its first one, also after a restart
+                    self._write_layout(step_dir.parent)
+            check = self._run_plan(plan)
+        except Exception as e:  # re-raised below, once every worker of the TP group knows
+            error = e
+        self._raise_if_any_failed(error, f"update to {step_dir.name}")
+        if check is not None:
+            self._check(*check)
+        self.weight_step = step
+
+    def _plan_update(self, step_dir: Path) -> dict:
+        """What every TP worker does to move to ``step_dir``: a dict of plain values, broadcast from TP rank 0."""
         has_delta = (step_dir / DELTA_DIR / MANIFEST).exists()
         base_step = read_manifest(step_dir)["base_step"] if has_delta else None
-        if has_delta or anchor_complete(step_dir):
-            self._build_index_map(step_dir)
-            if base_step is None:
-                # A full broadcast: the trainer takes the layout at its first one, also after a restart
-                self._write_layout(step_dir.parent)
-        # Unknown until this update succeeds: after a failure, the next update reloads and replays
-        held_step, self.weight_step = self.weight_step, -1
+        plan = {
+            "index_map": step_dir.as_posix() if has_delta or anchor_complete(step_dir) else None,
+            "write_layout": base_step is None,
+            "load": None,
+            "deltas": [],
+            # Steps that carry an anchor recompute every checksum from the params; others update them with the changes
+            "full_verify": any(step_dir.glob("anchor_rank*")),
+        }
         if not has_delta:
-            self._load_checkpoint(step_dir)
+            plan["load"] = step_dir.as_posix()
+        elif base_step is None:
+            # The run's first broadcast: the weights the engine booted with
+            if self.weight_step is not None:
+                plan["load"] = "boot"
+        elif base_step == self.weight_step:
+            plan["deltas"] = [step_dir.as_posix()]
         else:
-            if base_step is None:
-                # The run's first broadcast: the weights the engine booted with
-                if held_step is not None:
-                    self._load_checkpoint(self._boot_dir())
-            elif base_step == held_step:
-                self._apply_delta(step_dir)
-            else:
-                self._replay(step_dir)
-        self.weight_step = step
+            # Load the latest complete anchor at or before ``step_dir`` (or the boot checkpoint, for a chain from the
+            # run's first broadcast) and replay the deltas since
+            chain, source = [], step_dir
+            while not anchor_complete(source):
+                source_base = read_manifest(source)["base_step"]
+                if source_base is None:
+                    break
+                chain.append(source)
+                source = source.parent / f"step_{source_base}"
+            plan["load"] = source.as_posix() if anchor_complete(source) else "boot"
+            plan["deltas"] = [delta_dir.as_posix() for delta_dir in reversed(chain)]
+            logger.info(f"Replaying {len(chain)} sparse weight deltas on top of {Path(plan['load']).name}")
+        return plan
+
+    def _run_plan(self, plan: dict) -> tuple | None:
+        """Load and apply what ``plan`` says; returns the check of the final weights for ``_check``."""
+        check = None
+        if plan["load"] is not None:
+            check = self._load_checkpoint(self._boot_dir() if plan["load"] == "boot" else Path(plan["load"]))
+        for delta_dir in plan["deltas"]:
+            checksums = self._apply_delta(Path(delta_dir))
+            # Each delta's checksums cover every tensor, so the last one checks the whole chain
+            check = (Path(delta_dir).name, *self.index_map.checksums(full=plan["full_verify"]), checksums)
+        return check
+
+    def _raise_if_any_failed(self, error: Exception | None, what: str) -> None:
+        """Raise on every worker of the TP group if ``error`` is set on any of them."""
+        tp_group = get_tp_group()
+        if tp_group.world_size > 1:
+            failed = torch.tensor([error is not None], dtype=torch.int32)
+            torch.distributed.all_reduce(failed, op=torch.distributed.ReduceOp.MAX, group=tp_group.cpu_group)
+            if failed.item() and error is None:
+                raise RuntimeError(f"A TP peer failed the {what}")
+        if error is not None:
+            raise error
 
     def _boot_dir(self) -> Path:
         """The local dir of the checkpoint the engine booted from."""
@@ -253,22 +322,8 @@ def _write_layout(self, broadcast_dir: Path) -> None:
         write_file(path, json.dumps(layout).encode())
         logger.info(f"Wrote the engine layout: {len(layout['fp8'])} checkpoint tensors held in fp8")
 
-    def _replay(self, step_dir: Path) -> None:
-        """Load the latest complete anchor at or before ``step_dir`` (or the boot checkpoint, for a
-        chain from the run's first broadcast) and replay the deltas since."""
-        chain, source = [], step_dir
-        while not anchor_complete(source):
-            base_step = read_manifest(source)["base_step"]
-            if base_step is None:
-                break
-            chain.append(source)
-            source = source.parent / f"step_{base_step}"
-        logger.info(f"Replaying {len(chain)} sparse weight deltas on top of {source.name}")
-        self._load_checkpoint(source if anchor_complete(source) else self._boot_dir())
-        for delta_dir in reversed(chain):
-            self._apply_delta(delta_dir)
-
-    def _apply_delta(self, step_dir: Path) -> None:
+    def _apply_delta(self, step_dir: Path) -> dict[str, int]:
+        """Patch the params with the sparse delta in ``step_dir``; returns the trainer's checksums of the result."""
         start = time.perf_counter()
         prefetched, self.prefetched = self.prefetched, None
         # Always wait: a running prefetch writes into the same host buffer
@@ -284,67 +339,57 @@ def _apply_delta(self, step_dir: Path) -> None:
             self._apply_chunk(entries, gaps, values, tensors, device)
         update_mla_absorbed_weights(model)
         torch.cuda.synchronize()
-        apply_time = time.perf_counter() - start - read_time
-        self._verify(step_dir, checksums)
         logger.info(
             f"Applied sparse weight delta {step_dir.name}: {sum(len(gaps) for _, gaps, _ in delta)} changes, "
-            f"read {read_time:.2f}s (0 if prefetched), apply {apply_time:.2f}s, "
-            f"verify {time.perf_counter() - start - read_time - apply_time:.2f}s"
+            f"read {read_time:.2f}s (0 if prefetched), apply {time.perf_counter() - start - read_time:.2f}s"
         )
+        return checksums
 
-    def _verify(self, step_dir: Path, checksums: dict[str, int]) -> None:
-        """Check every tensor this worker holds against the trainer's checksum of the new weights.
-        Tensors split across TP workers are summed over the TP group first; without TP, tensors a worker
-        holds only part of are not checked."""
+    def _check(self, label: str, names: list[str], totals: torch.Tensor, held: torch.Tensor, expected: dict) -> None:
+        """Check the checksums of the elements this worker holds of each checkpoint tensor (``totals``, over ``held``
+        elements) against the trainer's (``expected``), on every worker of the TP group alike.
+
+        Workers sum their parts over the TP group, in one vector over every checkpoint tensor of the run, so each
+        runs the same collectives whatever it holds. A tensor is checked when its holders cover it the same number
+        of times each (TP slices: once; KV heads replicated over TP workers: once per replica); one held by workers
+        that overlap unevenly (fp8 tensors read in 128-row blocks) is checked only where a worker holds it whole."""
+        start = time.perf_counter()
         assert self.index_map is not None
-        names, totals, partial = self.index_map.held_checksums()
+        canonical = sorted(self.index_map.numel)
+        position = {name: index for index, name in enumerate(canonical)}
+        device = totals.device
+        ids = torch.tensor([position[name] for name in names], dtype=torch.long, device=device)
+        sums = torch.zeros(len(canonical), dtype=torch.int64, device=device).index_put_((ids,), totals)
+        counts = torch.zeros(len(canonical), dtype=torch.int64, device=device).index_put_((ids,), held)
+        numel = torch.tensor([self.index_map.numel[name] for name in canonical], device=device)
+        has_expected = torch.tensor([name in expected for name in canonical], device=device)
+        target = torch.tensor([expected.get(name, 0) for name in canonical], dtype=torch.int64, device=device)
+        whole = counts == numel
+        mismatched = whole & has_expected & (sums != target)
         tp_group = get_tp_group()
-        if tp_group.world_size > 1 and partial.any():
-            reduced = totals[partial]
-            torch.distributed.all_reduce(reduced, group=tp_group.device_group)
-            totals[partial] = reduced
-        elif partial.any():
-            names = [name if not is_partial else None for name, is_partial in zip(names, partial.tolist())]
-        mismatched = [
-            name
-            for name, value in zip(names, totals.tolist())
-            if name is not None and name in checksums and value != checksums[name]
-        ]
-        if mismatched:
-            raise RuntimeError(f"Sparse weight delta {step_dir.name} diverged from the trainer: {mismatched[:8]}")
-
-    def _verify_rows(
-        self, step_dir: Path, sums: dict[str, torch.Tensor], partial: dict[str, bool], expected: dict[str, int]
-    ) -> None:
-        """Check the rows read from an anchor or full checkpoint against the trainer's checksums. Tensors this
-        worker read only some rows of are summed over the TP group first (not for fp8 tensors, read in whole
-        128-row blocks that TP workers may share); without TP they are not checked."""
-        if not sums:
-            return
-        index_map, tp_group = self.index_map, get_tp_group()
-        names = sorted(sums)
-        totals = torch.stack([sums[name] for name in names])
-        partial = torch.tensor([partial[name] for name in names], device=totals.device)
-        if tp_group.world_size > 1 and partial.any():
-            reduced = totals[partial]
-            torch.distributed.all_reduce(reduced, group=tp_group.device_group)
-            totals[partial] = reduced
-        skip = (
-            partial
-            if tp_group.world_size == 1
-            else partial & torch.tensor([name in index_map.fp8 for name in names], device=totals.device)
-        )
-        mismatched = [
-            name
-            for name, value, skipped in zip(names, totals.tolist(), skip.tolist())
-            if not skipped and value != expected[name]
-        ]
-        if mismatched:
-            raise RuntimeError(f"Sparse weight checkpoint {step_dir.name} diverged from the trainer: {mismatched[:8]}")
-
-    def _load_checkpoint(self, step_dir: Path) -> None:
+        if tp_group.world_size > 1:
+            group = tp_group.device_group
+            total_sums, total_counts, most = sums.clone(), counts.clone(), counts.clone()
+            # Holders' fewest elements; non-holders report the most anyone could hold
+            fewest = torch.where(counts > 0, counts, numel)
+            torch.distributed.all_reduce(total_sums, group=group)
+            torch.distributed.all_reduce(total_counts, group=group)
+            torch.distributed.all_reduce(most, op=torch.distributed.ReduceOp.MAX, group=group)
+            torch.distributed.all_reduce(fewest, op=torch.distributed.ReduceOp.MIN, group=group)
+            even = (total_counts > 0) & (total_counts % numel == 0) & (most == fewest)
+            mismatched |= even & has_expected & (total_sums != total_counts // numel * target)
+            any_mismatch = mismatched.any().to(torch.int32)
+            torch.distributed.all_reduce(any_mismatch, op=torch.distributed.ReduceOp.MAX, group=group)
+        else:
+            any_mismatch = mismatched.any()
+        if any_mismatch.item():
+            names_off = [canonical[index] for index in mismatched.nonzero().flatten().tolist()[:8]]
+            raise RuntimeError(f"Weights at {label} diverged from the trainer: {names_off or 'on a TP peer'}")
+        logger.info(f"Verified the weights at {label} in {time.perf_counter() - start:.2f}s")
+
+    def _load_checkpoint(self, step_dir: Path) -> tuple | None:
         """Reload every weight from the HF checkpoint (or anchor) in ``step_dir``. In a sparse run, each worker reads
-        only the rows it holds, straight into its params."""
+        only the rows it holds, straight into its params, and returns their check for ``_check``."""
         model = self._model()
         if self.index_map is None:
             model_loader = get_model_loader(self.load_config)
@@ -358,14 +403,14 @@ def _load_checkpoint(self, step_dir: Path) -> None:
             )
             weights_iterator = model_loader._get_weights_iterator(local_source)
             load_weights_checkpoint_layerwise(model, weights_iterator, self.model_runner.model_config, self.vllm_config)
-            return
+            return None
         start = time.perf_counter()
         index_map, index = self.index_map, checkpoint_index(step_dir)
         device = next(model.parameters()).device
         # Checksums of anchors and full checkpoints written by the trainer (none for the boot checkpoint)
         expected = read_anchor_checksums(step_dir)
         sums: dict[str, torch.Tensor] = {}
-        partial: dict[str, bool] = {}
+        held: dict[str, int] = {}
 
         def read(name: str, begin: int, end: int) -> torch.Tensor:
             num_cols, dtype = index_map.num_cols[name], index[name][1]
@@ -384,11 +429,16 @@ def read(name: str, begin: int, end: int) -> torch.Tensor:
                     if name in expected:
                         num_cols = index_map.num_cols[name]
                         sums[name] = rows_checksum(as_bits(rows).reshape(-1), begin * num_cols)
-                        partial[name] = (end - begin) * num_cols != index_map.numel[name]
+                        held[name] = (end - begin) * num_cols
                     num_bytes += rows.nbytes
         update_mla_absorbed_weights(model)
         torch.cuda.synchronize()
-        self._verify_rows(step_dir, sums, partial, expected)
         logger.info(
             f"Loaded this worker's {num_bytes / 2**30:.2f} GiB of {step_dir.name} in {time.perf_counter() - start:.1f}s"
         )
+        names = sorted(sums)
+        totals = (
+            torch.stack([sums[name] for name in names]) if names else torch.zeros(0, dtype=torch.int64, device=device)
+        )
+        counts = torch.tensor([held[name] for name in names], dtype=torch.int64, device=device)
+        return step_dir.name, names, totals, counts, expected
diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py
index a68837a663c..f007d1da305 100644
--- a/src/prime_rl/inference/vllm/worker/weight_transfer.py
+++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py
@@ -122,19 +122,33 @@ def placeholders():
         self.num_cols = {name: (shape[-1] if shape else 1) for name, (shape, _) in tensors.items()}
         self.numel = {name: prod(shape) for name, (shape, _) in tensors.items()}
         self.checksum_plan = None
+        # ``held_checksums`` kept current by ``apply_delta``; None until the next full recompute
+        self.running: torch.Tensor | None = None
         self.fp8: dict[str, str] = {}
         self.unpatchable = unpatchable
         # Checkpoint tensors this worker holds; unpatchable ones too, so a change to them is caught in apply_delta
         self.names = set(self.records)
 
-    def held_checksums(self) -> tuple[list[str], torch.Tensor, torch.Tensor]:
-        """``bits_checksum`` of the live param elements each checkpoint tensor was written to, weighted by
-        their checkpoint positions: the names, their sums, and a mask of the ones this worker holds only
-        part of (TP slices)."""
+    def checksums(self, full: bool = False) -> tuple[list[str], torch.Tensor, torch.Tensor]:
+        """``held_checksums``, from the running sums ``apply_delta`` keeps unless ``full`` or nothing has
+        recomputed them since the last load."""
+        if full or self.running is None:
+            return self.held_checksums()
+        names, held, *_ = self._checksum_plan()
+        return names, self.running.clone(), held
+
+    def _checksum_plan(self):
         if self.checksum_plan is None:
             self.checksum_plan = self._plan_checksums()
-        names, partial, row_ranges, other_writes = self.checksum_plan
-        totals = torch.zeros(len(names), dtype=torch.int64, device=partial.device)
+        return self.checksum_plan
+
+    def held_checksums(self) -> tuple[list[str], torch.Tensor, torch.Tensor]:
+        """``bits_checksum`` of the live param elements each checkpoint tensor was written to, weighted by
+        their checkpoint positions: the names, their sums, and how many elements of each this worker holds
+        (fewer than the tensor's for TP slices, more workers than one for replicated KV heads). A full pass over
+        the params; it also resets the running sums ``apply_delta`` updates."""
+        names, held, row_ranges, other_writes, _ = self._checksum_plan()
+        totals = torch.zeros(len(names), dtype=torch.int64, device=held.device)
         for index, (starts, ends, ids, num_cols, firsts) in row_ranges.items():
             bits = as_bits(self.params[index].data)
             bits = bits.reshape(-1, bits.shape[-1])
@@ -161,7 +175,8 @@ def held_checksums(self) -> tuple[list[str], torch.Tensor, torch.Tensor]:
                 values = flat[offsets[part, None, None] + row * row_stride + col * col_stride].long()
                 weights = firsts[part, None, None] + num_cols[part, None, None] * row + col
                 totals.index_add_(0, ids[part], (values * weights).sum((1, 2)))
-        return names, totals, partial
+        self.running = totals.clone()
+        return names, totals, held
 
     def _plan_checksums(self):
         """Writes that cover whole rows of a param become row ranges, summed with cumsums over the rows
@@ -192,13 +207,15 @@ def _plan_checksums(self):
         row_ranges = {
             index: tuple(torch.tensor(values, device=device) for values in lists) for index, lists in ranges.items()
         }
-        held = [sum(rows * cols for _, rows, _, cols, *_ in self.records[name]) for name in names]
-        partial = torch.tensor([h < self.numel[name] for h, name in zip(held, names)], device=device)
+        held = torch.tensor(
+            [sum(rows * cols for _, rows, _, cols, *_ in self.records[name]) for name in names], device=device
+        )
         other_writes = {
             shape: tuple(torch.tensor(values, device=device) for values in lists)
             for shape, lists in other_writes.items()
         }
-        return names, partial, row_ranges, other_writes
+        plan_ids = {name: name_id for name_id, name in enumerate(names)}
+        return names, held, row_ranges, other_writes, plan_ids
 
     def fp8_layout(self, model: Module, read_tensor) -> dict[str, str]:
         """The vLLM per-block quantization (``"fp32"`` or ``"ue8m0"`` scales) that reproduces each
@@ -270,7 +287,7 @@ def fp8_layout(self, model: Module, read_tensor) -> dict[str, str]:
                 if self.fp8.setdefault(name, modes[index]) != modes[index]:
                     raise NotImplementedError(f"Sparse weight deltas cannot patch {name}: mixed fp8 layouts")
         self.names |= set(self.records)
-        self.checksum_plan = None
+        self.checksum_plan = self.running = None
         return self.fp8
 
     def _view(self, index: int, rows: int, cols: int, offset: int, strides) -> torch.Tensor:
@@ -296,6 +313,7 @@ def load_rows(self, name: str, rows: torch.Tensor, begin: int) -> None:
         """Set the params from rows ``begin:`` of checkpoint tensor ``name`` (the ``row_ranges`` rows, on the
         params' device): the counterpart of ``apply_delta`` for full checkpoints. fp8 tensors are quantized
         like the engine does, block scales included; tensors the engine casts on load are cast the same way."""
+        self.running = None
         sources = [(name, rows, begin)]
         if name in self.fp8:
             quantized, scales = quantize_fp8_blocks(rows, self.fp8[name] == "ue8m0")
@@ -333,6 +351,13 @@ def apply_delta(self, entries: list[tuple[str, int]], indices: torch.Tensor, del
             for k in range(max_records)
         ]
         entry_ends = torch.tensor([count for _, count in entries], device=device).cumsum(0)
+        # The running checksum of each entry's tensor (-1: not checksummed), updated with the changes alone
+        plan_ids = self._checksum_plan()[4] if self.running is not None else None
+        entry_ids = (
+            torch.tensor([plan_ids.get(name, -1) for name, _ in entries], device=device)
+            if plan_ids is not None
+            else None
+        )
         chunk = 1 << 25
         for start in range(0, indices.numel(), chunk):
             index = indices[start : start + chunk]
@@ -347,11 +372,22 @@ def apply_delta(self, entries: list[tuple[str, int]], indices: torch.Tensor, del
                 hit = (row >= 0) & (row < table[entry, 2]) & (col >= 0) & (col < table[entry, 4])
                 target = (table[entry, 6] + row * table[entry, 7] + col * table[entry, 8])[hit]
                 owner = table[entry, 5][hit]
-                self._scatter_add(owner, target, source[hit])
-
-    def _scatter_add(self, owner: torch.Tensor, target: torch.Tensor, source: torch.Tensor) -> None:
+                tracked = None if entry_ids is None else (entry_ids[entry][hit], index[hit] + 1)
+                self._scatter_add(owner, target, source[hit], tracked)
+
+    def _scatter_add(
+        self,
+        owner: torch.Tensor,
+        target: torch.Tensor,
+        source: torch.Tensor,
+        tracked: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> None:
+        """Add ``source`` to the bits at ``target`` of params ``owner``. ``tracked`` (each change's checksum id
+        and checkpoint position + 1) moves the running checksums by ``position * (new - old)``."""
         order = torch.argsort(owner, stable=True)
         owner, target, source = owner[order], target[order], source[order]
+        if tracked is not None:
+            ids, weights = (values[order] for values in tracked)
         counts = torch.bincount(owner, minlength=len(self.params)).tolist()
         end = 0
         for index, count in enumerate(counts):
@@ -360,7 +396,14 @@ def _scatter_add(self, owner: torch.Tensor, target: torch.Tensor, source: torch.
                 continue
             flat = as_bits(self.params[index].data).view(-1)
             positions = target[begin:end]
-            flat.index_put_((positions,), flat[positions] + source[begin:end])
+            old = flat[positions]
+            new = old + source[begin:end]
+            flat.index_put_((positions,), new)
+            if tracked is not None:
+                part_ids = ids[begin:end]
+                keep = part_ids >= 0
+                change = (new.long() - old.long()) * weights[begin:end]
+                self.running.index_add_(0, part_ids[keep], change[keep])
 
 
 def _restore_checkpoint_format(model: Module) -> list[tuple]:
diff --git a/tests/unit/inference/test_checkpoint_index_map.py b/tests/unit/inference/test_checkpoint_index_map.py
index d2cd0dbf2b2..7f81ed501ce 100644
--- a/tests/unit/inference/test_checkpoint_index_map.py
+++ b/tests/unit/inference/test_checkpoint_index_map.py
@@ -3,6 +3,7 @@
 from torch import nn
 
 from prime_rl.inference.vllm.worker.weight_transfer import CheckpointIndexMap
+from prime_rl.utils.sparse_weights import as_bits, bits_checksum
 
 
 class Toy(nn.Module):
@@ -31,3 +32,23 @@ def test_cast_tensors_load_by_value_and_reject_deltas():
 
     with pytest.raises(NotImplementedError):
         index_map.apply_delta([("norm", 1)], torch.tensor([0]), torch.tensor([1], dtype=torch.int16))
+
+
+def test_running_checksums_follow_deltas():
+    model = Toy()
+    index_map = CheckpointIndexMap(model, {"weight": ([4, 8], "bfloat16"), "norm": ([8], "bfloat16")})
+    weight = torch.randn(4, 8, dtype=torch.bfloat16)
+    index_map.load_rows("weight", weight, 0)
+    names, totals, held = index_map.checksums()
+    assert names == ["weight"] and held.tolist() == [32]
+    assert totals.tolist() == [bits_checksum(as_bits(weight))]
+
+    # Raw-bit differences, including ones that wrap around the int16 range
+    indices = torch.tensor([0, 5, 6, 31])
+    delta = torch.tensor([1, -3, 32767, -32768], dtype=torch.int16)
+    index_map.apply_delta([("weight", 4)], indices, delta)
+    new = as_bits(weight).clone().view(-1)
+    new[indices] += delta
+    _, running, _ = index_map.checksums()
+    assert running.tolist() == [bits_checksum(new)]
+    assert running.tolist() == index_map.checksums(full=True)[1].tolist()

From ca9cfa6d7f7b5c05a270f7fe1b307242d64a051c Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Tue, 3 Nov 2026 00:31:34 +0000
Subject: [PATCH 80/88] perf(weights): apply sparse deltas with one fused
 Triton kernel

Replaces the per-chunk tables, searchsorted, argsort and index_put of the torch
path on GPU: each change finds its tensor by binary search over the delta's
entries, maps through device tables built once per index map, and adds its
bits in place; it also returns its running-checksum contribution. The kernel
is compiled when the index map is built, outside any pause. The torch path
stays for CPU and as the reference.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../inference/vllm/worker/filesystem.py       |   2 +
 .../inference/vllm/worker/sparse_apply.py     | 140 ++++++++++++++++++
 .../inference/vllm/worker/weight_transfer.py  |  78 +++++++++-
 .../inference/test_checkpoint_index_map.py    |  57 +++++++
 4 files changed, 276 insertions(+), 1 deletion(-)
 create mode 100644 src/prime_rl/inference/vllm/worker/sparse_apply.py

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 85f3d9d9cad..b0a8af7aad4 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -309,6 +309,8 @@ def read_tensor(name: str) -> torch.Tensor:
             return torch.from_numpy(data).view(dtype).view(shape)
 
         index_map.fp8_layout(model, read_tensor)
+        if next(model.parameters()).is_cuda:
+            index_map.fused_tables(warm_up=True)
         self.index_map = index_map
 
     def _write_layout(self, broadcast_dir: Path) -> None:
diff --git a/src/prime_rl/inference/vllm/worker/sparse_apply.py b/src/prime_rl/inference/vllm/worker/sparse_apply.py
new file mode 100644
index 00000000000..a4b5352cce2
--- /dev/null
+++ b/src/prime_rl/inference/vllm/worker/sparse_apply.py
@@ -0,0 +1,140 @@
+"""One fused kernel that applies a sparse weight delta to the live params in place.
+
+Each change is a checkpoint tensor's flat index and a raw-bit difference. The kernel finds its tensor (a binary
+search over the delta's entries), maps the index through that tensor's writes (where the engine's loader copied
+each block of the checkpoint tensor, recorded once per ``CheckpointIndexMap``), and adds the difference to every
+param element it was written to, wrapping in the bits dtype. Indices are unique within a delta, so the
+read-modify-writes need no sort or atomics. It also returns each change's contribution to the position-weighted
+checksum of its tensor, ``(index + 1) * (new - old)``, so the running checksums move with the changes alone.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+# Fields of one write: first row, rows, first column, columns, param, offset, row stride, column stride
+WRITE_FIELDS = 8
+
+_TL_BITS = {torch.uint8: tl.uint8, torch.int16: tl.int16, torch.int32: tl.int32, torch.int64: tl.int64}
+
+
+@triton.jit
+def _sparse_apply_kernel(
+    indices_ptr,
+    deltas_ptr,
+    contributions_ptr,
+    entry_ends_ptr,
+    entry_names_ptr,
+    writes_ptr,
+    num_cols_ptr,
+    param_ptrs_ptr,
+    num_changes,
+    num_entries,
+    search_steps,
+    MAX_WRITES: tl.constexpr,
+    TRACK: tl.constexpr,
+    BITS: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
+    mask = offsets < num_changes
+    index = tl.load(indices_ptr + offsets, mask=mask, other=0)
+    delta = tl.load(deltas_ptr + offsets, mask=mask, other=0)
+
+    # The entry of each change: the first whose cumulative end exceeds its position
+    low = tl.zeros([BLOCK], dtype=tl.int64)
+    high = tl.full([BLOCK], num_entries, dtype=tl.int64)
+    for _ in range(search_steps):
+        active = low < high
+        middle = (low + high) // 2
+        end = tl.load(entry_ends_ptr + middle, mask=mask & active, other=0)
+        right = active & (end <= offsets)
+        low = tl.where(right, middle + 1, low)
+        high = tl.where(active & (end > offsets), middle, high)
+
+    name = tl.load(entry_names_ptr + low, mask=mask, other=-1)
+    valid = mask & (name >= 0)
+    num_cols = tl.load(num_cols_ptr + name, mask=valid, other=1)
+    row = index // num_cols
+    col = index % num_cols
+    total = tl.zeros([BLOCK], dtype=tl.int64)
+    for k in tl.static_range(MAX_WRITES):
+        write = writes_ptr + (name * MAX_WRITES + k) * 8
+        row0 = tl.load(write + 0, mask=valid, other=0)
+        rows = tl.load(write + 1, mask=valid, other=0)
+        col0 = tl.load(write + 2, mask=valid, other=0)
+        cols = tl.load(write + 3, mask=valid, other=0)
+        r = row - row0
+        c = col - col0
+        hit = valid & (r >= 0) & (r < rows) & (c >= 0) & (c < cols)
+        param = tl.load(write + 4, mask=hit, other=0)
+        offset = tl.load(write + 5, mask=hit, other=0)
+        row_stride = tl.load(write + 6, mask=hit, other=0)
+        col_stride = tl.load(write + 7, mask=hit, other=0)
+        base = tl.load(param_ptrs_ptr + param, mask=hit, other=0).to(tl.pointer_type(BITS))
+        target = base + (offset + r * row_stride + c * col_stride)
+        old = tl.load(target, mask=hit, other=0)
+        new = old + delta
+        tl.store(target, new, mask=hit)
+        if TRACK:
+            total += tl.where(hit, (new.to(tl.int64) - old.to(tl.int64)) * (index + 1), 0)
+    if TRACK:
+        tl.store(contributions_ptr + offsets, total, mask=mask)
+
+
+def sparse_apply(
+    indices: torch.Tensor,
+    deltas: torch.Tensor,
+    entry_ends: torch.Tensor,
+    entry_names: torch.Tensor,
+    writes: torch.Tensor,
+    num_cols: torch.Tensor,
+    param_ptrs: torch.Tensor,
+    track: bool,
+    block: int = 1024,
+) -> torch.Tensor | None:
+    """Apply ``deltas`` (raw-bit differences in one bits dtype) at the checkpoint ``indices`` of consecutive
+    entries ending at ``entry_ends`` (entry tensor ids ``entry_names``, -1 for tensors this worker does not hold).
+    ``writes[name]`` are the tensor's writes, ``WRITE_FIELDS`` each, padded with zero-row writes. Returns each
+    change's checksum contribution when ``track``."""
+    num_changes, num_entries = indices.numel(), entry_ends.numel()
+    contributions = torch.empty(num_changes if track else 1, dtype=torch.int64, device=indices.device)
+    if num_changes == 0:
+        return contributions[:0] if track else None
+    grid = (triton.cdiv(num_changes, block),)
+    _sparse_apply_kernel[grid](
+        indices,
+        deltas,
+        contributions,
+        entry_ends,
+        entry_names,
+        writes,
+        num_cols,
+        param_ptrs,
+        num_changes,
+        num_entries,
+        num_entries.bit_length() + 1,
+        MAX_WRITES=writes.shape[1],
+        TRACK=track,
+        BITS=_TL_BITS[deltas.dtype],
+        BLOCK=block,
+    )
+    return contributions if track else None
+
+
+def warm_up(writes: torch.Tensor, num_cols: torch.Tensor, param_ptrs: torch.Tensor) -> None:
+    """Compile the kernel for every bits dtype before it runs inside a weight-update pause: one change of a tensor
+    this worker does not hold, so nothing is written."""
+    device = writes.device
+    for bits in _TL_BITS:
+        for track in (False, True):
+            sparse_apply(
+                torch.zeros(1, dtype=torch.int64, device=device),
+                torch.zeros(1, dtype=bits, device=device),
+                torch.ones(1, dtype=torch.int64, device=device),
+                torch.full((1,), -1, dtype=torch.int64, device=device),
+                writes,
+                num_cols,
+                param_ptrs,
+                track,
+            )
diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py
index f007d1da305..9af0fbf7fac 100644
--- a/src/prime_rl/inference/vllm/worker/weight_transfer.py
+++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py
@@ -124,6 +124,8 @@ def placeholders():
         self.checksum_plan = None
         # ``held_checksums`` kept current by ``apply_delta``; None until the next full recompute
         self.running: torch.Tensor | None = None
+        # The fused kernel's device tables (``fused_tables``)
+        self.fused = None
         self.fp8: dict[str, str] = {}
         self.unpatchable = unpatchable
         # Checkpoint tensors this worker holds; unpatchable ones too, so a change to them is caught in apply_delta
@@ -287,7 +289,7 @@ def fp8_layout(self, model: Module, read_tensor) -> dict[str, str]:
                 if self.fp8.setdefault(name, modes[index]) != modes[index]:
                     raise NotImplementedError(f"Sparse weight deltas cannot patch {name}: mixed fp8 layouts")
         self.names |= set(self.records)
-        self.checksum_plan = self.running = None
+        self.checksum_plan = self.running = self.fused = None
         return self.fp8
 
     def _view(self, index: int, rows: int, cols: int, offset: int, strides) -> torch.Tensor:
@@ -327,6 +329,43 @@ def load_rows(self, name: str, rows: torch.Tensor, begin: int) -> None:
                 else:
                     self._view(index, rows_, cols, offset, strides).copy_(as_bits(block))
 
+    def fused_tables(self, warm_up: bool = False) -> dict:
+        """Device tables of ``sparse_apply``: tensor ids, every tensor's writes, columns and running-checksum id, and
+        the params' addresses. Built once (and again after ``fp8_layout`` adds the block scales); ``warm_up`` also
+        compiles the kernel, so neither happens inside a weight-update pause."""
+        if self.fused is None:
+            from prime_rl.inference.vllm.worker.sparse_apply import WRITE_FIELDS
+
+            device = self.params[0].device
+            names = sorted(self.records)
+            max_writes = max((len(writes) for writes in self.records.values()), default=1)
+            padding = [0] * WRITE_FIELDS  # zero rows: never hit
+            writes = torch.tensor(
+                [
+                    [list(write) for write in self.records[name]] + [padding] * (max_writes - len(self.records[name]))
+                    for name in names
+                ],
+                dtype=torch.int64,
+            ).view(len(names), max_writes, WRITE_FIELDS)
+            plan_ids = self._checksum_plan()[4]
+            self.fused = {
+                "ids": {name: name_id for name_id, name in enumerate(names)},
+                "writes": writes.to(device),
+                "num_cols": torch.tensor([self.num_cols[name] for name in names], dtype=torch.int64, device=device),
+                "plan_ids": torch.tensor([plan_ids.get(name, -1) for name in names], dtype=torch.int64, device=device),
+                "param_ptrs": torch.tensor(
+                    [param.data_ptr() for param in self.params], dtype=torch.int64, device=device
+                ),
+                "data_ptrs": [param.data_ptr() for param in self.params],
+            }
+            if warm_up:
+                from prime_rl.inference.vllm.worker.sparse_apply import warm_up as compile_kernel
+
+                compile_kernel(self.fused["writes"], self.fused["num_cols"], self.fused["param_ptrs"])
+        if [param.data_ptr() for param in self.params] != self.fused["data_ptrs"]:
+            raise RuntimeError("Sparse weight deltas need the params where the index map found them: they moved")
+        return self.fused
+
     @torch.no_grad()
     def apply_delta(self, entries: list[tuple[str, int]], indices: torch.Tensor, delta: torch.Tensor) -> None:
         """Add ``delta`` (raw-bit differences, wrapping in the bits dtype) to the params at the checkpoint
@@ -335,6 +374,14 @@ def apply_delta(self, entries: list[tuple[str, int]], indices: torch.Tensor, del
         if changed:
             raise NotImplementedError(f"Sparse weight deltas cannot patch {changed[:8]}: the engine casts them on load")
         device = indices.device
+        if device.type == "cuda":
+            self._apply_delta_fused(entries, indices, delta)
+        else:
+            self._apply_delta_torch(entries, indices, delta)
+
+    def _apply_delta_torch(self, entries: list[tuple[str, int]], indices: torch.Tensor, delta: torch.Tensor) -> None:
+        """``apply_delta`` in plain torch ops (CPU, and the reference for the fused kernel)."""
+        device = indices.device
         max_records = max((len(self.records.get(name, ())) for name, _ in entries), default=0)
         # Per entry and write k: [num_cols, row0, rows, col0, cols, param, offset, row stride, col stride].
         tables = [
@@ -375,6 +422,35 @@ def apply_delta(self, entries: list[tuple[str, int]], indices: torch.Tensor, del
                 tracked = None if entry_ids is None else (entry_ids[entry][hit], index[hit] + 1)
                 self._scatter_add(owner, target, source[hit], tracked)
 
+    def _apply_delta_fused(self, entries: list[tuple[str, int]], indices: torch.Tensor, delta: torch.Tensor) -> None:
+        from prime_rl.inference.vllm.worker.sparse_apply import sparse_apply
+
+        tables = self.fused_tables()
+        device = indices.device
+        counts = torch.tensor([count for _, count in entries], dtype=torch.int64, device=device)
+        entry_ends = counts.cumsum(0)
+        entry_names = torch.tensor(
+            [tables["ids"].get(name, -1) for name, _ in entries], dtype=torch.int64, device=device
+        )
+        track = self.running is not None
+        contributions = sparse_apply(
+            indices.long(),
+            delta,
+            entry_ends,
+            entry_names,
+            tables["writes"],
+            tables["num_cols"],
+            tables["param_ptrs"],
+            track,
+        )
+        if track and contributions.numel():
+            # Per entry: a difference of the cumulative contributions at its ends (int64 sums wrap like the checksums)
+            cumulative = torch.cat([contributions.new_zeros(1), contributions.cumsum(0)])
+            sums = cumulative[entry_ends] - cumulative[entry_ends - counts]
+            plan_ids = tables["plan_ids"][entry_names.clamp(min=0)]
+            keep = (entry_names >= 0) & (plan_ids >= 0)
+            self.running.index_add_(0, plan_ids[keep], sums[keep])
+
     def _scatter_add(
         self,
         owner: torch.Tensor,
diff --git a/tests/unit/inference/test_checkpoint_index_map.py b/tests/unit/inference/test_checkpoint_index_map.py
index 7f81ed501ce..cf4f391f458 100644
--- a/tests/unit/inference/test_checkpoint_index_map.py
+++ b/tests/unit/inference/test_checkpoint_index_map.py
@@ -52,3 +52,60 @@ def test_running_checksums_follow_deltas():
     _, running, _ = index_map.checksums()
     assert running.tolist() == [bits_checksum(new)]
     assert running.tolist() == index_map.checksums(full=True)[1].tolist()
+
+
+class Sliced(nn.Module):
+    """Checkpoint tensors written as row halves (like TP slices), through a column-strided view, and as fp32."""
+
+    def __init__(self, device):
+        super().__init__()
+        self.top = nn.Parameter(torch.zeros(4, 16, dtype=torch.bfloat16, device=device))
+        self.bottom = nn.Parameter(torch.zeros(4, 16, dtype=torch.bfloat16, device=device))
+        self.packed = nn.Parameter(torch.zeros(6, 10, dtype=torch.bfloat16, device=device))
+        self.scale = nn.Parameter(torch.zeros(32, dtype=torch.float32, device=device))
+
+    def load_weights(self, weights):
+        for name, tensor in weights:
+            if name == "w":
+                self.top.data.copy_(tensor[:4])
+                self.bottom.data.copy_(tensor[4:])
+            elif name == "v":
+                self.packed.data[:, 2:7].copy_(tensor)
+            else:
+                self.scale.data.copy_(tensor)
+
+
+@pytest.mark.gpu
+def test_fused_apply_matches_torch_apply():
+    tensors = {"w": ([8, 16], "bfloat16"), "v": ([6, 5], "bfloat16"), "s": ([32], "float32")}
+    checkpoint = {
+        "w": torch.randn(8, 16, dtype=torch.bfloat16),
+        "v": torch.randn(6, 5, dtype=torch.bfloat16),
+        "s": torch.randn(32, dtype=torch.float32),
+    }
+    maps = {}
+    for device in ("cpu", "cuda"):
+        model = Sliced(device)
+        index_map = CheckpointIndexMap(model, tensors)
+        for name, tensor in checkpoint.items():
+            index_map.load_rows(name, tensor.view(-1, tensors[name][0][-1]).to(device), 0)
+        index_map.checksums()
+        maps[device] = (model, index_map)
+    maps["cuda"][1].fused_tables(warm_up=True)
+
+    generator = torch.Generator().manual_seed(0)
+    for _ in range(3):
+        w = torch.randperm(128, generator=generator)[:40].sort().values
+        v = torch.randperm(30, generator=generator)[:7].sort().values
+        s = torch.randperm(32, generator=generator)[:5].sort().values
+        bf16 = torch.randint(-32768, 32768, (47,), generator=generator, dtype=torch.int32).to(torch.int16)
+        fp32 = torch.randint(-(2**31), 2**31, (5,), generator=generator, dtype=torch.int64).to(torch.int32)
+        for device, (_, index_map) in maps.items():
+            index_map.apply_delta([("w", 40), ("v", 7)], torch.cat([w, v]).to(device), bf16.to(device))
+            index_map.apply_delta([("s", 5)], s.to(device), fp32.to(device))
+    (cpu_model, cpu_map), (gpu_model, gpu_map) = maps["cpu"], maps["cuda"]
+    for (name, cpu_param), gpu_param in zip(cpu_model.named_parameters(), gpu_model.parameters()):
+        assert torch.equal(as_bits(cpu_param.data), as_bits(gpu_param.data).cpu()), name
+    running = gpu_map.checksums()[1].cpu()
+    assert running.tolist() == cpu_map.checksums()[1].tolist()
+    assert running.tolist() == gpu_map.checksums(full=True)[1].cpu().tolist()

From f1652fec70e6048e84c63c6d2c0bb8c3811cdab1 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Tue, 3 Nov 2026 00:56:11 +0000
Subject: [PATCH 81/88] feat(weights): offset tables for writes into permuted
 kernel layouts

A write can carry row and/or column offset tables instead of strides (element
(r, c) at offset + rows[r] + cols[c]). Loads, deltas (torch path and fused
kernel) and checksums all go through them. This is what kernel layouts that
shuffle rows or swizzle block scales need (Blackwell MXFP8/NVFP4 MoE and
linear backends); nothing creates such maps yet.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../inference/vllm/worker/sparse_apply.py     |  26 +++--
 .../inference/vllm/worker/weight_transfer.py  | 100 +++++++++++++++---
 .../inference/test_checkpoint_index_map.py    |  27 +++++
 3 files changed, 133 insertions(+), 20 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/sparse_apply.py b/src/prime_rl/inference/vllm/worker/sparse_apply.py
index a4b5352cce2..c379d13cf61 100644
--- a/src/prime_rl/inference/vllm/worker/sparse_apply.py
+++ b/src/prime_rl/inference/vllm/worker/sparse_apply.py
@@ -12,8 +12,9 @@
 import triton
 import triton.language as tl
 
-# Fields of one write: first row, rows, first column, columns, param, offset, row stride, column stride
-WRITE_FIELDS = 8
+# Fields of one write: first row, rows, first column, columns, param, offset, row stride, column stride, and where its
+# row and column offset tables start in the maps buffer (-1: the stride applies)
+WRITE_FIELDS = 10
 
 _TL_BITS = {torch.uint8: tl.uint8, torch.int16: tl.int16, torch.int32: tl.int32, torch.int64: tl.int64}
 
@@ -28,6 +29,7 @@ def _sparse_apply_kernel(
     writes_ptr,
     num_cols_ptr,
     param_ptrs_ptr,
+    maps_ptr,
     num_changes,
     num_entries,
     search_steps,
@@ -59,7 +61,7 @@ def _sparse_apply_kernel(
     col = index % num_cols
     total = tl.zeros([BLOCK], dtype=tl.int64)
     for k in tl.static_range(MAX_WRITES):
-        write = writes_ptr + (name * MAX_WRITES + k) * 8
+        write = writes_ptr + (name * MAX_WRITES + k) * 10
         row0 = tl.load(write + 0, mask=valid, other=0)
         rows = tl.load(write + 1, mask=valid, other=0)
         col0 = tl.load(write + 2, mask=valid, other=0)
@@ -71,8 +73,16 @@ def _sparse_apply_kernel(
         offset = tl.load(write + 5, mask=hit, other=0)
         row_stride = tl.load(write + 6, mask=hit, other=0)
         col_stride = tl.load(write + 7, mask=hit, other=0)
+        row_table = tl.load(write + 8, mask=hit, other=-1)
+        col_table = tl.load(write + 9, mask=hit, other=-1)
+        row_part = tl.where(
+            row_table >= 0, tl.load(maps_ptr + row_table + r, mask=hit & (row_table >= 0), other=0), r * row_stride
+        )
+        col_part = tl.where(
+            col_table >= 0, tl.load(maps_ptr + col_table + c, mask=hit & (col_table >= 0), other=0), c * col_stride
+        )
         base = tl.load(param_ptrs_ptr + param, mask=hit, other=0).to(tl.pointer_type(BITS))
-        target = base + (offset + r * row_stride + c * col_stride)
+        target = base + (offset + row_part + col_part)
         old = tl.load(target, mask=hit, other=0)
         new = old + delta
         tl.store(target, new, mask=hit)
@@ -90,12 +100,14 @@ def sparse_apply(
     writes: torch.Tensor,
     num_cols: torch.Tensor,
     param_ptrs: torch.Tensor,
+    maps: torch.Tensor,
     track: bool,
     block: int = 1024,
 ) -> torch.Tensor | None:
     """Apply ``deltas`` (raw-bit differences in one bits dtype) at the checkpoint ``indices`` of consecutive
     entries ending at ``entry_ends`` (entry tensor ids ``entry_names``, -1 for tensors this worker does not hold).
-    ``writes[name]`` are the tensor's writes, ``WRITE_FIELDS`` each, padded with zero-row writes. Returns each
+    ``writes[name]`` are the tensor's writes, ``WRITE_FIELDS`` each, padded with zero-row writes; ``maps`` holds their
+    row and column offset tables. Returns each
     change's checksum contribution when ``track``."""
     num_changes, num_entries = indices.numel(), entry_ends.numel()
     contributions = torch.empty(num_changes if track else 1, dtype=torch.int64, device=indices.device)
@@ -111,6 +123,7 @@ def sparse_apply(
         writes,
         num_cols,
         param_ptrs,
+        maps,
         num_changes,
         num_entries,
         num_entries.bit_length() + 1,
@@ -122,7 +135,7 @@ def sparse_apply(
     return contributions if track else None
 
 
-def warm_up(writes: torch.Tensor, num_cols: torch.Tensor, param_ptrs: torch.Tensor) -> None:
+def warm_up(writes: torch.Tensor, num_cols: torch.Tensor, param_ptrs: torch.Tensor, maps: torch.Tensor) -> None:
     """Compile the kernel for every bits dtype before it runs inside a weight-update pause: one change of a tensor
     this worker does not hold, so nothing is written."""
     device = writes.device
@@ -136,5 +149,6 @@ def warm_up(writes: torch.Tensor, num_cols: torch.Tensor, param_ptrs: torch.Tens
                 writes,
                 num_cols,
                 param_ptrs,
+                maps,
                 track,
             )
diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py
index 9af0fbf7fac..8e9f0555b2f 100644
--- a/src/prime_rl/inference/vllm/worker/weight_transfer.py
+++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py
@@ -126,11 +126,53 @@ def placeholders():
         self.running: torch.Tensor | None = None
         # The fused kernel's device tables (``fused_tables``)
         self.fused = None
+        # (name, write index) -> (row offsets, column offsets) for writes into a kernel layout that permutes rows or
+        # columns (e.g. swizzled block scales): element (r, c) of the write is at its offset + rows[r] + cols[c], with
+        # the write's stride for a missing table
+        self.maps: dict[tuple[str, int], tuple[torch.Tensor | None, torch.Tensor | None]] = {}
+        self.map_cache = None
         self.fp8: dict[str, str] = {}
         self.unpatchable = unpatchable
         # Checkpoint tensors this worker holds; unpatchable ones too, so a change to them is caught in apply_delta
         self.names = set(self.records)
 
+    def map_write(
+        self, name: str, k: int, row_offsets: torch.Tensor | None = None, col_offsets: torch.Tensor | None = None
+    ) -> None:
+        """Place write ``k`` of ``name`` through offset tables (``maps``) instead of its strides."""
+        _, rows, _, cols, *_ = self.records[name][k]
+        assert row_offsets is None or row_offsets.shape == (rows,)
+        assert col_offsets is None or col_offsets.shape == (cols,)
+        self.maps[(name, k)] = (row_offsets, col_offsets)
+        self.map_cache = self.checksum_plan = self.running = self.fused = None
+
+    def _map_tables(self) -> tuple[torch.Tensor, dict[tuple[str, int], tuple[int, int]]]:
+        """Every offset table in one int64 device buffer, and each mapped write's table starts (-1: none)."""
+        if self.map_cache is None:
+            device = self.params[0].device
+            parts, starts, size = [torch.zeros(1, dtype=torch.int64)], {}, 1
+            for key, tables in self.maps.items():
+                first = []
+                for table in tables:
+                    if table is None:
+                        first.append(-1)
+                        continue
+                    first.append(size)
+                    parts.append(table.to(torch.int64).cpu())
+                    size += table.numel()
+                starts[key] = tuple(first)
+            self.map_cache = (torch.cat(parts).to(device), starts)
+        return self.map_cache
+
+    def _targets(self, name: str, k: int) -> torch.Tensor:
+        """Flat element offsets of write ``k`` of ``name`` in its param, ``[rows, cols]``."""
+        _, rows, _, cols, _, offset, row_stride, col_stride = self.records[name][k]
+        device = self.params[0].device
+        row_table, col_table = self.maps.get((name, k), (None, None))
+        row_part = row_table.to(device) if row_table is not None else torch.arange(rows, device=device) * row_stride
+        col_part = col_table.to(device) if col_table is not None else torch.arange(cols, device=device) * col_stride
+        return offset + row_part[:, None] + col_part[None, :]
+
     def checksums(self, full: bool = False) -> tuple[list[str], torch.Tensor, torch.Tensor]:
         """``held_checksums``, from the running sums ``apply_delta`` keeps unless ``full`` or nothing has
         recomputed them since the last load."""
@@ -149,7 +191,7 @@ def held_checksums(self) -> tuple[list[str], torch.Tensor, torch.Tensor]:
         their checkpoint positions: the names, their sums, and how many elements of each this worker holds
         (fewer than the tensor's for TP slices, more workers than one for replicated KV heads). A full pass over
         the params; it also resets the running sums ``apply_delta`` updates."""
-        names, held, row_ranges, other_writes, _ = self._checksum_plan()
+        names, held, row_ranges, other_writes, _, mapped = self._checksum_plan()
         totals = torch.zeros(len(names), dtype=torch.int64, device=held.device)
         for index, (starts, ends, ids, num_cols, firsts) in row_ranges.items():
             bits = as_bits(self.params[index].data)
@@ -177,6 +219,14 @@ def held_checksums(self) -> tuple[list[str], torch.Tensor, torch.Tensor]:
                 values = flat[offsets[part, None, None] + row * row_stride + col * col_stride].long()
                 weights = firsts[part, None, None] + num_cols[part, None, None] * row + col
                 totals.index_add_(0, ids[part], (values * weights).sum((1, 2)))
+        for name_id, name, k in mapped:
+            row0, rows, col0, cols, index, *_ = self.records[name][k]
+            values = as_bits(self.params[index].data).view(-1)[self._targets(name, k)].long()
+            num_cols = self.num_cols[name]
+            row = torch.arange(rows, device=values.device)[:, None]
+            col = torch.arange(cols, device=values.device)
+            weights = (row0 + row) * num_cols + col0 + col + 1
+            totals[name_id] += (values * weights).sum()
         self.running = totals.clone()
         return names, totals, held
 
@@ -192,11 +242,14 @@ def _plan_checksums(self):
         device = self.params[0].device
         ranges: dict[int, list[list[int]]] = defaultdict(lambda: [[], [], [], [], []])
         other_writes: dict[tuple, list[list[int]]] = defaultdict(lambda: [[], [], [], []])
+        mapped = []
         for name_id, name in enumerate(names):
             num_cols = self.num_cols[name]
-            for row0, rows, col0, cols, index, offset, row_stride, col_stride in self.records[name]:
+            for k, (row0, rows, col0, cols, index, offset, row_stride, col_stride) in enumerate(self.records[name]):
                 width = self.params[index].shape[-1]
-                if col_stride == 1 and row_stride == width and cols == width and offset % width == 0:
+                if (name, k) in self.maps:
+                    mapped.append((name_id, name, k))
+                elif col_stride == 1 and row_stride == width and cols == width and offset % width == 0:
                     start = offset // width
                     for values, value in zip(
                         ranges[index], (start, start + rows, name_id, num_cols, (row0 - start) * num_cols + col0 + 1)
@@ -217,7 +270,7 @@ def _plan_checksums(self):
             for shape, lists in other_writes.items()
         }
         plan_ids = {name: name_id for name_id, name in enumerate(names)}
-        return names, held, row_ranges, other_writes, plan_ids
+        return names, held, row_ranges, other_writes, plan_ids, mapped
 
     def fp8_layout(self, model: Module, read_tensor) -> dict[str, str]:
         """The vLLM per-block quantization (``"fp32"`` or ``"ue8m0"`` scales) that reproduces each
@@ -321,9 +374,12 @@ def load_rows(self, name: str, rows: torch.Tensor, begin: int) -> None:
             quantized, scales = quantize_fp8_blocks(rows, self.fp8[name] == "ue8m0")
             sources = [(name, quantized, begin), (f"{name}_scale_inv", scales, begin // FP8_BLOCK[0])]
         for source_name, source, first in sources:
-            for row0, rows_, col0, cols, index, offset, *strides in self.records[source_name]:
+            for k, (row0, rows_, col0, cols, index, offset, *strides) in enumerate(self.records[source_name]):
                 block = source[row0 - first : row0 - first + rows_, col0 : col0 + cols]
-                if source_name in self.unpatchable:
+                if (source_name, k) in self.maps:
+                    flat = as_bits(self.params[index].data).view(-1)
+                    flat.index_put_((self._targets(source_name, k).view(-1),), as_bits(block).reshape(-1))
+                elif source_name in self.unpatchable:
                     flat = self.params[index].data.view(-1)
                     flat.as_strided((rows_, cols), tuple(strides), offset).copy_(block)
                 else:
@@ -339,10 +395,12 @@ def fused_tables(self, warm_up: bool = False) -> dict:
             device = self.params[0].device
             names = sorted(self.records)
             max_writes = max((len(writes) for writes in self.records.values()), default=1)
-            padding = [0] * WRITE_FIELDS  # zero rows: never hit
+            map_buffer, map_starts = self._map_tables()
+            padding = [0] * (WRITE_FIELDS - 2) + [-1, -1]  # zero rows: never hit
             writes = torch.tensor(
                 [
-                    [list(write) for write in self.records[name]] + [padding] * (max_writes - len(self.records[name]))
+                    [[*write, *map_starts.get((name, k), (-1, -1))] for k, write in enumerate(self.records[name])]
+                    + [padding] * (max_writes - len(self.records[name]))
                     for name in names
                 ],
                 dtype=torch.int64,
@@ -357,11 +415,14 @@ def fused_tables(self, warm_up: bool = False) -> dict:
                     [param.data_ptr() for param in self.params], dtype=torch.int64, device=device
                 ),
                 "data_ptrs": [param.data_ptr() for param in self.params],
+                "maps": map_buffer,
             }
             if warm_up:
                 from prime_rl.inference.vllm.worker.sparse_apply import warm_up as compile_kernel
 
-                compile_kernel(self.fused["writes"], self.fused["num_cols"], self.fused["param_ptrs"])
+                compile_kernel(
+                    self.fused["writes"], self.fused["num_cols"], self.fused["param_ptrs"], self.fused["maps"]
+                )
         if [param.data_ptr() for param in self.params] != self.fused["data_ptrs"]:
             raise RuntimeError("Sparse weight deltas need the params where the index map found them: they moved")
         return self.fused
@@ -383,13 +444,15 @@ def _apply_delta_torch(self, entries: list[tuple[str, int]], indices: torch.Tens
         """``apply_delta`` in plain torch ops (CPU, and the reference for the fused kernel)."""
         device = indices.device
         max_records = max((len(self.records.get(name, ())) for name, _ in entries), default=0)
-        # Per entry and write k: [num_cols, row0, rows, col0, cols, param, offset, row stride, col stride].
+        map_buffer, map_starts = self._map_tables()
+        # Per entry and write k: [num_cols, row0, rows, col0, cols, param, offset, row stride, col stride, row table,
+        # column table].
         tables = [
             torch.tensor(
                 [
-                    [self.num_cols[name], *self.records[name][k]]
+                    [self.num_cols[name], *self.records[name][k], *map_starts.get((name, k), (-1, -1))]
                     if len(self.records.get(name, ())) > k
-                    else [1, 0, 0, 0, 0, 0, 0, 0, 0]
+                    else [1, 0, 0, 0, 0, 0, 0, 0, 0, -1, -1]
                     for name, _ in entries
                 ],
                 dtype=torch.int64,
@@ -417,8 +480,16 @@ def _apply_delta_torch(self, entries: list[tuple[str, int]], indices: torch.Tens
                 row = index // num_cols - table[entry, 1]
                 col = index % num_cols - table[entry, 3]
                 hit = (row >= 0) & (row < table[entry, 2]) & (col >= 0) & (col < table[entry, 4])
-                target = (table[entry, 6] + row * table[entry, 7] + col * table[entry, 8])[hit]
-                owner = table[entry, 5][hit]
+                write, row, col = table[entry[hit]], row[hit], col[hit]
+                last = map_buffer.numel() - 1
+                row_part = torch.where(
+                    write[:, 9] >= 0, map_buffer[(write[:, 9] + row).clamp(0, last)], row * write[:, 7]
+                )
+                col_part = torch.where(
+                    write[:, 10] >= 0, map_buffer[(write[:, 10] + col).clamp(0, last)], col * write[:, 8]
+                )
+                target = write[:, 6] + row_part + col_part
+                owner = write[:, 5]
                 tracked = None if entry_ids is None else (entry_ids[entry][hit], index[hit] + 1)
                 self._scatter_add(owner, target, source[hit], tracked)
 
@@ -441,6 +512,7 @@ def _apply_delta_fused(self, entries: list[tuple[str, int]], indices: torch.Tens
             tables["writes"],
             tables["num_cols"],
             tables["param_ptrs"],
+            tables["maps"],
             track,
         )
         if track and contributions.numel():
diff --git a/tests/unit/inference/test_checkpoint_index_map.py b/tests/unit/inference/test_checkpoint_index_map.py
index cf4f391f458..1e869a981b3 100644
--- a/tests/unit/inference/test_checkpoint_index_map.py
+++ b/tests/unit/inference/test_checkpoint_index_map.py
@@ -54,6 +54,31 @@ def test_running_checksums_follow_deltas():
     assert running.tolist() == index_map.checksums(full=True)[1].tolist()
 
 
+def test_mapped_writes_load_patch_and_checksum_through_their_tables():
+    model = Toy()
+    index_map = CheckpointIndexMap(model, {"weight": ([4, 8], "bfloat16"), "norm": ([8], "bfloat16")})
+    # A kernel layout with rows in the order 2, 0, 3, 1 and the two column halves swapped
+    rows, cols = torch.tensor([2, 0, 3, 1]), torch.cat([torch.arange(4, 8), torch.arange(4)])
+    index_map.map_write("weight", 0, row_offsets=rows * 8, col_offsets=cols)
+    weight = torch.randn(4, 8, dtype=torch.bfloat16)
+    index_map.load_rows("weight", weight, 0)
+    kernel = torch.empty_like(weight)
+    kernel[rows[:, None], cols[None, :]] = weight
+    assert torch.equal(model.weight.data, kernel)
+    assert index_map.checksums()[1].tolist() == [bits_checksum(as_bits(weight))]
+
+    indices = torch.tensor([1, 9, 30])
+    delta = torch.tensor([5, -7, 11], dtype=torch.int16)
+    index_map.apply_delta([("weight", 3)], indices, delta)
+    new = as_bits(weight).clone().view(-1)
+    new[indices] += delta
+    kernel_bits = torch.empty_like(as_bits(weight))
+    kernel_bits[rows[:, None], cols[None, :]] = new.view(4, 8)
+    assert torch.equal(as_bits(model.weight.data), kernel_bits)
+    assert index_map.checksums()[1].tolist() == [bits_checksum(new)]
+    assert index_map.checksums(full=True)[1].tolist() == [bits_checksum(new)]
+
+
 class Sliced(nn.Module):
     """Checkpoint tensors written as row halves (like TP slices), through a column-strided view, and as fp32."""
 
@@ -87,6 +112,8 @@ def test_fused_apply_matches_torch_apply():
     for device in ("cpu", "cuda"):
         model = Sliced(device)
         index_map = CheckpointIndexMap(model, tensors)
+        # "v" through a kernel layout that reverses its rows
+        index_map.map_write("v", 0, row_offsets=(5 - torch.arange(6)) * 10)
         for name, tensor in checkpoint.items():
             index_map.load_rows(name, tensor.view(-1, tensors[name][0][-1]).to(device), 0)
         index_map.checksums()

From 7d35cf43cd501b157fffd4f5789be361b1dad4b4 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Tue, 3 Nov 2026 01:23:43 +0000
Subject: [PATCH 82/88] perf(weights): keep the static half of the TP checksum
 check across updates

Positions, counts and how evenly the TP group covers each tensor depend only on
the index map's checksum plan, so they (and three of the four reductions) are
computed once; an update reduces only the sums and the mismatch flag.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../inference/vllm/worker/filesystem.py       | 52 +++++++++++++------
 1 file changed, 36 insertions(+), 16 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index b0a8af7aad4..89ba06e11c9 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -356,30 +356,17 @@ def _check(self, label: str, names: list[str], totals: torch.Tensor, held: torch
         of times each (TP slices: once; KV heads replicated over TP workers: once per replica); one held by workers
         that overlap unevenly (fp8 tensors read in 128-row blocks) is checked only where a worker holds it whole."""
         start = time.perf_counter()
-        assert self.index_map is not None
-        canonical = sorted(self.index_map.numel)
-        position = {name: index for index, name in enumerate(canonical)}
+        canonical, ids, counts, whole, multiplicity = self._coverage(names, held)
         device = totals.device
-        ids = torch.tensor([position[name] for name in names], dtype=torch.long, device=device)
         sums = torch.zeros(len(canonical), dtype=torch.int64, device=device).index_put_((ids,), totals)
-        counts = torch.zeros(len(canonical), dtype=torch.int64, device=device).index_put_((ids,), held)
-        numel = torch.tensor([self.index_map.numel[name] for name in canonical], device=device)
         has_expected = torch.tensor([name in expected for name in canonical], device=device)
         target = torch.tensor([expected.get(name, 0) for name in canonical], dtype=torch.int64, device=device)
-        whole = counts == numel
         mismatched = whole & has_expected & (sums != target)
         tp_group = get_tp_group()
         if tp_group.world_size > 1:
             group = tp_group.device_group
-            total_sums, total_counts, most = sums.clone(), counts.clone(), counts.clone()
-            # Holders' fewest elements; non-holders report the most anyone could hold
-            fewest = torch.where(counts > 0, counts, numel)
-            torch.distributed.all_reduce(total_sums, group=group)
-            torch.distributed.all_reduce(total_counts, group=group)
-            torch.distributed.all_reduce(most, op=torch.distributed.ReduceOp.MAX, group=group)
-            torch.distributed.all_reduce(fewest, op=torch.distributed.ReduceOp.MIN, group=group)
-            even = (total_counts > 0) & (total_counts % numel == 0) & (most == fewest)
-            mismatched |= even & has_expected & (total_sums != total_counts // numel * target)
+            torch.distributed.all_reduce(sums, group=group)
+            mismatched |= (multiplicity > 0) & has_expected & (sums != multiplicity * target)
             any_mismatch = mismatched.any().to(torch.int32)
             torch.distributed.all_reduce(any_mismatch, op=torch.distributed.ReduceOp.MAX, group=group)
         else:
@@ -389,6 +376,39 @@ def _check(self, label: str, names: list[str], totals: torch.Tensor, held: torch
             raise RuntimeError(f"Weights at {label} diverged from the trainer: {names_off or 'on a TP peer'}")
         logger.info(f"Verified the weights at {label} in {time.perf_counter() - start:.2f}s")
 
+    def _coverage(self, names: list[str], held: torch.Tensor) -> tuple:
+        """What ``_check`` needs of who holds what: every checkpoint tensor of the run (sorted), the positions of
+        ``names`` in it, this worker's element counts and whole tensors, and how many times the TP group's holders
+        cover each tensor evenly (0: unevenly, or not held). Static for an index map, so kept for the ``names`` and
+        ``held`` of its checksum plan; every TP worker follows the same plan, so they compute it (and run its
+        collectives) at the same updates."""
+        cached = getattr(self, "coverage_cache", None)
+        if cached is not None and cached[0] is names and cached[1] is held:
+            return cached[2]
+        assert self.index_map is not None
+        canonical = sorted(self.index_map.numel)
+        position = {name: index for index, name in enumerate(canonical)}
+        device = held.device
+        ids = torch.tensor([position[name] for name in names], dtype=torch.long, device=device)
+        counts = torch.zeros(len(canonical), dtype=torch.int64, device=device).index_put_((ids,), held)
+        numel = torch.tensor([self.index_map.numel[name] for name in canonical], device=device)
+        whole = counts == numel
+        multiplicity = torch.zeros_like(counts)
+        tp_group = get_tp_group()
+        if tp_group.world_size > 1:
+            group = tp_group.device_group
+            total_counts, most = counts.clone(), counts.clone()
+            # Holders' fewest elements; non-holders report the most anyone could hold
+            fewest = torch.where(counts > 0, counts, numel)
+            torch.distributed.all_reduce(total_counts, group=group)
+            torch.distributed.all_reduce(most, op=torch.distributed.ReduceOp.MAX, group=group)
+            torch.distributed.all_reduce(fewest, op=torch.distributed.ReduceOp.MIN, group=group)
+            even = (total_counts > 0) & (total_counts % numel == 0) & (most == fewest)
+            multiplicity = torch.where(even, total_counts // numel, 0)
+        coverage = (canonical, ids, counts, whole, multiplicity)
+        self.coverage_cache = (names, held, coverage)
+        return coverage
+
     def _load_checkpoint(self, step_dir: Path) -> tuple | None:
         """Reload every weight from the HF checkpoint (or anchor) in ``step_dir``. In a sparse run, each worker reads
         only the rows it holds, straight into its params, and returns their check for ``_check``."""

From ef154483f9077e4fa83d32d8ee907b9179a52aa9 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Tue, 3 Nov 2026 02:14:10 +0000
Subject: [PATCH 83/88] perf(weights): one compiled sparse-apply variant per
 bits dtype

Triton specialized the kernel on its sizes and pointer alignments, so a delta with
new ones compiled inside the weight-update pause (0.7-2 s apply spikes on an
MXFP8 Qwen3-0.6B engine, against 0.02 s otherwise). do_not_specialize keeps the
variants the index map compiles ahead of time.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../inference/vllm/worker/sparse_apply.py     | 20 ++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/prime_rl/inference/vllm/worker/sparse_apply.py b/src/prime_rl/inference/vllm/worker/sparse_apply.py
index c379d13cf61..9e79cb9d7de 100644
--- a/src/prime_rl/inference/vllm/worker/sparse_apply.py
+++ b/src/prime_rl/inference/vllm/worker/sparse_apply.py
@@ -19,7 +19,25 @@
 _TL_BITS = {torch.uint8: tl.uint8, torch.int16: tl.int16, torch.int32: tl.int32, torch.int64: tl.int64}
 
 
-@triton.jit
+# Triton otherwise compiles a variant per value class of these (== 1, divisible by 16, pointer alignment), which a new
+# delta can hit inside a weight-update pause: one variant per bits dtype serves every delta.
+_RUNTIME_ARGS = (
+    "indices_ptr",
+    "deltas_ptr",
+    "contributions_ptr",
+    "entry_ends_ptr",
+    "entry_names_ptr",
+    "writes_ptr",
+    "num_cols_ptr",
+    "param_ptrs_ptr",
+    "maps_ptr",
+    "num_changes",
+    "num_entries",
+    "search_steps",
+)
+
+
+@triton.jit(do_not_specialize=_RUNTIME_ARGS)
 def _sparse_apply_kernel(
     indices_ptr,
     deltas_ptr,

From 45f15415b1b37914223ac1fdc96e039aaf355ce4 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Tue, 3 Nov 2026 03:18:03 +0000
Subject: [PATCH 84/88] feat(weights): log where a sparse delta's apply time
 goes

CUDA events split each chunk into copies, index decode and patch, next to the
host time of the loop.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../inference/vllm/worker/filesystem.py       | 33 +++++++++++++++----
 1 file changed, 27 insertions(+), 6 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 89ba06e11c9..064b477c7dd 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -268,12 +268,23 @@ def _boot_dir(self) -> Path:
         return Path(folder)
 
     def _apply_chunk(
-        self, entries: list[tuple[str, int]], gaps: torch.Tensor, values: torch.Tensor, tensors: dict, device
+        self,
+        entries: list[tuple[str, int]],
+        gaps: torch.Tensor,
+        values: torch.Tensor,
+        tensors: dict,
+        device,
+        events: list[torch.cuda.Event],
     ) -> None:
+        """Apply one chunk of a delta, recording CUDA events after its copies, its index decode and its apply."""
         assert self.index_map is not None
+        gaps, values = gaps.to(device, non_blocking=True), values.to(device, non_blocking=True)
+        events.append(torch.cuda.Event(enable_timing=True))
+        events[-1].record()
         counts = torch.tensor([count for _, count in entries], device=device)
-        indices = decode_gaps(gaps.to(device, non_blocking=True), counts)
-        values = values.to(device, non_blocking=True)
+        indices = decode_gaps(gaps, counts)
+        events.append(torch.cuda.Event(enable_timing=True))
+        events[-1].record()
         offset = num_bytes = 0
         for bits, group in itertools.groupby(entries, key=lambda entry: bits_dtype(tensors[entry[0]][1])):
             run = list(group)
@@ -285,6 +296,8 @@ def _apply_chunk(
             self.index_map.apply_delta(run, indices[offset : offset + num_changes], delta)
             offset += num_changes
             num_bytes += num_changes * itemsize
+        events.append(torch.cuda.Event(enable_timing=True))
+        events[-1].record()
 
     def _build_index_map(self, step_dir: Path) -> None:
         """At a sparse run's first update, while the engine still holds its boot weights: map checkpoint tensors to
@@ -337,13 +350,21 @@ def _apply_delta(self, step_dir: Path) -> dict[str, int]:
         tensors = read_manifest(step_dir)["tensors"]
         model = self._model()
         device = next(model.parameters()).device
+        events = [torch.cuda.Event(enable_timing=True)]
+        events[0].record()
         for entries, gaps, values in delta:
-            self._apply_chunk(entries, gaps, values, tensors, device)
+            self._apply_chunk(entries, gaps, values, tensors, device, events)
+        host_time = time.perf_counter() - start - read_time
         update_mla_absorbed_weights(model)
         torch.cuda.synchronize()
+        # Per chunk: copies, index decode, apply
+        phases = [0.0, 0.0, 0.0]
+        for index, (begin, end) in enumerate(zip(events, events[1:])):
+            phases[index % 3] += begin.elapsed_time(end) / 1000
         logger.info(
-            f"Applied sparse weight delta {step_dir.name}: {sum(len(gaps) for _, gaps, _ in delta)} changes, "
-            f"read {read_time:.2f}s (0 if prefetched), apply {time.perf_counter() - start - read_time:.2f}s"
+            f"Applied sparse weight delta {step_dir.name}: {sum(len(gaps) for _, gaps, _ in delta)} changes in "
+            f"{len(delta)} chunks, read {read_time:.2f}s (0 if prefetched), apply {time.perf_counter() - start - read_time:.2f}s "
+            f"(GPU: copy {phases[0]:.3f}s, decode {phases[1]:.3f}s, patch {phases[2]:.3f}s; host loop {host_time:.3f}s)"
         )
         return checksums
 

From 05d6f00f4c5101fe7c41dc1c77741171f48e70d1 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Tue, 3 Nov 2026 06:56:43 +0000
Subject: [PATCH 85/88] fix(trainer): per-token fp8 casts multiply by the
 reciprocal, like vLLM 0.33

vllm-project/vllm#61343 quantizes per-token groups as x * (1/s) instead of x / s,
so the trainer's Triton casts (meant to be bit-identical to vLLM's) differed in
0.03-0.08% of fp8 values. Multiply by the exact reciprocal (div_rn(1, s)) on the
non-ue8m0 path; test_fp8_utils passes again (25/25 on an RTX PRO 6000).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/prime_rl/trainer/models/kernels/fp8_utils.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/prime_rl/trainer/models/kernels/fp8_utils.py b/src/prime_rl/trainer/models/kernels/fp8_utils.py
index bcce6db3f90..6dceaed8f1e 100644
--- a/src/prime_rl/trainer/models/kernels/fp8_utils.py
+++ b/src/prime_rl/trainer/models/kernels/fp8_utils.py
@@ -146,7 +146,8 @@ def _per_token_fp8_kernel(
         scale = tl.exp2(tl.ceil(tl.log2(scale)))
         y = tl.clamp(x * (1.0 / scale)[:, None], FP8_MIN, FP8_MAX)
     else:
-        y = tl.clamp(tl.math.div_rn(x, scale[:, None]), FP8_MIN, FP8_MAX)
+        # vLLM's per-token-group quant multiplies by the exact reciprocal (vllm-project/vllm#61343)
+        y = tl.clamp(x * tl.math.div_rn(1.0, scale)[:, None], FP8_MIN, FP8_MAX)
     tl.store(
         out_ptr + row_offsets_i64[:, None] * stride_ym + col_offsets_i64[None, :] * stride_yn,
         y.to(tl.float8e4nv),
@@ -203,7 +204,8 @@ def _grouped_per_token_fp8_kernel(
         scale = tl.exp2(tl.ceil(tl.log2(scale)))
         y = tl.clamp(x * (1.0 / scale)[:, None], FP8_MIN, FP8_MAX)
     else:
-        y = tl.clamp(tl.math.div_rn(x, scale[:, None]), FP8_MIN, FP8_MAX)
+        # vLLM's per-token-group quant multiplies by the exact reciprocal (vllm-project/vllm#61343)
+        y = tl.clamp(x * tl.math.div_rn(1.0, scale)[:, None], FP8_MIN, FP8_MAX)
     tl.store(
         out_ptr + dst_rows_i64[:, None] * stride_ym + col_offsets_i64[None, :] * stride_yn,
         y.to(tl.float8e4nv),
@@ -258,7 +260,8 @@ def _grouped_per_channel_fp8_kernel(
         scale = tl.exp2(tl.ceil(tl.log2(scale)))
         y = tl.clamp(x * (1.0 / scale)[None, :], FP8_MIN, FP8_MAX)
     else:
-        y = tl.clamp(tl.math.div_rn(x, scale[None, :]), FP8_MIN, FP8_MAX)
+        # vLLM's per-token-group quant multiplies by the exact reciprocal (vllm-project/vllm#61343)
+        y = tl.clamp(x * tl.math.div_rn(1.0, scale)[None, :], FP8_MIN, FP8_MAX)
     flat_base = block_start.to(tl.int64) * BLOCK_K * cols
     if K_MAJOR:
         out_ptrs = out_ptr + flat_base + col_offsets_i64[:, None] * aligned_m + row_offsets_i64[None, :]

From 9280e96d028f8668593a9f4e97dcca2f81a69d66 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Tue, 3 Nov 2026 13:06:13 +0000
Subject: [PATCH 86/88] fix(weights): revalidate before probing a delta; check
 loaded anchors from the params

The update plan probed the delta manifest without revalidating the step dir,
so a stale NFS lookup cache could turn a sparse step into a full load. Anchor
and full-checkpoint loads checked the rows read from disk, not the params
after load_rows; they now check the live params like the delta path.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../inference/vllm/worker/filesystem.py       | 19 +++++--------------
 1 file changed, 5 insertions(+), 14 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 064b477c7dd..bd5067f8569 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -23,7 +23,6 @@
     ENGINE_LAYOUT,
     MANIFEST,
     anchor_complete,
-    as_bits,
     bits_dtype,
     checkpoint_index,
     decode_gaps,
@@ -32,7 +31,7 @@
     read_checkpoint_range,
     read_delta,
     read_manifest,
-    rows_checksum,
+    revalidate,
     unpin,
     write_file,
 )
@@ -203,6 +202,8 @@ def update_weights_from_path(self, weight_path: str) -> None:
 
     def _plan_update(self, step_dir: Path) -> dict:
         """What every TP worker does to move to ``step_dir``: a dict of plain values, broadcast from TP rank 0."""
+        # A stale NFS lookup cache could hide a manifest written before ``.finished``, turning a delta into a full load
+        revalidate(step_dir / DELTA_DIR)
         has_delta = (step_dir / DELTA_DIR / MANIFEST).exists()
         base_step = read_manifest(step_dir)["base_step"] if has_delta else None
         plan = {
@@ -452,8 +453,6 @@ def _load_checkpoint(self, step_dir: Path) -> tuple | None:
         device = next(model.parameters()).device
         # Checksums of anchors and full checkpoints written by the trainer (none for the boot checkpoint)
         expected = read_anchor_checksums(step_dir)
-        sums: dict[str, torch.Tensor] = {}
-        held: dict[str, int] = {}
 
         def read(name: str, begin: int, end: int) -> torch.Tensor:
             num_cols, dtype = index_map.num_cols[name], index[name][1]
@@ -469,19 +468,11 @@ def read(name: str, begin: int, end: int) -> torch.Tensor:
                 for (name, (begin, end)), rows in zip(window, pool.map(lambda item: read(item[0], *item[1]), window)):
                     rows = rows.to(device)
                     index_map.load_rows(name, rows, begin)
-                    if name in expected:
-                        num_cols = index_map.num_cols[name]
-                        sums[name] = rows_checksum(as_bits(rows).reshape(-1), begin * num_cols)
-                        held[name] = (end - begin) * num_cols
                     num_bytes += rows.nbytes
         update_mla_absorbed_weights(model)
         torch.cuda.synchronize()
         logger.info(
             f"Loaded this worker's {num_bytes / 2**30:.2f} GiB of {step_dir.name} in {time.perf_counter() - start:.1f}s"
         )
-        names = sorted(sums)
-        totals = (
-            torch.stack([sums[name] for name in names]) if names else torch.zeros(0, dtype=torch.int64, device=device)
-        )
-        counts = torch.tensor([held[name] for name in names], dtype=torch.int64, device=device)
-        return step_dir.name, names, totals, counts, expected
+        # Checked from the params as loaded, so a wrong mapping or write in ``load_rows`` fails the check
+        return step_dir.name, *index_map.held_checksums(), expected

From 1de91f164685f7e371990043433da95fd27b2ef3 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 4 Nov 2026 05:01:58 +0000
Subject: [PATCH 87/88] fix(weights): match the index map to the engine's boot
 layout

vLLM can move blocks within a param after loading: FlashInfer CUTLASS
(unquantized MoE, e.g. on SM120) swaps each expert's fused w13 halves to
[up; gate]. The index map, recorded from the loader's writes, kept them at
[gate; up], so anchor loads and replays served swapped experts while the
checks (which read through the index map) passed.

At the index map's build, while the engine holds its boot weights, compare
the bytes at every write against the boot checkpoint's rows (fp8: either
scale format; casts like the loader). A write whose block sits at another
write's place in the same param takes that place; a block found nowhere, or
a layout that is not a permutation of the writes, raises.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../inference/vllm/worker/filesystem.py       | 50 +++++++-----
 .../inference/vllm/worker/weight_transfer.py  | 78 ++++++++++++++++++-
 .../inference/test_checkpoint_index_map.py    | 48 ++++++++++++
 3 files changed, 155 insertions(+), 21 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index bd5067f8569..8e1606dfae8 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -4,7 +4,7 @@
 from concurrent.futures import Future, ThreadPoolExecutor
 from math import prod
 from pathlib import Path
-from typing import TYPE_CHECKING, NamedTuple
+from typing import TYPE_CHECKING, Iterator, NamedTuple
 
 import numpy as np
 import torch
@@ -302,7 +302,8 @@ def _apply_chunk(
 
     def _build_index_map(self, step_dir: Path) -> None:
         """At a sparse run's first update, while the engine still holds its boot weights: map checkpoint tensors to
-        the params, and find which ones this worker holds in fp8 (and their scale format)."""
+        the params (where the boot weights are, after the engine's processing), and find which ones this worker holds
+        in fp8 (and their scale format)."""
         if self.index_map is not None:
             return
         if (step_dir / DELTA_DIR / MANIFEST).exists():
@@ -314,7 +315,11 @@ def _build_index_map(self, step_dir: Path) -> None:
             }
         model = self._model()
         index_map = CheckpointIndexMap(model, tensors)
-        boot = checkpoint_index(self._boot_dir())
+        boot_dir = self._boot_dir()
+        moved = index_map.match_boot(_read_rows(index_map, boot_dir, next(model.parameters()).device))
+        if moved:
+            logger.info(f"Moved {moved} index map writes to where the engine's kernel layout holds them")
+        boot = checkpoint_index(boot_dir)
 
         def read_tensor(name: str) -> torch.Tensor:
             shape, dtype, _ = boot[name]
@@ -449,26 +454,13 @@ def _load_checkpoint(self, step_dir: Path) -> tuple | None:
             load_weights_checkpoint_layerwise(model, weights_iterator, self.model_runner.model_config, self.vllm_config)
             return None
         start = time.perf_counter()
-        index_map, index = self.index_map, checkpoint_index(step_dir)
-        device = next(model.parameters()).device
+        index_map = self.index_map
         # Checksums of anchors and full checkpoints written by the trainer (none for the boot checkpoint)
         expected = read_anchor_checksums(step_dir)
-
-        def read(name: str, begin: int, end: int) -> torch.Tensor:
-            num_cols, dtype = index_map.num_cols[name], index[name][1]
-            data = np.empty((end - begin) * num_cols * dtype.itemsize, np.uint8)
-            read_checkpoint_range(index[name], begin * num_cols, end * num_cols, data)
-            return torch.from_numpy(data).view(dtype).view(end - begin, num_cols)
-
-        ranges = list(index_map.row_ranges().items())
         num_bytes = 0
-        with ThreadPoolExecutor(16) as pool:
-            # A bounded window of reads ahead of the GPU writes
-            for window in itertools.batched(ranges, 64):
-                for (name, (begin, end)), rows in zip(window, pool.map(lambda item: read(item[0], *item[1]), window)):
-                    rows = rows.to(device)
-                    index_map.load_rows(name, rows, begin)
-                    num_bytes += rows.nbytes
+        for name, begin, rows in _read_rows(index_map, step_dir, next(model.parameters()).device):
+            index_map.load_rows(name, rows, begin)
+            num_bytes += rows.nbytes
         update_mla_absorbed_weights(model)
         torch.cuda.synchronize()
         logger.info(
@@ -476,3 +468,21 @@ def read(name: str, begin: int, end: int) -> torch.Tensor:
         )
         # Checked from the params as loaded, so a wrong mapping or write in ``load_rows`` fails the check
         return step_dir.name, *index_map.held_checksums(), expected
+
+
+def _read_rows(index_map: CheckpointIndexMap, step_dir: Path, device) -> Iterator[tuple[str, int, torch.Tensor]]:
+    """``(name, begin, rows)``: the ``row_ranges`` rows of each checkpoint tensor in ``step_dir``, on ``device``."""
+    index = checkpoint_index(step_dir)
+
+    def read(name: str, begin: int, end: int) -> torch.Tensor:
+        num_cols, dtype = index_map.num_cols[name], index[name][1]
+        data = np.empty((end - begin) * num_cols * dtype.itemsize, np.uint8)
+        read_checkpoint_range(index[name], begin * num_cols, end * num_cols, data)
+        return torch.from_numpy(data).view(dtype).view(end - begin, num_cols)
+
+    ranges = list(index_map.row_ranges().items())
+    with ThreadPoolExecutor(16) as pool:
+        # A bounded window of reads ahead of the GPU writes
+        for window in itertools.batched(ranges, 64):
+            for (name, (begin, _)), rows in zip(window, pool.map(lambda item: read(item[0], *item[1]), window)):
+                yield name, begin, rows.to(device)
diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py
index 8e9f0555b2f..7480d2f0bea 100644
--- a/src/prime_rl/inference/vllm/worker/weight_transfer.py
+++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py
@@ -345,6 +345,68 @@ def fp8_layout(self, model: Module, read_tensor) -> dict[str, str]:
         self.checksum_plan = self.running = self.fused = None
         return self.fp8
 
+    @torch.no_grad()
+    def match_boot(self, boot_rows: Iterable[tuple[str, int, torch.Tensor]]) -> int:
+        """Point every write at where the params hold its block, from the boot checkpoint's ``row_ranges`` rows
+        (``(name, begin, rows)`` on the params' device) while the params still hold them. Kernel processing after
+        loading can move blocks within a param (FlashInfer CUTLASS swaps the gate and up halves of fused expert
+        weights); a write whose block is at another write's place in the same param takes that place. Raises if a
+        block is nowhere. Before ``fp8_layout``; returns how many writes moved."""
+        groups: dict[tuple[int, int, int], list[tuple[str, int]]] = defaultdict(list)
+        for name, writes in self.records.items():
+            for k, (_, rows, _, cols, index, *_) in enumerate(writes):
+                groups[(index, rows, cols)].append((name, k))
+        hashes: dict[tuple[int, int, int], dict[int, list[tuple[str, int]]]] = {}
+        moves: dict[tuple[str, int], tuple[str, int]] = {}
+        taken: set[tuple[str, int]] = set()
+        for name, begin, rows in boot_rows:
+            sources = {}
+            for k, (row0, num_rows, col0, cols, index, *_) in enumerate(self.records[name]):
+                dtype = self.params[index].dtype
+                if dtype not in sources:
+                    sources[dtype] = _boot_bits(rows, dtype)
+                expected = [bits[row0 - begin : row0 - begin + num_rows, col0 : col0 + cols] for bits in sources[dtype]]
+                if any(torch.equal(block, self._live(name, k)) for block in expected):
+                    continue
+                group = (index, num_rows, cols)
+                if group not in hashes:
+                    hashes[group] = defaultdict(list)
+                    for slot in groups[group]:
+                        hashes[group][_block_hash(self._live(*slot))].append(slot)
+                found = [
+                    slot
+                    for block in expected
+                    for slot in hashes[group].get(_block_hash(block), ())
+                    if slot not in taken and torch.equal(block, self._live(*slot))
+                ]
+                if not found:
+                    raise NotImplementedError(
+                        f"Sparse weight deltas cannot patch {name}: the engine's {self.param_names[index]} does not "
+                        "hold the boot checkpoint's block of it anywhere a write of the index map goes"
+                    )
+                moves[(name, k)] = found[0]
+                taken.add(found[0])
+        # The moved writes trade places among themselves; taking an unmoved write's place would put two writes there
+        if taken != set(moves):
+            raise NotImplementedError(
+                f"Sparse weight deltas cannot patch {sorted({name for name, _ in moves})[:8]}: the engine's params "
+                "hold their blocks in a layout that is not a permutation of the index map's writes"
+            )
+        places = {slot: (self.records[slot[0]][slot[1]][4:], self.maps.get(slot)) for slot in moves}
+        for (name, k), slot in moves.items():
+            place, tables = places[slot]
+            self.records[name][k] = (*self.records[name][k][:4], *place)
+            self.maps.pop((name, k), None)
+            if tables is not None:
+                self.maps[(name, k)] = tables
+        self.map_cache = self.checksum_plan = self.running = self.fused = None
+        return len(moves)
+
+    def _live(self, name: str, k: int) -> torch.Tensor:
+        """The raw bits the params hold at write ``k`` of ``name``, ``[rows, cols]``."""
+        index = self.records[name][k][4]
+        return as_bits(self.params[index].data).view(-1)[self._targets(name, k)]
+
     def _view(self, index: int, rows: int, cols: int, offset: int, strides) -> torch.Tensor:
         """The raw bits of one write of a checkpoint block into param ``index``."""
         return as_bits(self.params[index].data).view(-1).as_strided((rows, cols), tuple(strides), offset)
@@ -357,7 +419,7 @@ def row_ranges(self) -> dict[str, tuple[int, int]]:
             if name.endswith("_scale_inv"):
                 continue  # scales are set with their weights
             begin, end = min(write[0] for write in writes), max(write[0] + write[1] for write in writes)
-            if name in self.fp8:
+            if any(self.params[write[4]].dtype == torch.float8_e4m3fn for write in writes):
                 num_rows = self.numel[name] // self.num_cols[name]
                 begin, end = begin // FP8_BLOCK[0] * FP8_BLOCK[0], min(-(-end // FP8_BLOCK[0]) * FP8_BLOCK[0], num_rows)
             ranges[name] = (begin, end)
@@ -554,6 +616,20 @@ def _scatter_add(
                 self.running.index_add_(0, part_ids[keep], change[keep])
 
 
+def _boot_bits(rows: torch.Tensor, dtype: torch.dtype) -> list[torch.Tensor]:
+    """The raw bits ``load_rows`` writes for checkpoint ``rows`` into a param of ``dtype``: fp8 in either scale
+    format (``fp8_layout`` picks one later), others cast like the engine's loader."""
+    if dtype == torch.float8_e4m3fn:
+        return [quantize_fp8_blocks(rows, ue8m0)[0] for ue8m0 in (True, False)]
+    return [as_bits(rows.to(dtype))]
+
+
+def _block_hash(bits: torch.Tensor) -> int:
+    """A position-weighted sum of a block's raw bits, to find a block among a param's writes."""
+    flat = bits.reshape(-1).long()
+    return int((flat * torch.arange(1, flat.numel() + 1, device=flat.device)).sum())
+
+
 def _restore_checkpoint_format(model: Module) -> list[tuple]:
     """Put the checkpoint-format params (on meta, with their unwrapped loaders) back into
     every layer whose live params were transformed after loading (online quantization)."""
diff --git a/tests/unit/inference/test_checkpoint_index_map.py b/tests/unit/inference/test_checkpoint_index_map.py
index 1e869a981b3..62b3cfb32a1 100644
--- a/tests/unit/inference/test_checkpoint_index_map.py
+++ b/tests/unit/inference/test_checkpoint_index_map.py
@@ -136,3 +136,51 @@ def test_fused_apply_matches_torch_apply():
     running = gpu_map.checksums()[1].cpu()
     assert running.tolist() == cpu_map.checksums()[1].tolist()
     assert running.tolist() == gpu_map.checksums(full=True)[1].cpu().tolist()
+
+
+class FusedExperts(nn.Module):
+    """Two experts' gate and up projections loaded into one fused ``[experts, gate; up, hidden]`` param, then, like
+    FlashInfer CUTLASS's processing after loading, swapped to ``[up; gate]``."""
+
+    def __init__(self):
+        super().__init__()
+        self.w13 = nn.Parameter(torch.zeros(2, 8, 4, dtype=torch.bfloat16))
+
+    def load_weights(self, weights):
+        for name, tensor in weights:
+            expert, projection = int(name.split(".")[1]), name.split(".")[2]
+            self.w13.data[expert, (0 if projection == "gate" else 4) : (4 if projection == "gate" else 8)].copy_(tensor)
+
+    def swap_halves(self):
+        self.w13.data.copy_(self.w13.data.view(2, 2, 4, 4).flip(1).view(2, 8, 4))
+
+
+def fused_experts_boot():
+    tensors = {f"experts.{e}.{p}.weight": ([4, 4], "bfloat16") for e in range(2) for p in ("gate", "up")}
+    checkpoint = {name: torch.randn(4, 4, dtype=torch.bfloat16) for name in tensors}
+    model = FusedExperts()
+    model.load_weights(checkpoint.items())
+    model.swap_halves()
+    return model, tensors, checkpoint
+
+
+def test_match_boot_follows_a_post_load_gate_up_swap():
+    model, tensors, checkpoint = fused_experts_boot()
+    index_map = CheckpointIndexMap(model, tensors)
+    boot = model.w13.data.clone()
+    assert index_map.match_boot((name, 0, tensor) for name, tensor in checkpoint.items()) == 4
+    # Moved writes stay put on a second check, and a load through the index map keeps the engine's layout
+    assert index_map.match_boot((name, 0, tensor) for name, tensor in checkpoint.items()) == 0
+    for name, tensor in checkpoint.items():
+        index_map.load_rows(name, tensor, 0)
+    assert torch.equal(model.w13.data, boot)
+    names, totals, _ = index_map.checksums()
+    assert totals.tolist() == [bits_checksum(as_bits(checkpoint[name])) for name in names]
+
+
+def test_match_boot_raises_on_a_block_it_cannot_find():
+    model, tensors, checkpoint = fused_experts_boot()
+    index_map = CheckpointIndexMap(model, tensors)
+    model.w13.data[1, 0, 0] += 1
+    with pytest.raises(NotImplementedError, match="experts.1.up.weight"):
+        index_map.match_boot((name, 0, tensor) for name, tensor in checkpoint.items())

From c70aa3c342e1eb76fdbda6102321c324489265ee Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Mon, 9 Nov 2026 14:00:57 +0000
Subject: [PATCH 88/88] fix(weights): skip tied names a checkpoint omits when
 reading rows

The index map records a write for every name the engine's loader consumes, including a tied weight's
second name (vLLM's Qwen3 loads `lm_head.weight` into the tied embedding). HF checkpoints of tied models
may omit it (PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT does), so the first sparse update's boot read
(and a replay from the boot checkpoint) raised a KeyError. Names that write the same elements are now
known as tied; `_read_rows` skips one the checkpoint omits when it has a tied name, and still raises
for any other missing name.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../inference/vllm/worker/filesystem.py       | 10 ++++--
 .../inference/vllm/worker/weight_transfer.py  | 11 +++++++
 .../inference/test_checkpoint_index_map.py    | 33 +++++++++++++++++++
 3 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 8e1606dfae8..44b6903ebad 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -471,8 +471,15 @@ def _load_checkpoint(self, step_dir: Path) -> tuple | None:
 
 
 def _read_rows(index_map: CheckpointIndexMap, step_dir: Path, device) -> Iterator[tuple[str, int, torch.Tensor]]:
-    """``(name, begin, rows)``: the ``row_ranges`` rows of each checkpoint tensor in ``step_dir``, on ``device``."""
+    """``(name, begin, rows)``: the ``row_ranges`` rows of each checkpoint tensor in ``step_dir``, on ``device``. A
+    tied name the checkpoint omits is skipped: its elements come with a tied name it has."""
     index = checkpoint_index(step_dir)
+    ranges = []
+    for name, rows in index_map.row_ranges().items():
+        if name in index:
+            ranges.append((name, rows))
+        elif not any(tied in index for tied in index_map.tied.get(name, ())):
+            raise KeyError(f"{step_dir} has no {name}")
 
     def read(name: str, begin: int, end: int) -> torch.Tensor:
         num_cols, dtype = index_map.num_cols[name], index[name][1]
@@ -480,7 +487,6 @@ def read(name: str, begin: int, end: int) -> torch.Tensor:
         read_checkpoint_range(index[name], begin * num_cols, end * num_cols, data)
         return torch.from_numpy(data).view(dtype).view(end - begin, num_cols)
 
-    ranges = list(index_map.row_ranges().items())
     with ThreadPoolExecutor(16) as pool:
         # A bounded window of reads ahead of the GPU writes
         for window in itertools.batched(ranges, 64):
diff --git a/src/prime_rl/inference/vllm/worker/weight_transfer.py b/src/prime_rl/inference/vllm/worker/weight_transfer.py
index 7480d2f0bea..eac0d34061a 100644
--- a/src/prime_rl/inference/vllm/worker/weight_transfer.py
+++ b/src/prime_rl/inference/vllm/worker/weight_transfer.py
@@ -135,6 +135,17 @@ def placeholders():
         self.unpatchable = unpatchable
         # Checkpoint tensors this worker holds; unpatchable ones too, so a change to them is caught in apply_delta
         self.names = set(self.records)
+        # Names that write the same elements (a tied weight's names, e.g. ``lm_head.weight`` of tied embeddings,
+        # which HF checkpoints may omit) -> the other names
+        by_writes = defaultdict(list)
+        for name, writes in self.records.items():
+            by_writes[tuple(writes)].append(name)
+        self.tied = {
+            name: [other for other in names if other != name]
+            for names in by_writes.values()
+            if len(names) > 1
+            for name in names
+        }
 
     def map_write(
         self, name: str, k: int, row_offsets: torch.Tensor | None = None, col_offsets: torch.Tensor | None = None
diff --git a/tests/unit/inference/test_checkpoint_index_map.py b/tests/unit/inference/test_checkpoint_index_map.py
index 62b3cfb32a1..55953aa8ce4 100644
--- a/tests/unit/inference/test_checkpoint_index_map.py
+++ b/tests/unit/inference/test_checkpoint_index_map.py
@@ -1,7 +1,9 @@
 import pytest
 import torch
+from safetensors.torch import save_file
 from torch import nn
 
+from prime_rl.inference.vllm.worker.filesystem import _read_rows
 from prime_rl.inference.vllm.worker.weight_transfer import CheckpointIndexMap
 from prime_rl.utils.sparse_weights import as_bits, bits_checksum
 
@@ -184,3 +186,34 @@ def test_match_boot_raises_on_a_block_it_cannot_find():
     model.w13.data[1, 0, 0] += 1
     with pytest.raises(NotImplementedError, match="experts.1.up.weight"):
         index_map.match_boot((name, 0, tensor) for name, tensor in checkpoint.items())
+
+
+class TiedToy(nn.Module):
+    """An engine model whose head is tied to its embedding: its loader loads both names into one param."""
+
+    def __init__(self):
+        super().__init__()
+        self.embed = nn.Parameter(torch.zeros(4, 8, dtype=torch.bfloat16))
+
+    def load_weights(self, weights):
+        for _, tensor in weights:
+            self.embed.data.copy_(tensor)
+
+
+def test_read_rows_skips_only_tied_names_the_checkpoint_omits(tmp_path):
+    model = TiedToy()
+    index_map = CheckpointIndexMap(model, {"embed": ([4, 8], "bfloat16"), "head": ([4, 8], "bfloat16")})
+    assert index_map.tied == {"embed": ["head"], "head": ["embed"]}
+    embed = torch.randn(4, 8, dtype=torch.bfloat16)
+    # A HF checkpoint of tied embeddings may store only the embedding
+    save_file({"embed": embed}, tmp_path / "model.safetensors")
+    model.embed.data.copy_(embed)
+    rows = list(_read_rows(index_map, tmp_path, "cpu"))
+    assert [(name, begin) for name, begin, _ in rows] == [("embed", 0)]
+    assert torch.equal(rows[0][2], embed)
+    assert index_map.match_boot(rows) == 0
+
+    index_map = CheckpointIndexMap(Toy(), {"weight": ([4, 8], "bfloat16"), "norm": ([8], "bfloat16")})
+    save_file({"weight": embed}, tmp_path / "model.safetensors")
+    with pytest.raises(KeyError, match="norm"):
+        list(_read_rows(index_map, tmp_path, "cpu"))