[diffusion] Add --dit-layerwise-residency-policy for strided DiT residency (#34534)

Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>

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
+        # Layers held on GPU across denoise steps, instead of being re-streamed
-        # re-streamed every step like the tail.
+        # 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)
+        # The resident set first: it has to be there for the whole step, and the
-        for i in range(num_prefetch_layers):
+        # caller decides whether to block on it.
-            self.prefetch_layer(i, non_blocking=non_blocking)
+        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
+        """True if this manager keeps a resident layer set beyond the streaming
-        streaming prefetch window, so it must be denoise-stage-scoped."""
+        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(

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

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"):

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