From 993e24df750d2730d673350d05569efdeb778102 Mon Sep 17 00:00:00 2001 From: triple-mu Date: Thu, 13 Aug 2026 19:49:56 +0800 Subject: [PATCH] [diffusion] Add --dit-layerwise-residency-policy for strided DiT residency (#34534) Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> --- .../memory_managers/layerwise_offload.py | 169 +++++++++++-- .../layerwise_offload_components.py | 7 + .../runtime/server_args/server_args.py | 53 +++- .../test/unit/test_layerwise_offload.py | 230 +++++++++++++++++- 4 files changed, 441 insertions(+), 18 deletions(-) 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 494ab4947..1262dcea0 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,3 +1,4 @@ +import bisect import re from collections.abc import Mapping, Sequence from typing import Any, Dict, List, Set, Tuple @@ -8,6 +9,9 @@ from torch.distributed.tensor import DTensor from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import ( LAYERWISE_OFFLOAD_ALL_COMPONENTS, LAYERWISE_OFFLOAD_DIT_GROUP, + RESIDENCY_POLICIES, + RESIDENCY_POLICY_LEADING, + RESIDENCY_POLICY_STRIDED, layerwise_component_matches_any_selection, normalize_layerwise_offload_components, ) @@ -18,6 +22,56 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) +def compute_streamed_layers( + *, num_layers: int, resident_layers: int, policy: str +) -> tuple[int, ...]: + """Which layer indices are streamed rather than held on the GPU. + + Both policies stream the same *count* of layers, so they cost the same + memory and move the same bytes. They differ only in when those bytes move: + + ``leading`` keeps layers ``0..r-1`` and streams the tail. Every streamed + layer sits next to another streamed layer, so the transfers + arrive as one burst confined to the last ``(n-r)/n`` of the + step, and each has exactly one layer of compute to hide behind. + + ``strided`` spreads the streamed layers evenly across the whole step, so + the same bytes move over ``n`` layers instead of ``n-r`` and + the peak concurrent traffic drops by ``n/(n-r)``. + + What that buys is contention, not bandwidth and not stalls. Profiling the + two policies on an 8-GPU run shows the same HtoD volume to within 0.1%, the + copy engines about half idle in both, and only a handful of long gaps in + either. What differs is how much traffic is in flight beside the compute: + under ``strided`` the GEMM, the attention and the sequence-parallel + all-to-all each run measurably faster (-1.5%, -0.7%, -0.5%) without any + kernel changing, which is the whole of the -0.5% end to end. + + Returned sorted, and always exactly ``num_layers - resident_layers`` long. + """ + if policy not in RESIDENCY_POLICIES: + raise ValueError( + f"unknown residency policy {policy!r}, expected one of {RESIDENCY_POLICIES}" + ) + resident = min(max(0, resident_layers), num_layers) + streamed_count = num_layers - resident + if streamed_count <= 0: + return () + if resident <= 0: + return tuple(range(num_layers)) + + if policy == RESIDENCY_POLICY_LEADING: + return tuple(range(resident, num_layers)) + + # The step num_layers / streamed_count is >= 1 here (resident > 0 was + # handled above), so round() of the ramp is strictly increasing and the + # indices cannot collide -- the partition is total by construction, which + # test_both_policies_partition_the_stack pins. + return tuple( + round(index * num_layers / streamed_count) for index in range(streamed_count) + ) + + # Adapted from skywork AI Infra diffusion optimize class LayerwiseOffloadManager: """A lightweight layerwise CPU offload manager. @@ -43,15 +97,26 @@ class LayerwiseOffloadManager: pin_cpu_memory: bool = True, prefetch_size: int = 1, resident_layers: int = 0, + residency_policy: str = RESIDENCY_POLICY_LEADING, ) -> None: self.model = model self.layers_attr_str = layers_attr_str self.num_layers = num_layers self.pin_cpu_memory = pin_cpu_memory self.prefetch_size = min(max(1, prefetch_size), self.num_layers) - # Leading layers held on GPU across denoise steps, instead of being - # re-streamed every step like the tail. + # Layers held on GPU across denoise steps, instead of being re-streamed + # every step. `residency_policy` picks *which* layers those are; see + # compute_streamed_layers for why the choice is not cosmetic. self.resident_layers = min(max(0, int(resident_layers)), self.num_layers) + self.residency_policy = residency_policy + self._streamed_order = compute_streamed_layers( + num_layers=self.num_layers, + resident_layers=self.resident_layers, + policy=residency_policy, + ) + self._resident_set = frozenset(range(self.num_layers)) - set( + self._streamed_order + ) # Armed on the first denoise forward, so that the load-time prefetch below # does not pin the whole resident set before the DiT is the active component. self._residency_active = False @@ -257,36 +322,93 @@ class LayerwiseOffloadManager: if not self._has_dtensor_weights: self.model.to(self.device) - # prefetch the first layer for warm-up + # prefetch the head of the stream for warm-up; residency is not armed + # yet, so this is layer 0 regardless of policy self.prepare_for_next_req(non_blocking=False) self.register_forward_hooks() self._configured = True logger.info( - f"LayerwiseOffloadManager initialized with num prefetched layer: {self.prefetch_size}, num resident layers: {self.resident_layers}, total num layers: {self.num_layers}" + f"LayerwiseOffloadManager initialized with num prefetched layer: {self.prefetch_size}, num resident layers: {self.resident_layers}, total num layers: {self.num_layers}, residency policy: {self.residency_policy}" ) + if self.residency_policy == RESIDENCY_POLICY_STRIDED and self._streamed_order: + # Printed because the layout is the whole point of the policy, and + # "did it actually stride?" is otherwise only answerable from a + # profile. + logger.info( + "Strided residency streams layers %s (%d of %d)", + list(self._streamed_order), + len(self._streamed_order), + self.num_layers, + ) + + def _head_of_stream(self) -> list[int]: + """The first layers the coming forward will have to stream in. + + Before residency is armed nothing is pinned, so the forward starts at + layer 0 like any other; afterwards the first streamed layer is whichever + the policy put first, which under `strided` is not necessarily layer 0. + """ + count = min(self.prefetch_size, self.num_layers) + if not self._residency_active: + return list(range(count)) + return self._next_streamed(after=-1, count=count) def prepare_for_next_req(self, non_blocking=True): """ Prepare for the next round of denoising loop with prefetching the necessary layers """ - num_prefetch_layers = max(self.prefetch_size, self._retained_layers) - for i in range(num_prefetch_layers): - self.prefetch_layer(i, non_blocking=non_blocking) + # The resident set first: it has to be there for the whole step, and the + # caller decides whether to block on it. + for layer_idx in sorted(self._retained_set): + self.prefetch_layer(layer_idx, non_blocking=non_blocking) if not non_blocking and self.copy_stream is not None: torch.get_device_module().current_stream().wait_stream(self.copy_stream) + # The head of the stream is issued after that wait, and always + # asynchronously. wait_stream drains the whole copy stream, so issuing + # it first would make the caller block on a layer it does not need yet: + # this runs from the layer-0 pre-hook on every denoise step, and under + # `leading` the first streamed layer is `resident_layers` away, one full + # transfer (~48 ms on the 50-layer H3 DiT) ahead of a layer 0 that is + # already pinned. The per-layer wait_event in the pre-hook blocks + # exactly when the weights are needed and no earlier. + for layer_idx in self._head_of_stream(): + self.prefetch_layer(layer_idx, non_blocking=True) + @property def holds_residents(self) -> bool: - """True if this manager keeps a resident leading-layer set beyond the - streaming prefetch window, so it must be denoise-stage-scoped.""" + """True if this manager keeps a resident layer set beyond the streaming + prefetch window, so it must be denoise-stage-scoped.""" return self.enabled and self.resident_layers > 0 @property def _retained_layers(self) -> int: - """Leading layers currently held across denoise steps; 0 until armed.""" + """How many layers are currently held across denoise steps; 0 until armed.""" return self.resident_layers if self._residency_active else 0 + @property + def _retained_set(self) -> frozenset[int]: + """Which layers are currently held across denoise steps; empty until armed.""" + return self._resident_set if self._residency_active else frozenset() + + def _next_streamed(self, *, after: int, count: int) -> List[int]: + """The next ``count`` streamed layers after ``after``, wrapping around. + + Under ``leading`` this is just the following indices, but under + ``strided`` the immediate successor is usually resident, so prefetching + ``after + 1`` would be a no-op and the real next transfer would not + start until its own layer was already running. + """ + total = len(self._streamed_order) + if total == 0: + return [] + start = bisect.bisect_right(self._streamed_order, after) + return [ + self._streamed_order[(start + offset) % total] + for offset in range(min(count, total)) + ] + @torch.compiler.disable def _activate_residency(self) -> None: """Arm the resident set on the first denoise forward. The pinning itself is @@ -372,12 +494,12 @@ class LayerwiseOffloadManager: lightweight release layer weights Basically set the reference count to the gpu weight tensor to zero. The weights on cpu is untouched - Leading resident layers are kept across denoise steps + Resident layers are kept across denoise steps """ if not self.enabled or self.device is None: return - if not force and layer_idx < self._retained_layers: + if not force and layer_idx in self._retained_set: return # clear prefetch event, since it's useless and needs to be reset @@ -572,8 +694,22 @@ class LayerwiseOffloadManager: self._prefetch_events[i] ) + if self.residency_policy == RESIDENCY_POLICY_STRIDED: + # Top up the stream at every layer rather than in bursts of + # prefetch_size. Under `strided` the next streamed layer can + # be several layers away, so a burst schedule keyed on index + # arithmetic would either skip it or issue it late; asking + # for "the next N streamed layers" is the same request every + # layer and prefetch_layer is idempotent, so the repeats are + # free. This is what buys the wider hiding window: the + # transfer is issued as soon as the previous streamed layer + # is done with, not one layer before it is needed. + for layer_to_prefetch in self._next_streamed( + after=i, count=self.prefetch_size + ): + self.prefetch_layer(layer_to_prefetch, non_blocking=True) # trigger batch prefetch (i + prefetch_size ~ i + 2 * prefetch_size) if needed - if i % self.prefetch_size == 0: + elif i % self.prefetch_size == 0: for j in range(i + self.prefetch_size, i + 2 * self.prefetch_size): layer_to_prefetch = j % self.num_layers self.prefetch_layer(layer_to_prefetch, non_blocking=True) @@ -654,6 +790,11 @@ class LayerwiseOffloadableModuleMixin: pin_cpu_memory=server_args.pin_cpu_memory, prefetch_size=prefetch_size, resident_layers=resident_layers, + residency_policy=( + server_args.dit_layerwise_residency_policy + if dit_tuning_enabled + else RESIDENCY_POLICY_LEADING + ), ) self.layerwise_offload_managers.append(manager) configured_layer_names.append(layer_name) @@ -744,7 +885,7 @@ def is_layerwise_offloaded_module(module: torch.nn.Module) -> bool: def is_resident_layerwise_module(module: torch.nn.Module) -> bool: - """True if the module keeps leading DiT layers resident beyond the streaming + """True if the module keeps a resident DiT layer set beyond the streaming prefetch window. """ return isinstance(module, LayerwiseOffloadableModuleMixin) and any( diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload_components.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload_components.py index a2ac0b6fc..4e803c8f5 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload_components.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload_components.py @@ -7,6 +7,13 @@ LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP = "image_encoder" LAYERWISE_OFFLOAD_VAE_GROUP = "vae" LAYERWISE_OFFLOAD_DEFAULT_GROUP = "default" +# Which layers --dit-layerwise-resident-layers keeps on the GPU. Lives here +# rather than in layerwise_offload.py because server_args needs it for the CLI +# choices and layerwise_offload.py imports server_args, which would be a cycle. +RESIDENCY_POLICY_LEADING = "leading" +RESIDENCY_POLICY_STRIDED = "strided" +RESIDENCY_POLICIES = (RESIDENCY_POLICY_LEADING, RESIDENCY_POLICY_STRIDED) + # Components whose layerwise policy has been validated as a better default than # component-level CPU offload when the user has not pinned their placement. LAYERWISE_OFFLOAD_DEFAULT_GROUP_COMPONENTS = ( diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index 5036b3cfc..6c99daead 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -35,6 +35,8 @@ from sglang.multimodal_gen.runtime.loader.utils import BYTES_PER_GB from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import ( LAYERWISE_OFFLOAD_ALL_COMPONENTS, LAYERWISE_OFFLOAD_DIT_GROUP, + RESIDENCY_POLICIES, + RESIDENCY_POLICY_LEADING, cpu_offload_component_matches, cpu_offload_flags_for_layerwise_components, is_dit_component_name, @@ -317,8 +319,10 @@ class ServerArgs(DisaggServerArgsMixin): dit_layerwise_offload: bool | None = None layerwise_offload_components: list[str] | None = None dit_offload_prefetch_size: float = 0.0 - # If set, keep this many leading DiT layers resident on GPU + # If set, keep this many DiT layers resident on GPU dit_layerwise_resident_layers: float = 0.0 + # Which layers those are: the leading ones, or spread evenly over the stack. + dit_layerwise_residency_policy: str = RESIDENCY_POLICY_LEADING offload_during_compile: bool = True text_encoder_cpu_offload: bool | None = None image_encoder_cpu_offload: bool | None = None @@ -1912,14 +1916,29 @@ class ServerArgs(DisaggServerArgsMixin): "--dit-layerwise-resident-layers", type=float, default=ServerArgs.dit_layerwise_resident_layers, - help="With --dit-layerwise-offload, keep this many leading DiT layers " + help="With --dit-layerwise-offload, keep this many DiT layers " "permanently resident on GPU (retained across denoise steps) and stream " - "only the tail with --dit-offload-prefetch-size. 0.0 = off (pure " + "the rest with --dit-offload-prefetch-size; which layers stay resident " + "is --dit-layerwise-residency-policy. 0.0 = off (pure " "streaming). Between 0.0 and 1.0 = ratio of layers; >= 1 = absolute " "count. Unlike raising the prefetch size, resident layers are transferred " "once (not re-streamed every step), so this trades VRAM for lower denoise " "latency when memory is available.", ) + parser.add_argument( + "--dit-layerwise-residency-policy", + type=str, + choices=RESIDENCY_POLICIES, + default=ServerArgs.dit_layerwise_residency_policy, + help="Which layers --dit-layerwise-resident-layers keeps resident. " + "'leading' (default) keeps the first N, which crams the whole " + "weight stream into the tail of each step. 'strided' spreads the " + "resident layers evenly over the stack so the same bytes move over " + "the whole step instead: same VRAM, same bytes, only a different " + "schedule. Worth trying when weight streaming overlaps " + "memory-bound compute -- the transfers stop competing with it for " + "L2 and DRAM bandwidth, which is where the gain comes from.", + ) # offload flags parser.add_argument( @@ -2704,6 +2723,34 @@ class ServerArgs(DisaggServerArgsMixin): "--dit-layerwise-offload (or 'dit' in --layerwise-offload-components)." ) + if self.dit_layerwise_residency_policy not in RESIDENCY_POLICIES: + # argparse's choices= only covers the CLI; ServerArgs is also + # constructed directly by the Python API, and without this the bad + # value would surface as a ValueError inside the GPU worker at + # model-load time. + raise ValueError( + f"Invalid --dit-layerwise-residency-policy " + f"{self.dit_layerwise_residency_policy!r}; expected one of " + f"{RESIDENCY_POLICIES}." + ) + + if self.dit_layerwise_residency_policy != RESIDENCY_POLICY_LEADING: + if not self.is_dit_layerwise_offload_selected: + logger.warning( + "--dit-layerwise-residency-policy has no effect because the DiT is " + "not layerwise-offloaded. It only applies together with " + "--dit-layerwise-offload (or 'dit' in " + "--layerwise-offload-components)." + ) + elif self.dit_layerwise_resident_layers <= 0: + # With nothing resident every layer streams, so there is no + # layout to choose and the policies are the same run. + logger.warning( + "--dit-layerwise-residency-policy has no effect because " + "--dit-layerwise-resident-layers is 0: every layer is streamed, " + "so there is no resident set to place." + ) + # validate layerwise offload conflicts if envs.SGLANG_CACHE_DIT_ENABLED and self.use_fsdp_inference: if self.is_arg_explicitly_set("use_fsdp_inference"): 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 49e0c363d..38ad0f286 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -28,11 +28,16 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_resident_s from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, LayerwiseOffloadManager, + compute_streamed_layers, configure_layerwise_offload_modules, get_layerwise_offload_component_names_for_pipeline, is_layerwise_offloaded_module, is_resident_layerwise_module, ) +from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import ( + RESIDENCY_POLICY_LEADING, + RESIDENCY_POLICY_STRIDED, +) from sglang.multimodal_gen.runtime.server_args import ServerArgs @@ -176,6 +181,7 @@ def _server_args(**kwargs): vae_cpu_offload=False, dit_offload_prefetch_size=1, dit_layerwise_resident_layers=0.0, + dit_layerwise_residency_policy=RESIDENCY_POLICY_LEADING, pin_cpu_memory=False, ) defaults.update(kwargs) @@ -627,7 +633,14 @@ def _patch_fake_device(monkeypatch): monkeypatch.setattr(layerwise_offload_mod.current_platform, "device_type", "cpu") -def _resident_manager(model, *, num_layers, prefetch_size=1, resident_layers=0): +def _resident_manager( + model, + *, + num_layers, + prefetch_size=1, + resident_layers=0, + residency_policy=RESIDENCY_POLICY_LEADING, +): return LayerwiseOffloadManager( model=model, layers_attr_str="blocks", @@ -636,6 +649,7 @@ def _resident_manager(model, *, num_layers, prefetch_size=1, resident_layers=0): pin_cpu_memory=False, prefetch_size=prefetch_size, resident_layers=resident_layers, + residency_policy=residency_policy, ) @@ -698,6 +712,87 @@ def test_prepare_for_next_req_repins_residents(monkeypatch): assert {0, 1, 2} <= manager._gpu_layers +def _record_prepare(manager, monkeypatch): + """Log the order of prefetches and stream waits inside prepare_for_next_req. + + The order is the contract: prepare_for_next_req runs from the layer-0 + pre-hook on every denoise step, so anything issued before the blocking + wait_stream is something the compute stream will sit and wait for. + """ + log: list = [] + original = manager.prefetch_layer + + def spy(layer_idx, non_blocking=True): + log.append(("prefetch", layer_idx, non_blocking)) + return original(layer_idx, non_blocking=non_blocking) + + class _RecordingStream(_FakeStream): + def wait_stream(self, _stream) -> None: + log.append(("wait_stream",)) + + monkeypatch.setattr(manager, "prefetch_layer", spy) + monkeypatch.setattr( + _FakeDeviceModule, "current_stream", staticmethod(_RecordingStream) + ) + return log + + +def test_blocking_prepare_waits_for_residents_only(monkeypatch): + # Regression: the head of the stream used to be issued before the + # wait_stream. wait_stream drains the whole copy stream, so under `leading` + # that made every denoise step block on layer `resident_layers` -- a full + # transfer that is not needed for another `resident_layers` layers, while + # layer 0 was already pinned. Costs one layer transfer per step on the + # default path, which is the path nobody passes a flag to get. + _patch_fake_device(monkeypatch) + manager = _resident_manager( + _MultiBlockModel(6), num_layers=6, prefetch_size=1, resident_layers=3 + ) + _arm_residency(manager) + manager.release_all() + + log = _record_prepare(manager, monkeypatch) + manager.prepare_for_next_req(non_blocking=False) + + wait = log.index(("wait_stream",)) + before = [entry for entry in log[:wait] if entry[0] == "prefetch"] + after = [entry for entry in log[wait:] if entry[0] == "prefetch"] + assert [entry[1] for entry in before] == [0, 1, 2] + assert all(entry[2] is False for entry in before) + # Layer 3 is the first streamed layer, and it is issued after the wait and + # asynchronously, so the pre-hook's own wait_event is what blocks on it. + assert [entry[1] for entry in after] == [3] + assert all(entry[2] is True for entry in after) + + +def test_warmup_prepare_prefetches_the_layer_that_runs_first(monkeypatch): + # At load time residency is not armed yet, so there is no resident set and + # the forward will start at layer 0 whatever the policy says. Deriving the + # head of the stream from the policy here would warm layer + # `resident_layers` under `leading`, i.e. not the one about to run. + _patch_fake_device(monkeypatch) + for policy in (RESIDENCY_POLICY_LEADING, RESIDENCY_POLICY_STRIDED): + manager = _resident_manager( + _MultiBlockModel(6), + num_layers=6, + prefetch_size=1, + resident_layers=3, + residency_policy=policy, + ) + assert manager._head_of_stream() == [0], policy + + +def test_configure_resolves_residency_policy(monkeypatch): + _patch_fake_device(monkeypatch) + comp = _ResidentComponent(8) + comp.configure_layerwise_offload( + _server_args(dit_layerwise_residency_policy=RESIDENCY_POLICY_STRIDED) + ) + assert comp.layerwise_offload_managers[0].residency_policy == ( + RESIDENCY_POLICY_STRIDED + ) + + def test_holds_residents_reflects_configuration(monkeypatch): _patch_fake_device(monkeypatch) resident = _resident_manager(_MultiBlockModel(3), num_layers=3, resident_layers=2) @@ -814,3 +909,136 @@ def test_enable_offload_rearms_after_disable(monkeypatch): assert tuple(model.blocks[2].weight.shape) == (1,) manager.prefetch_layer(2, non_blocking=False) assert torch.equal(model.blocks[2].weight.data, original) + + +# --------------------------------------------------------------------------- +# --dit-layerwise-residency-policy: which layers stay resident. +# +# `leading` keeps 0..r-1, so every transfer lands in one burst at the tail of +# the step. `strided` spreads them, so the same bytes move at 1/(n/(n-r)) of the +# peak rate and each transfer gets that many layers of compute to hide behind. +# That matters when the model also runs a collective per layer: measured on +# 8 GPUs the ulysses SendRecv total went 2589.7 -> 4016.0 ms once the DiT +# streamed, for identical collective volume. +# --------------------------------------------------------------------------- +def test_leading_policy_keeps_todays_prefix_layout(): + # Regression guard: `leading` is the default, and changing it would silently + # re-time every existing deployment. + for num_layers, resident in ((4, 2), (50, 35), (12, 1), (7, 6)): + assert compute_streamed_layers( + num_layers=num_layers, + resident_layers=resident, + policy=RESIDENCY_POLICY_LEADING, + ) == tuple(range(resident, num_layers)) + + +def test_strided_policy_layout_is_pinned_for_the_h3_dit(): + # 50 layers with 35 resident is the measured MiniMax-H3 operating point. + # Pinned exactly so a later "simplification" of the ramp cannot quietly + # change the schedule this policy exists to produce. + assert compute_streamed_layers( + num_layers=50, resident_layers=35, policy=RESIDENCY_POLICY_STRIDED + ) == (0, 3, 7, 10, 13, 17, 20, 23, 27, 30, 33, 37, 40, 43, 47) + + +def test_both_policies_partition_the_stack(): + for num_layers in range(1, 33): + for resident in range(0, num_layers + 1): + for policy in (RESIDENCY_POLICY_LEADING, RESIDENCY_POLICY_STRIDED): + streamed = compute_streamed_layers( + num_layers=num_layers, resident_layers=resident, policy=policy + ) + # Every layer is either streamed or resident, never both and + # never neither -- a gap here would strand a layer with no + # weights at forward time. + assert len(streamed) == len(set(streamed)) == num_layers - resident + assert set(streamed) <= set(range(num_layers)) + + +def test_policies_agree_at_the_degenerate_ends(): + for num_layers in (1, 4, 50): + for resident in (0, num_layers): + assert compute_streamed_layers( + num_layers=num_layers, + resident_layers=resident, + policy=RESIDENCY_POLICY_LEADING, + ) == compute_streamed_layers( + num_layers=num_layers, + resident_layers=resident, + policy=RESIDENCY_POLICY_STRIDED, + ) + + +def test_strided_release_keeps_the_spread_resident_set(monkeypatch): + _patch_fake_device(monkeypatch) + # 8 layers, 4 resident -> streams every other layer. + manager = _resident_manager( + _MultiBlockModel(8), + num_layers=8, + resident_layers=4, + residency_policy=RESIDENCY_POLICY_STRIDED, + ) + streamed = set(manager._streamed_order) + resident = set(range(8)) - streamed + assert streamed == {0, 2, 4, 6} + + _arm_residency(manager) + for layer_idx in range(8): + manager.prefetch_layer(layer_idx, non_blocking=False) + for layer_idx in range(8): + manager.release_layer(layer_idx) + + # Non-force release frees exactly the streamed layers, whatever their index. + assert manager._gpu_layers == resident + manager.release_all() + assert not manager._gpu_layers + + +def test_next_streamed_skips_residents_and_wraps(monkeypatch): + _patch_fake_device(monkeypatch) + manager = _resident_manager( + _MultiBlockModel(8), + num_layers=8, + resident_layers=4, + residency_policy=RESIDENCY_POLICY_STRIDED, + ) + # Streamed = {0, 2, 4, 6}. From layer 1 the next transfer to issue is 2, not + # 1+1 handled as "the following index" -- under this policy the immediate + # successor is usually resident and prefetching it would be a no-op. + assert manager._next_streamed(after=1, count=1) == [2] + assert manager._next_streamed(after=2, count=2) == [4, 6] + # Past the last streamed layer it wraps into the next step. + assert manager._next_streamed(after=6, count=2) == [0, 2] + # -1 is the "before the step starts" probe used when priming. + assert manager._next_streamed(after=-1, count=1) == [0] + + +class _RunnableBlockModel(torch.nn.Module): + """Blocks that actually run, so the registered hooks fire for real.""" + + def __init__(self, n: int) -> None: + super().__init__() + self.blocks = torch.nn.ModuleList([_OrderedLinearLayer(1.0) for _ in range(n)]) + + +def test_strided_forward_leaves_exactly_the_resident_set(monkeypatch): + _patch_fake_device(monkeypatch) + model = _RunnableBlockModel(8) + manager = _resident_manager( + model, + num_layers=8, + resident_layers=4, + residency_policy=RESIDENCY_POLICY_STRIDED, + ) + + hidden = torch.ones(1, 2) + for _ in range(2): # two denoise steps: residents must survive the first + for layer in model.blocks: + hidden = layer(hidden) + + # The pre-hook on layer 0 arms residency and primes; the post-hooks release + # only streamed layers. After two full steps the GPU should hold the + # resident set plus whatever the prefetch window pulled in ahead. + resident = set(range(8)) - set(manager._streamed_order) + assert resident <= manager._gpu_layers + assert len(manager._gpu_layers) <= len(resident) + manager.prefetch_size