[diffusion] optimization: transfer mapped layers through a courier thread (#35882)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

python/sglang/multimodal_gen/envs.py

@@ -30,6 +30,9 @@ if TYPE_CHECKING:
     CMAKE_BUILD_TYPE: str | None = None
     VERBOSE: bool = False
     SGLANG_DIFFUSION_SERVER_DEV_MODE: bool = False
+    SGLANG_DIFFUSION_DISABLE_MAPPED_COURIER: bool = False
+    SGLANG_DIFFUSION_TEST_FORCE_HOST_AVAILABLE_GIB: float | None = None
+    SGLANG_DIFFUSION_TEST_CAP_DEVICE_MEMORY_GIB: float | None = None
     SGLANG_DIFFUSION_STAGE_LOGGING: bool = False
     SGLANG_DIFFUSION_CFG_GATE_STEP: float = 1.0
     # cache-dit env vars (primary transformer)
@@ -111,6 +114,14 @@ def _lazy_float(key: str, default: str | float) -> Callable[[], float]:
     return lambda: float(os.getenv(key, str(default)))
 
 
+def _lazy_optional_float(key: str) -> Callable[[], float | None]:
+    def _getter():
+        val = os.getenv(key)
+        return float(val) if val is not None else None
+
+    return _getter
+
+
 def _lazy_bool(key: str, default: str = "false") -> Callable[[], bool]:
     return lambda: get_bool_env_var(key, default)
 
@@ -222,6 +233,27 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # some additional endpoints for developing and debugging,
     # e.g. `/reset_prefix_cache`
     "SGLANG_DIFFUSION_SERVER_DEV_MODE": _lazy_bool("SGLANG_DIFFUSION_SERVER_DEV_MODE"),
+    # Kill-switch for the courier thread that ships checkpoint-mapped layers to
+    # the device off the compute thread. The courier already falls back to the
+    # synchronous copy on any failure; this forces that path up front.
+    "SGLANG_DIFFUSION_DISABLE_MAPPED_COURIER": _lazy_bool(
+        "SGLANG_DIFFUSION_DISABLE_MAPPED_COURIER"
+    ),
+    # Test hook: make the host memory budget behave as if the machine had this
+    # many GiB of RAM (available = this figure minus the process's own
+    # anonymous memory). CI uses it to exercise the constrained placement
+    # paths -- mapped weights, partial pinning, the courier -- on runners whose
+    # real hosts are never short of memory.
+    "SGLANG_DIFFUSION_TEST_FORCE_HOST_AVAILABLE_GIB": _lazy_optional_float(
+        "SGLANG_DIFFUSION_TEST_FORCE_HOST_AVAILABLE_GIB"
+    ),
+    # Test-only: cap the CUDA caching allocator at this many GiB, so a large
+    # CI card behaves like the consumer card a case is written for. Without
+    # the cap the allocator is free to reserve past the pretended budget and
+    # a peak-VRAM baseline stops meaning "fits the card".
+    "SGLANG_DIFFUSION_TEST_CAP_DEVICE_MEMORY_GIB": _lazy_optional_float(
+        "SGLANG_DIFFUSION_TEST_CAP_DEVICE_MEMORY_GIB"
+    ),
     # If set, sgl_diffusion will enable stage logging, which will print the time
     # taken for each stage
     "SGLANG_DIFFUSION_STAGE_LOGGING": _lazy_bool("SGLANG_DIFFUSION_STAGE_LOGGING"),

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -231,10 +231,31 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
             )
         return self.memory_occupation
 
