feat(weights): sparse XOR deltas for the filesystem weight broadcast

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 \
@@ -251,10 +252,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"

docs/scaling.md

@@ -291,6 +291,23 @@ 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.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 = "filesystem"
+
+[weight_broadcast.delta]
+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 `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
 
 For unusual partitions, module loads, or environment setup, supply your own Jinja2 template:

packages/prime-rl-configs/src/prime_rl/configs/shared.py

@@ -45,9 +45,22 @@ class BaseWeightBroadcastConfig(BaseConfig):
     offered version in time. Raise it for large models on slow shared filesystems."""
 
 
+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."""
+
+    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):
     type: Literal["filesystem"] = "filesystem"
 
+    delta: DeltaWeightBroadcastConfig | None = None
+    """Publish sparse deltas instead of full checkpoints."""
+
 
 class InMemoryWeightBroadcastConfig(BaseWeightBroadcastConfig):
     host: str = "localhost"

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."""
@@ -726,6 +726,8 @@ 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.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")

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-runs[train]>=0.1.6",
@@ -51,44 +51,41 @@ 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 = [
     "fastapi>=0.115",
     "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'"]
 
@@ -98,7 +95,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 = [
@@ -152,10 +149,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",
 ]
 
@@ -178,8 +176,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",
@@ -203,8 +201,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.
@@ -225,9 +223,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
@@ -251,7 +249,6 @@ deep-gemm = false
 nixl-cu13 = false
 torchao = false
 prime-kernels = false
-flash-mla = false
 deep-select = false
 
 [tool.uv.sources]
@@ -278,41 +275,39 @@ 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'" },
-    { 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'" },
 ]
-# 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" }

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 ==="