From 382343f860c7df8dc43ecefd7902c940d81c99a9 Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 22 Aug 2026 17:17:14 +0800 Subject: [PATCH] [diffusion] optimization: transfer mapped layers through a courier thread (#35882) Co-authored-by: Claude Opus 5 --- python/sglang/multimodal_gen/envs.py | 32 +++ .../runtime/managers/gpu_worker.py | 21 ++ .../memory_managers/host_memory_budget.py | 17 ++ .../memory_managers/layerwise_offload.py | 236 +++++++++++++++++- .../multimodal_gen/test/server/gpu_cases.py | 68 +++++ .../test/server/perf_baselines/5090.json | 23 ++ .../test/unit/test_host_memory_budget.py | 14 ++ .../test/unit/test_layerwise_offload.py | 81 ++++++ 8 files changed, 490 insertions(+), 2 deletions(-) diff --git a/python/sglang/multimodal_gen/envs.py b/python/sglang/multimodal_gen/envs.py index 4da5b3f84..8e15b209f 100644 --- a/python/sglang/multimodal_gen/envs.py +++ b/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"), diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index e1c6d35a4..29b753352 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py index 8b057cc29..82d5f4fa5 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py index 161a56e30..eb416647d 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py +++ b/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) diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index 30d9726a9..042ee18fc 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/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). diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json b/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json index 47f249009..04ab0907e 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json +++ b/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 } } } diff --git a/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py b/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py index 3b8d99737..646c738a1 100644 --- a/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py +++ b/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" diff --git a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py index 2fb394904..51dee69d0 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/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__()