+    def _cap_device_memory_for_tests(self) -> None:
+        """Make a large CI card behave like the consumer card a case targets.
+
+        The caching allocator otherwise reserves past the pretended budget
+        whenever the physical card has room, and a peak-VRAM baseline stops
+        meaning "fits the card". OOM inside the cap is the intended signal.
+        """
+        cap_gib = envs.SGLANG_DIFFUSION_TEST_CAP_DEVICE_MEMORY_GIB
+        if cap_gib is None or not current_platform.is_cuda():
+            return
+        device = torch.cuda.current_device()
+        total = torch.cuda.get_device_properties(device).total_memory
+        fraction = min(1.0, cap_gib * 1024**3 / total)
+        torch.cuda.set_per_process_memory_fraction(fraction, device)
+        logger.info(
+            "Test hook: CUDA allocator capped at %.1f GiB (fraction %.4f)",
+            cap_gib,
+            fraction,
+        )
+
     def init_device_and_model(self) -> None:
         """Initialize the device and load the model."""
         if not current_platform.is_mps():
             current_platform.set_device(current_platform.get_device(self.local_rank))
+        self._cap_device_memory_for_tests()
         # num_gpus is the total world size across every node; the co-located,
         # CPU-contending worker count on THIS host is num_gpus // nnodes.
         local_num_gpus = self.server_args.num_gpus // self.server_args.nnodes

python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py

@@ -17,6 +17,7 @@ import os
 
 import psutil
 
