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 f851966cc..0bbdde9ea 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 @@ -154,6 +154,16 @@ def cgroup_memory_limit_bytes( return None +def physical_host_memory_available_bytes() -> int: + """What the machine itself can still give: kernel available under any cgroup cap.""" + available = int(psutil.virtual_memory().available) + capped = cgroup_memory_limit_bytes() + if capped is None: + return available + limit, usage = capped + return min(available, max(0, limit - usage)) + + def host_memory_available_bytes() -> int: """Bytes this process can still commit without hitting a wall. @@ -175,13 +185,7 @@ def host_memory_available_bytes() -> int: 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: - return available - limit, usage = capped - return min(available, max(0, limit - usage)) + return physical_host_memory_available_bytes() def shared_pool_available_bytes() -> int: @@ -211,6 +215,19 @@ def host_copies_are_redundant() -> bool: return current_platform.device_shares_host_memory() +def page_cache_cannot_hold(mapped_bytes: int) -> bool: + """Whether the kernel's page cache cannot keep a mapping of this size. + + Read against the machine, not the test-only forced view: a mapping the + cache holds is re-read from memory however small the pretend host is. + """ + if mapped_bytes <= 0: + return False + return ( + mapped_bytes >= physical_host_memory_available_bytes() - HOST_COPY_RESERVE_BYTES + ) + + def host_copies_would_not_fit(weight_bytes: int) -> bool: """Whether copying `weight_bytes` into host memory would run the host out. 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 ee96fe374..d50542a17 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 @@ -41,6 +41,7 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget i host_copies_would_not_fit, host_memory_available_bytes, module_weight_bytes, + page_cache_cannot_hold, pin_benefit_bytes, ) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import ( @@ -567,6 +568,36 @@ class _DirectReader: self._fds.clear() +_LSB = bytes(b & 1 for b in range(256)) + + +def _resident_fraction(data_ptr: int, nbytes: int, *, samples: int = 8) -> float: + """Share of a mapping's pages the page cache holds, sampled in 4 MiB windows. + + -1.0 when it cannot be asked (no libc, mincore failed). + """ + if _libc is None or nbytes <= 0: + return -1.0 + page = mmap.PAGESIZE + window = 4 << 20 + start = data_ptr & ~(page - 1) + total = data_ptr + nbytes - start + # mincore wants a page-aligned address: step in whole pages + step = max(total // samples, window) & ~(page - 1) + resident = checked = 0 + offset = 0 + while offset < total: + length = min(window, total - offset) + pages = (length + page - 1) // page + vec = (ctypes.c_ubyte * pages)() + if _libc.mincore(ctypes.c_void_p(start + offset), ctypes.c_size_t(length), vec): + return -1.0 + resident += bytes(vec).translate(_LSB).count(1) + checked += pages + offset += step + return resident / checked if checked else -1.0 + + def _aligned_span(file_offset: int, nbytes: int) -> tuple[int, int]: """(aligned start, span) covering [file_offset, file_offset + nbytes) at 4 KiB granularity.""" start = file_offset & ~(_DIRECT_ALIGN - 1) @@ -680,11 +711,15 @@ class MappedLayerCourier: await_populated: Optional[Callable[[int], bool]] = None, direct_copy: bool = False, direct_read: bool = False, + direct_read_always: bool = False, ) -> None: self._mapped_cpu_weights = mapped_cpu_weights # Read each layer's bytes from the checkpoint file with O_DIRECT into the # pinned slot instead of copying them out of the page cache. self.direct_read = bool(direct_read) and hasattr(os, "O_DIRECT") + # On a shared pool every read goes to the drive anyway; elsewhere a + # layer the page cache already holds is a memcpy, not a re-read. + self._direct_read_always = bool(direct_read_always) self._reader: Optional[_DirectReader] = ( _DirectReader() if self.direct_read else None ) @@ -703,6 +738,9 @@ class MappedLayerCourier: "h2d_issue_s": 0.0, "direct_read_s": 0.0, "direct_read_bytes": 0, + "cached_layers": 0, + "probes": 0, + "probe_fraction_sum": 0.0, } # Blocks until a populator thread has faulted the layer in; True if one # did, so the courier does not fault the same range a second time. @@ -864,6 +902,16 @@ class MappedLayerCourier: located = ( self._reader.locate(cpu_tensor) if self.direct_read else None ) + if located is not None and not self._direct_read_always: + # a layer the cache holds -- or one we cannot ask about -- + # keeps the memcpy path; only pages known to be cold go + # to the drive + fraction = _resident_fraction(cpu_tensor.data_ptr(), nbytes) + stats["probes"] += 1 + stats["probe_fraction_sum"] += max(fraction, 0.0) + if fraction < 0 or fraction >= 0.5: + located = None + stats["cached_layers"] += 1 if located is not None: path, file_offset, _ = located aligned_start, span = _aligned_span(file_offset, nbytes) @@ -1669,6 +1717,40 @@ class LayerwiseOffloadManager: continue populator.submit(ahead, self._mapped_cpu_weights.get(ahead, {}).values()) + def _log_direct_read_summary(self) -> None: + """One line per pass that went to the drive: how much, and what the probe saw.""" + courier = self._mapped_courier + if courier is None: + return + stats = courier.stats + seen = getattr( + self, + "_direct_read_seen", + {"bytes": 0, "cached": 0, "probes": 0, "fraction_sum": 0.0}, + ) + direct_bytes = stats["direct_read_bytes"] - seen["bytes"] + cached = stats["cached_layers"] - seen["cached"] + probes = stats["probes"] - seen["probes"] + fraction_sum = stats["probe_fraction_sum"] - seen["fraction_sum"] + self._direct_read_seen = { + "bytes": stats["direct_read_bytes"], + "cached": stats["cached_layers"], + "probes": stats["probes"], + "fraction_sum": stats["probe_fraction_sum"], + } + if direct_bytes <= 0: + return + logger.info( + "Layerwise offload: %s read %.1f GiB straight from the drive this pass " + "(%d tensors served from the page cache; mean sampled residency %.2f " + "over %d probes).", + self.layers_attr_str, + direct_bytes / 1024**3, + cached, + fraction_sum / probes if probes else -1.0, + probes, + ) + def _log_debug_timing(self) -> None: """Debug: where this stage's layer traffic spent its time.""" if not envs.SGLANG_DIFFUSION_DEBUG_LAYERWISE_TIMING: @@ -1789,13 +1871,16 @@ class LayerwiseOffloadManager: # 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): + if self._mapped_cpu_weights.get(layer_idx) and ( + non_blocking or self._blocking_load_via_courier() + ): courier = self._ensure_mapped_courier() if courier is not None and courier.submit(layer_idx): self._courier_inflight.add(layer_idx) ship_mapped = True if ( - not envs.SGLANG_DIFFUSION_DISABLE_MAPPED_WILLNEED + non_blocking + and not envs.SGLANG_DIFFUSION_DISABLE_MAPPED_WILLNEED and not courier.direct_read ): # Schedule the disk read for this layer's pages now, in @@ -1884,6 +1969,20 @@ class LayerwiseOffloadManager: if not ship_mapped: self._gpu_layers.add(layer_idx) + elif not non_blocking: + self._collect_mapped_layer(layer_idx) + + def _blocking_load_via_courier(self) -> bool: + """Whether a blocking load should still go through the courier. + + A caller that needs the layer now (arming a resident set, the + materialization of a permanent placement) otherwise faults the + mapping in on this thread; when the courier reads directly, cold + pages arrive at the drive's rate instead (measured 4.7 s vs 12 s for + the same 47 GiB on one NVMe). + """ + courier = self._ensure_mapped_courier() + return courier is not None and courier.direct_read def _ensure_mapped_courier(self) -> Optional[MappedLayerCourier]: """The courier, built on first use; None where it cannot help.""" @@ -1891,6 +1990,9 @@ class LayerwiseOffloadManager: return self._mapped_courier if envs.SGLANG_DIFFUSION_DISABLE_MAPPED_COURIER: return None + if getattr(self, "_courier_retired", False): + # a courier that failed stays retired: its layers keep the synchronous copy + return None if self.copy_stream is None or self._synchronous_mps: return None if not self._mapped_bytes: @@ -1906,11 +2008,25 @@ class LayerwiseOffloadManager: # page) and the process's anonymous memory grew past 100 GiB; # the pinned slots stay even on a shared pool. direct_copy=False, + # A host that cannot cache the mapping re-reads it from the + # drive every pass anyway, through 4 KiB faults at the mercy + # of readahead; O_DIRECT into the slots reads at the drive's + # sequential rate (9.4 vs ~1.1 GiB/s on a GB10 NVMe). direct_read=( - host_copies_are_redundant() + ( + host_copies_are_redundant() + or page_cache_cannot_hold(self._mapped_bytes) + ) and not envs.SGLANG_DIFFUSION_DISABLE_MAPPED_DIRECT_READ - and self._mapped_bytes >= MAPPED_DIRECT_READ_MIN_BYTES + # the size floor guards components re-streamed many times + # per request; one armed once (every layer resident) has + # no such pass to protect + and ( + self._mapped_bytes >= MAPPED_DIRECT_READ_MIN_BYTES + or not self._streamed_order + ) ), + direct_read_always=host_copies_are_redundant(), cold_source=self._mapped_source_is_cold, populate_source=self._mapped_source_may_be_cold, await_populated=self._await_mapped_populated, @@ -1929,6 +2045,7 @@ class LayerwiseOffloadManager: exc, ) self._mapped_courier = None + self._courier_retired = True self._mapped_bytes = self._mapped_bytes # unchanged; direct path return self._mapped_courier @@ -1948,6 +2065,7 @@ class LayerwiseOffloadManager: exc, ) self._mapped_courier = None + self._courier_retired = True self._courier_inflight.discard(layer_idx) self.prefetch_layer(layer_idx, non_blocking=False) return @@ -2007,6 +2125,7 @@ class LayerwiseOffloadManager: def release_all(self) -> None: """Release every layer, including the resident ones: this ends the denoise stage that the resident set is scoped to.""" + self._log_direct_read_summary() self._log_debug_timing() if self._mapped_populator is not None: self._mapped_populator.reset() 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 4888bbf44..727653b37 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 @@ -299,3 +299,30 @@ def test_the_forced_host_size_behaves_like_a_machine_of_that_size(monkeypatch): assert abs((larger - available) - 32 * 1024**3) < 512 * 1024**2, ( "the same process on a machine twice the size has one machine more room" ) + + +def test_the_physical_reading_ignores_the_forced_host_view(monkeypatch): + from types import SimpleNamespace + + monkeypatch.setattr( + host_memory_budget.psutil, + "virtual_memory", + lambda: SimpleNamespace(available=200 * host_memory_budget.GIB_BYTES), + ) + monkeypatch.setattr( + host_memory_budget, "cgroup_memory_limit_bytes", lambda *a, **k: None + ) + monkeypatch.setenv("SGLANG_DIFFUSION_TEST_FORCE_HOST_AVAILABLE_GIB", "32") + # the pretend host sizes our own copies ... + assert ( + host_memory_budget.host_memory_available_bytes() + <= 32 * host_memory_budget.GIB_BYTES + ) + # ... but not what the kernel's page cache can hold + assert ( + host_memory_budget.physical_host_memory_available_bytes() + == 200 * host_memory_budget.GIB_BYTES + ) + assert not host_memory_budget.page_cache_cannot_hold( + 45 * host_memory_budget.GIB_BYTES + ) 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 b3f3006fc..f7fb99366 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -1,4 +1,5 @@ import gc +import os import pathlib from contextlib import nullcontext from types import SimpleNamespace @@ -2268,6 +2269,153 @@ def test_mixed_scm_and_dbcache_step_schedule(monkeypatch, step_kinds): manager.prepare_for_next_req(non_blocking=False) +@pytest.mark.skipif(not hasattr(os, "O_DIRECT"), reason="needs O_DIRECT") +def test_mapped_layers_read_directly_when_the_host_cannot_cache_them( + tmp_path, monkeypatch +): + if not pathlib.Path("/proc/self/maps").exists(): + pytest.skip("needs /proc to tell a mapping from anonymous memory") + monkeypatch.setattr(layerwise_offload_mod, "MAPPED_DIRECT_READ_MIN_BYTES", 1) + monkeypatch.setattr( + layerwise_offload_mod, "host_copies_are_redundant", lambda: False + ) + + # the page cache cannot hold the mapping: it is re-read from the drive every pass + monkeypatch.setattr( + layerwise_offload_mod, "page_cache_cannot_hold", lambda _bytes: True + ) + manager = _mapped_manager(tmp_path, monkeypatch, available_gib=0.001) + assert manager._mapped_cpu_weights[0], "expected the weight to stay mapped" + assert manager._ensure_mapped_courier().direct_read + + # the same mapping on a host that can cache it keeps the page-cache path + monkeypatch.setattr( + layerwise_offload_mod, "page_cache_cannot_hold", lambda _bytes: False + ) + manager._mapped_courier = None + assert not manager._ensure_mapped_courier().direct_read + + +@pytest.mark.skipif(layerwise_offload_mod._libc is None, reason="needs libc mincore") +def test_resident_fraction_sees_the_pages_the_cache_holds(tmp_path): + path = tmp_path / "cached.bin" + path.write_bytes(b"\x01" * (16 << 20)) + mapped = torch.from_file(str(path), shared=True, size=16 << 20, dtype=torch.uint8) + mapped.sum() # touch every page + fraction = layerwise_offload_mod._resident_fraction( + mapped.data_ptr(), mapped.numel() + ) + assert fraction >= 0.9 + + +@pytest.mark.skipif(not hasattr(os, "O_DIRECT"), reason="needs O_DIRECT") +def test_cached_mapped_layers_are_copied_rather_than_re_read(tmp_path, monkeypatch): + if not pathlib.Path("/proc/self/maps").exists(): + pytest.skip("needs /proc to tell a mapping from anonymous memory") + monkeypatch.setattr(layerwise_offload_mod, "MAPPED_DIRECT_READ_MIN_BYTES", 1) + monkeypatch.setattr( + layerwise_offload_mod, "host_copies_are_redundant", lambda: False + ) + monkeypatch.setattr( + layerwise_offload_mod, "host_copies_would_not_fit", lambda _bytes: True + ) + monkeypatch.setattr( + layerwise_offload_mod, "page_cache_cannot_hold", lambda _bytes: True + ) + # the page cache holds the layer: shipping it is a memcpy, not a drive read + monkeypatch.setattr( + layerwise_offload_mod, "_resident_fraction", lambda *_a, **_k: 1.0 + ) + manager = _mapped_manager(tmp_path, monkeypatch, available_gib=0.001) + courier = manager._ensure_mapped_courier() + assert courier.direct_read + manager.prefetch_layer(0, non_blocking=False) + assert 0 in manager._gpu_layers + assert courier.stats["cached_layers"] >= 1 + assert courier.stats["direct_read_bytes"] == 0 + + +@pytest.mark.skipif(not hasattr(os, "O_DIRECT"), reason="needs O_DIRECT") +def test_blocking_loads_of_cold_mapped_layers_go_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") + monkeypatch.setattr(layerwise_offload_mod, "MAPPED_DIRECT_READ_MIN_BYTES", 1) + monkeypatch.setattr( + layerwise_offload_mod, "host_copies_are_redundant", lambda: False + ) + monkeypatch.setattr( + layerwise_offload_mod, "host_copies_would_not_fit", lambda _bytes: True + ) + monkeypatch.setattr( + layerwise_offload_mod, "page_cache_cannot_hold", lambda _bytes: True + ) + monkeypatch.setattr( + layerwise_offload_mod, "_resident_fraction", lambda *_a, **_k: 0.0 + ) + manager = _mapped_manager(tmp_path, monkeypatch, available_gib=0.001) + courier = manager._ensure_mapped_courier() + assert courier.direct_read + # a blocking load (how a resident set is armed) is shipped by the courier too + manager.prefetch_layer(0, non_blocking=False) + assert 0 in manager._gpu_layers and not manager._courier_inflight + assert courier.stats["layers"] == 1 + assert torch.equal(manager.model.blocks[0].weight.detach().cpu(), torch.zeros(8, 8)) + + +@pytest.mark.skipif(not hasattr(os, "O_DIRECT"), reason="needs O_DIRECT") +def test_a_fully_resident_small_component_may_still_read_directly( + tmp_path, monkeypatch +): + if not pathlib.Path("/proc/self/maps").exists(): + pytest.skip("needs /proc to tell a mapping from anonymous memory") + monkeypatch.setattr( + layerwise_offload_mod, "host_copies_are_redundant", lambda: False + ) + monkeypatch.setattr( + layerwise_offload_mod, "host_copies_would_not_fit", lambda _bytes: True + ) + monkeypatch.setattr( + layerwise_offload_mod, "page_cache_cannot_hold", lambda _bytes: True + ) + # far below the size floor, but every layer is resident: it is armed once + # per request, so there is no re-streamed pass for the floor to protect + monkeypatch.setattr( + layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule + ) + monkeypatch.setattr(layerwise_offload_mod.current_platform, "device_type", "cpu") + monkeypatch.setattr( + host_memory_budget, "host_memory_available_bytes", lambda: 1 << 20 + ) + model = _FileBackedModel(tmp_path / "weights.bin", num_blocks=2) + manager = LayerwiseOffloadManager( + model=model, + layers_attr_str="blocks", + num_layers=2, + enabled=True, + pin_cpu_memory=True, + prefetch_size=1, + resident_layers=2, + ) + assert manager._mapped_cpu_weights[0] and not manager._streamed_order + assert manager._ensure_mapped_courier().direct_read + + +@pytest.mark.skipif(layerwise_offload_mod._libc is None, reason="needs libc mincore") +def test_resident_fraction_samples_a_large_mapping_at_page_aligned_offsets(tmp_path): + # large enough that the sampling stride exceeds one window: every window + # must start on a page boundary or mincore rejects it and the answer is -1 + path = tmp_path / "large.bin" + path.write_bytes(b"\x01" * (96 << 20)) + mapped = torch.from_file(str(path), shared=True, size=96 << 20, dtype=torch.uint8) + mapped.sum() + fraction = layerwise_offload_mod._resident_fraction( + mapped.data_ptr() + 1000, (96 << 20) - 1000 + ) + assert fraction >= 0.9 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_large_pinned_stores_are_registered_in_place_at_exact_size(monkeypatch): pooled = []