+from sglang.multimodal_gen import envs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)
@@ -130,6 +131,22 @@ def host_memory_available_bytes() -> int:
     The smaller of what the kernel reports free and what the cgroup still
     allows, so a container does not plan against the whole machine.
     """
+    forced_gib = envs.SGLANG_DIFFUSION_TEST_FORCE_HOST_AVAILABLE_GIB
+    if forced_gib is not None:
+        # Behave like a machine of that size: what such a host would still
+        # have free is the pretend total minus what this process has already
+        # taken in anonymous memory.
+        own_anonymous = 0
+        try:
+            with open("/proc/self/status") as handle:
+                for line in handle:
+                    if line.startswith("RssAnon:"):
+                        own_anonymous = int(line.split()[1]) * 1024
+                        break
+        except OSError:
+            pass
+        return max(0, int(forced_gib * GIB_BYTES) - own_anonymous)
+
     available = int(psutil.virtual_memory().available)
     capped = cgroup_memory_limit_bytes()
     if capped is None:

python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py

@@ -1,12 +1,15 @@
 import bisect
+import queue
 import re
+import threading
 from collections.abc import Mapping, Sequence
 from contextlib import nullcontext
-from typing import Any, Dict, List, Set, Tuple
+from typing import Any, Dict, List, Optional, Set, Tuple
 
 import torch
 from torch.distributed.tensor import DTensor
 
+from sglang.multimodal_gen import envs
 from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
 from sglang.multimodal_gen.runtime.loader.utils import MappedRegions
 from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
@@ -182,6 +185,147 @@ def _install_host_gather_hooks(
     module.register_forward_hook(_output_to_device)
 
 
+class MappedLayerCourier:
+    """Ships a mapped layer's weights to the device off the compute thread.
+
+    A copy whose source is a checkpoint mapping is synchronous however it is
+    requested -- the driver stages unpinned memory through its own buffer -- and
+    the prefetch hooks run on the compute thread, so every such copy stalls the
+    step. This worker thread reads the mapped bytes into a pinned slot (a plain
+    memcpy when the page cache holds them) and issues the device copy from
+    there on its own stream, where it is genuinely asynchronous. The compute
+    thread's large tensor copies release the GIL, so reading layer i+1 really
+    does overlap computing layer i.
+
+    The thread only prepares device tensors and records an event; parameters
+    are rebound on the compute thread at collect time, so module state is never
+    touched concurrently.
+    """
+
+    _NUM_SLOTS = 2
+
+    def __init__(
+        self,
+        *,
+        mapped_cpu_weights: Dict[int, Dict[str, torch.Tensor]],
+        weight_metadata: Dict[int, Dict[str, Dict[str, Any]]],
+        device: torch.device,
+        pin_slots: bool,
+    ) -> None:
+        self._mapped_cpu_weights = mapped_cpu_weights
+        self._weight_metadata = weight_metadata
+        self._device = device
+        slot_bytes = max(
+            (
+                sum(t.numel() * t.element_size() for t in weights.values())
+                for weights in mapped_cpu_weights.values()
+                if weights
+            ),
+            default=0,
+        )
+        if slot_bytes <= 0:
+            raise ValueError("no mapped weights to ship")
+        self._slots = [
+            torch.empty(slot_bytes, dtype=torch.uint8, pin_memory=pin_slots)
+            for _ in range(self._NUM_SLOTS)
+        ]
+        self._slot_events: List[Optional[Any]] = [None] * self._NUM_SLOTS
+        self._stream = torch.get_device_module().Stream()
+        self._tasks: queue.Queue[Optional[int]] = queue.Queue()
+        self._results: Dict[int, Any] = {}
+        self._ready = threading.Condition()
+        self._pending: Set[int] = set()
+        self._broken = False
+        self._thread = threading.Thread(
+            target=self._run, name="mapped-layer-courier", daemon=True
+        )
+        self._thread.start()
+
+    def submit(self, layer_idx: int) -> bool:
+        """Queue a layer. False when the courier is out of service."""
+        if self._broken:
+            return False
+        with self._ready:
+            if layer_idx in self._pending or layer_idx in self._results:
+                return True
+            self._pending.add(layer_idx)
+        self._tasks.put(layer_idx)
+        return True
+
+    def pending(self, layer_idx: int) -> bool:
+        with self._ready:
+            return layer_idx in self._pending or layer_idx in self._results
+
+    def collect(self, layer_idx: int):
+        """Block until the layer is shipped; (event, {name: gpu_tensor})."""
+        with self._ready:
+            while layer_idx not in self._results:
+                if self._broken and layer_idx not in self._results:
+                    raise RuntimeError("mapped-layer courier stopped")
+                self._ready.wait(timeout=1.0)
+            outcome = self._results.pop(layer_idx)
+        if isinstance(outcome, BaseException):
+            raise outcome
+        return outcome
+
+    def close(self) -> None:
+        self._tasks.put(None)
+        self._thread.join(timeout=5.0)
+
+    def _run(self) -> None:
+        slot_turn = 0
+        while True:
+            layer_idx = self._tasks.get()
+            if layer_idx is None:
+                return
+            try:
+                outcome = self._ship(layer_idx, slot_turn)
+                slot_turn = (slot_turn + 1) % self._NUM_SLOTS
+            except BaseException as exc:  # published, never swallowed
+                outcome = exc
+                self._broken = True
+            with self._ready:
+                self._pending.discard(layer_idx)
+                self._results[layer_idx] = outcome
+                self._ready.notify_all()
+            if self._broken:
+                return
+
+    def _ship(self, layer_idx: int, slot_turn: int):
+        slot = self._slots[slot_turn]
+        previous = self._slot_events[slot_turn]
+        if previous is not None:
+            # the previous transfer through this slot must land before reuse
+            previous.synchronize()
+        tensors: Dict[str, torch.Tensor] = {}
+        offset = 0
+        with torch.inference_mode(False), torch.no_grad():
+            staged = []
+            for name, cpu_tensor in self._mapped_cpu_weights[layer_idx].items():
+                width = cpu_tensor.element_size()
+                if offset % width:
+                    offset += width - (offset % width)
+                start = offset // width
+                window = slot.view(cpu_tensor.dtype)[
+                    start : start + cpu_tensor.numel()
+                ].view(cpu_tensor.shape)
+                window.copy_(cpu_tensor)
+                offset += cpu_tensor.numel() * width
+                staged.append((name, window))
+            event = torch.get_device_module().Event()
+            with torch.get_device_module().stream(self._stream):
+                for name, window in staged:
+                    meta = self._weight_metadata[layer_idx][name]
+                    gpu_tensor = torch.empty(
+                        meta["shape"], dtype=meta["dtype"], device=self._device
+                    )
+                    gpu_tensor.copy_(window, non_blocking=True)
+                    tensors[name] = gpu_tensor
+                event.record(self._stream)
+        self._slot_events[slot_turn] = event
+        return event, tensors
+
+
 class LayerwiseOffloadManager:
     """A lightweight layerwise CPU offload manager.
 
@@ -296,6 +440,9 @@ class LayerwiseOffloadManager:
         self._weight_metadata: Dict[int, Dict[str, Dict[str, Any]]] = {}
         # layer indices that are already in gpu
         self._gpu_layers: Set[int] = set()
+        # mapped layers handed to the courier and not yet collected
+        self._mapped_courier: Optional[MappedLayerCourier] = None
+        self._courier_inflight: Set[int] = set()
         # layer_idx -> torch.get_device_module().Event for fine-grained sync, to make sure the weight is resident in pre-hook
         self._prefetch_events: Dict[int, torch.get_device_module().Event] = {}
 
@@ -790,6 +937,11 @@ class LayerwiseOffloadManager:
             return
         if layer_idx in self._gpu_layers:
             return
+        if layer_idx in self._courier_inflight:
+            if non_blocking:
+                return
+            self._collect_mapped_layer(layer_idx)
+            return
         if self._synchronous_mps:
             cpu_weights = self._mps_cpu_weights.get(layer_idx)
             if not cpu_weights:
@@ -815,6 +967,17 @@ class LayerwiseOffloadManager:
             non_blocking = False
             stream_context = nullcontext()
 
+        # A mapped source is synchronous on this thread however the copy is
+        # requested, so hand those weights to the courier and let it overlap
+        # this layer's transfer with the previous layer's compute. Blocking
+        # callers keep the direct path: they need the weights now.
+        ship_mapped = False
+        if non_blocking and self._mapped_cpu_weights.get(layer_idx):
+            courier = self._ensure_mapped_courier()
+            if courier is not None and courier.submit(layer_idx):
+                self._courier_inflight.add(layer_idx)
+                ship_mapped = True
+
         # create gpu buffer and load from CPU buffer
         gpu_buffers: Dict[torch.dtype, torch.Tensor] = {}
         with (
@@ -836,8 +999,11 @@ class LayerwiseOffloadManager:
             for name, meta in self._weight_metadata[layer_idx].items():
                 target = self.get_target_with_name(name)
                 if meta.get("mapped", False):
+                    if ship_mapped:
+                        # the courier stages and ships these; bound at collect
+                        continue
                     # Straight from the mapping. Not pinned, so this copy runs
-                    # on the compute stream rather than ahead of it, and a page
+                    # on the compute thread rather than ahead of it, and a page
                     # the kernel has reclaimed is faulted back in here.
                     cpu_tensor = self._mapped_cpu_weights[layer_idx][name]
                     gpu_tensor = torch.empty(
@@ -878,6 +1044,67 @@ class LayerwiseOffloadManager:
             event.record(self.copy_stream)
             self._prefetch_events[layer_idx] = event
 
+        if not ship_mapped:
+            self._gpu_layers.add(layer_idx)
+
+    def _ensure_mapped_courier(self) -> Optional[MappedLayerCourier]:
+        """The courier, built on first use; None where it cannot help."""
+        if self._mapped_courier is not None:
+            return self._mapped_courier
+        if envs.SGLANG_DIFFUSION_DISABLE_MAPPED_COURIER:
+            return None
+        if self.copy_stream is None or self._synchronous_mps:
+            return None
+        if not self._mapped_bytes:
+            return None
+        try:
+            self._mapped_courier = MappedLayerCourier(
+                mapped_cpu_weights=self._mapped_cpu_weights,
+                weight_metadata=self._weight_metadata,
+                device=self.device,
+                pin_slots=current_platform.is_cuda(),
+            )
+            logger.info(
+                "Layerwise offload: %s ships mapped layers through a courier "
+                "thread with %d pinned slots, so their device copies overlap "
+                "compute instead of stalling it.",
+                self.layers_attr_str,
+                MappedLayerCourier._NUM_SLOTS,
+            )
+        except (RuntimeError, MemoryError, ValueError) as exc:
+            logger.info(
+                "Layerwise offload: no courier for mapped layers (%s); they "
+                "keep the synchronous copy.",
+                exc,
+            )
+            self._mapped_courier = None
+            self._mapped_bytes = self._mapped_bytes  # unchanged; direct path
+        return self._mapped_courier
+
+    def _collect_mapped_layer(self, layer_idx: int) -> None:
+        """Bind a shipped layer's tensors on the compute thread."""
+        courier = self._mapped_courier
+        try:
+            event, tensors = courier.collect(layer_idx)
+        except BaseException as exc:
+            # The courier is out of service: fall back to the direct
+            # synchronous path for this and every later layer.
+            logger.warning(
+                "Layerwise offload: courier failed for layer %d (%s); mapped "
+                "layers return to the synchronous copy.",
+                layer_idx,
+                exc,
+            )
+            self._mapped_courier = None
+            self._courier_inflight.discard(layer_idx)
+            self.prefetch_layer(layer_idx, non_blocking=False)
+            return
+        with torch.inference_mode(False), torch.no_grad():
+            for name, gpu_tensor in tensors.items():
+                target = self.get_target_with_name(name)
+                target.data = self._wrap_for_target(target, gpu_tensor)
+        torch.get_device_module().current_stream().wait_event(event)
+        self._courier_inflight.discard(layer_idx)
         self._gpu_layers.add(layer_idx)
 
     @torch.compiler.disable
@@ -924,6 +1151,11 @@ class LayerwiseOffloadManager:
         if self.copy_stream is not None:
             torch.get_device_module().current_stream().wait_stream(self.copy_stream)
 
+        # A layer still in flight holds device tensors inside the courier;
+        # collecting binds and accounts for them so the release below sees them.
+        for layer_idx in list(self._courier_inflight):
+            self._collect_mapped_layer(layer_idx)
+
         for layer_idx in list(self._gpu_layers):
             self.release_layer(layer_idx, force=True)
 

python/sglang/multimodal_gen/test/server/gpu_cases.py

@@ -1152,6 +1152,73 @@ def _make_5090_flux_layerwise_cpu_offload_case() -> DiffusionTestCase:
     )
 
 
+def _make_5090_h3_consumer_budget_case() -> DiffusionTestCase:
+    """MiniMax-H3 on a pretend 12 GiB card with a pretend 32 GiB host.
+
+    This is the cookbook's consumer recipe, and it guards the whole constrained
+    placement stack at once: the deployment-size gate that keeps the VAE on its
+    checkpoint mapping, per-layer pinning under a small budget, the courier
+    thread that ships still-mapped layers, and the decode-scoped VAE residency
+    that decodes inside the VRAM the denoise just vacated. The runner's real
+    card and host are never short, so the pretend sizes are what route the run
+    onto those paths: the allocator cap makes "fits 12 GiB" enforced rather
+    than inferred, and the runtime peak baseline is that budget.
+    """
+    return DiffusionTestCase(
+        "minimax_h3_t2va_consumer_budget_1gpu_5090",
+        DiffusionServerArgs(
+            model_path="MiniMaxAI/MiniMax-H3",
+            modality="video",
+            extras=[
+                "--model-variant",
+                "fl2va",
+                "--revision",
+                "42ed227ee7df40d41602854ae760620d6eb651fe",
+                "--performance-mode",
+                "memory",
+                "--layerwise-offload-components",
+                "dit,text_encoder,vae",
+                "--layerwise-resident-layers",
+                "video_vae=24",
+            ],
+            env_vars={
+                "SGLANG_DIFFUSION_TEST_FORCE_HOST_AVAILABLE_GIB": "32",
+                "SGLANG_DIFFUSION_TEST_CAP_DEVICE_MEMORY_GIB": "12",
+                # the decode holds two thirds of the decoder against the 12 GiB
+                # cap; without expandable segments, fragmentation tips the
+                # last hundred MiB over
+                "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
+            },
+        ),
+        DiffusionSamplingParams(
+            prompt=("A cat walking on a sunny beach, gentle waves, soft camera pan."),
+            output_size="672x384",
+            seconds=4,
+            output_format="mp4",
+            expect_audio_output=True,
+            num_outputs_per_prompt=1,
+            extras={
+                "task": "t2va",
+                "conditions": [],
+                "target": {
+                    "short_edge": 384,
+                    "aspect_ratio": "16:9",
+                    "duration_seconds": 4.0,
+                },
+                "num_inference_steps": 8,
+                "flow_shift": 12.0,
+                "audio_flow_shift": 3.0,
+                "seed": 1101,
+            },
+        ),
+        run_perf_check=True,
+        run_consistency_check=False,
+        run_component_accuracy_check=False,
+        run_models_api_check=False,
+        run_t2v_input_reference_check=False,
+    )
+
+
 ONE_GPU_5090_CANARY_CASE_IDS = (
     "zimage_image_t2i",
     "flux_2_klein_base_image_t2i",
@@ -1162,6 +1229,7 @@ if not current_platform.is_hip():
 
 ONE_GPU_5090_CASES = _select_5090_canary_cases(ONE_GPU_5090_CANARY_CASE_IDS)
 ONE_GPU_5090_CASES.append(_make_5090_flux_layerwise_cpu_offload_case())
+ONE_GPU_5090_CASES.append(_make_5090_h3_consumer_budget_case())
 
 
 # Nested unit/ tests verified to pass on AMD/ROCm as-is (no code change).

python/sglang/multimodal_gen/test/server/perf_baselines/5090.json

@@ -133,6 +133,29 @@
             "expected_avg_denoise_ms": 2200.8,
             "expected_median_denoise_ms": 2200.8,
             "estimated_full_test_time_s": 90.0
+        },
+        "minimax_h3_t2va_consumer_budget_1gpu_5090": {
+            "stages_ms": {
+                "MiniMaxH3TextEncodingStage": 11900.0,
+                "MiniMaxH3DenoisingStage": 76100.0,
+                "MiniMaxH3DecodingStage": 26550.0
+            },
+            "denoise_step_ms": {
+                "0": 14000.0,
+                "1": 14000.0,
+                "2": 14000.0,
+                "3": 14000.0,
+                "4": 14000.0,
+                "5": 14000.0,
+                "6": 14000.0,
+                "7": 14000.0
+            },
+            "expected_e2e_ms": 125000.0,
+            "expected_avg_denoise_ms": 12000.0,
+            "expected_median_denoise_ms": 11000.0,
+            "load_peak_vram_mb": 6000.0,
+            "runtime_peak_vram_mb": 12288.0,
+            "estimated_full_test_time_s": 900.0
         }
     }
 }

python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py

@@ -265,3 +265,17 @@ class TestPinBenefit:
 
     def test_missing_step_count_is_treated_as_one_use(self):
         assert pin_benefit_bytes(weight_bytes=1000, uses_per_request=0) == 1000
+
+
+def test_the_forced_host_size_behaves_like_a_machine_of_that_size(monkeypatch):
+    monkeypatch.setenv("SGLANG_DIFFUSION_TEST_FORCE_HOST_AVAILABLE_GIB", "32")
+    available = host_memory_budget.host_memory_available_bytes()
+    # available is the pretend total minus this process's own anonymous
+    # memory, so it must sit strictly inside the pretend machine
+    assert 0 < available <= 32 * 1024**3
+    # and a larger pretend machine reports more room, same process
+    monkeypatch.setenv("SGLANG_DIFFUSION_TEST_FORCE_HOST_AVAILABLE_GIB", "64")
+    larger = host_memory_budget.host_memory_available_bytes()
+    assert (
+        abs((larger - available) - 32 * 1024**3) < 512 * 1024**2
+    ), "the same process on a machine twice the size has one machine more room"

python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py

@@ -65,6 +65,9 @@ class _FakeEvent:
     def record(self, _stream) -> None:
         return None
 
+    def synchronize(self) -> None:
+        return None
+
 
 class _FakeDeviceModule:
     Stream = _FakeStream
@@ -1308,6 +1311,84 @@ class _MixedBlock(torch.nn.Module):
         )
 
 
+def test_mapped_layers_ship_through_the_courier(tmp_path, monkeypatch):
+    if not pathlib.Path("/proc/self/maps").exists():
+        pytest.skip("needs /proc to tell a mapping from anonymous memory")
+    manager = _mapped_manager(tmp_path, monkeypatch, available_gib=0.001)
+    model = manager.model
+    # the manager has already swapped the parameter for a placeholder; the
+    # checkpoint file itself holds zeros, so that is the expected content
+    expected = torch.zeros(8, 8)
+
+    manager.prefetch_layer(0, non_blocking=True)
+    assert 0 in manager._courier_inflight, "an async prefetch hands the layer over"
+    assert (
+        0 not in manager._gpu_layers
+    ), "the layer is not ready until its tensors are bound on this thread"
+
+    manager.prefetch_layer(0, non_blocking=False)
+    assert 0 in manager._gpu_layers and not manager._courier_inflight
+    assert torch.equal(
+        model.blocks[0].weight.detach().cpu(), expected
+    ), "the bytes that went through the courier's slot must be the checkpoint's"
+
+
+def test_the_courier_kill_switch_forces_the_synchronous_path(tmp_path, monkeypatch):
+    if not pathlib.Path("/proc/self/maps").exists():
+        pytest.skip("needs /proc to tell a mapping from anonymous memory")
+    # The switch is read when the courier would first be built, and
+    # initialization itself prefetches -- so set it before the manager exists,
+    # the way a deployment sets it before starting the server.
+    monkeypatch.setenv("SGLANG_DIFFUSION_DISABLE_MAPPED_COURIER", "1")
+    manager = _mapped_manager(tmp_path, monkeypatch, available_gib=0.001)
+    manager.release_all()
+    manager.prefetch_layer(0, non_blocking=True)
+    assert not manager._courier_inflight and manager._mapped_courier is None
+    assert (
+        0 in manager._gpu_layers
+    ), "with the courier disabled the direct synchronous path serves the layer"
+
+
+def test_release_all_drains_the_courier(tmp_path, monkeypatch):
+    if not pathlib.Path("/proc/self/maps").exists():
+        pytest.skip("needs /proc to tell a mapping from anonymous memory")
+    manager = _mapped_manager(tmp_path, monkeypatch, available_gib=0.001)
+    manager.prefetch_layer(0, non_blocking=True)
+
+    manager.release_all()
+
+    assert not manager._courier_inflight
+    assert not manager._gpu_layers
+    courier = manager._mapped_courier
+    assert courier is not None and not courier._results, (
+        "an uncollected layer would keep its device tensors alive inside the "
+        "courier for the rest of the process"
+    )
+
+
+def test_a_broken_courier_falls_back_to_the_synchronous_copy(tmp_path, monkeypatch):
+    if not pathlib.Path("/proc/self/maps").exists():
+        pytest.skip("needs /proc to tell a mapping from anonymous memory")
+    manager = _mapped_manager(tmp_path, monkeypatch, available_gib=0.001)
+    model = manager.model
+    expected = torch.zeros(8, 8)
+
+    manager.prefetch_layer(0, non_blocking=True)
+    # the courier dies with the layer still in flight
+    courier = manager._mapped_courier
+    courier._broken = True
+    with courier._ready:
+        courier._results.pop(0, None)
+        courier._pending.discard(0)
+        courier._ready.notify_all()
+
+    manager.prefetch_layer(0, non_blocking=False)
+
+    assert 0 in manager._gpu_layers
+    assert manager._mapped_courier is None, "a failed courier is retired"
+    assert torch.equal(model.blocks[0].weight.detach().cpu(), expected)
+
+
 class _MixedModel(torch.nn.Module):
     def __init__(self, path: pathlib.Path, num_blocks: int) -> None:
         super().__init__()