vllm-omni/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Distributed Layerwise Offload backend with double-buffered H2D.

This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that:

* Optionally shards model weights across DP ranks and reconstructs each block
  with AllGather, or streams a complete rank-local block without a collective.
* Can retain compatible checkpoint tensors as node-shared mmap sources instead
  of creating a persistent private host copy in every rank.
* Uses a fixed double-buffer scheme that keeps only two layers' worth of
  weights on each device at any time.
* Pipelines H2D transfers and, when enabled, AllGather communications on
  dedicated streams, overlapping them with computation.
* Is hardware-agnostic, supporting both NVIDIA GPU (CUDA) and Ascend NPU
  (CANN) platforms via vLLM-Omni's platform abstraction layer.
"""

from __future__ import annotations

from itertools import chain
from typing import Any

import torch
import torch.distributed
from torch import nn
from vllm.logger import init_logger

from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
from vllm_omni.diffusion.model_loader.host_weight_plan import (
    HostWeightPlan,
)
from vllm_omni.platforms import current_omni_platform

from .base import OffloadBackend, OffloadConfig
from .block_discovery import get_blocks_from_dit
from .module_collector import ModuleDiscovery
from .offload_plan import (
    OffloadPlan,
    get_offload_plan,
)
from .tensor_utils import (
    dtype_size as _dtype_size,
)
from .tensor_utils import (
    is_materialized_tensor,
    make_offload_placeholder,
    set_tensor_storage,
)

logger = init_logger(__name__)

# Threshold (in MB) for deciding whether a non-block DiT submodule should
# use layerwise offload (streaming hooks) or be moved to GPU as a resident
# module.  Submodules larger than this are offloaded to save HBM; smaller
# ones stay resident for lower latency.
_ON_DEMAND_THRESHOLD_MB = 1024


class DistributedLayerwiseOffloadHook(ModelHook):
    """Hook for distributed layerwise offloading with fixed double-buffer.

    Each rank stores only a shard of each block's weights on host memory.
    Two device slots alternate: one holds current weights, one holds next weights.
    H2D and AllGather run asynchronously on dedicated streams, overlapped
    with computation.

    Supports both NVIDIA GPU (CUDA) and Ascend NPU (CANN) platforms.
    """

    _HOOK_NAME = "distributed_layerwise_offload"

    def __init__(
        self,
        next_block: nn.Module,
        device: torch.device,
        dp_group: torch.distributed.ProcessGroup | None,
        dp_size: int,
        rank: int,
        copy_stream: Any | None = None,
        comm_stream: Any | None = None,
        pin_memory: bool = True,
        shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] | None = None,
        rank_local_mmap: bool = False,
        tensor_transforms: dict[int, Any] | None = None,
    ):
        assert isinstance(next_block, nn.Module), "transformer block must be type `torch.nn.Module`"

        self.next_block = next_block
        self.device = device
        self.dp_group = dp_group
        self.dp_size = dp_size
        self.rank = rank
        self.pin_memory = pin_memory
        self.rank_local_mmap = rank_local_mmap
        self.tensor_transforms = tensor_transforms or {}

        self.copy_stream = copy_stream or current_omni_platform.Stream()
        self.comm_stream = comm_stream or current_omni_platform.Stream()

        # Double buffers: either shared (from backend) or self-allocated (lazy)
        if shared_buffers is not None:
            self.gpu_buffers: list[dict[torch.dtype, torch.Tensor] | None] = shared_buffers
            self._owns_buffers = False
        else:
            self.gpu_buffers = [None, None]
            self._owns_buffers = True
        self.ready_events: list[Any | None] = [None, None]

        # Sharded host weights for the next block, keyed by dtype
        self.cpu_shards: dict[torch.dtype, torch.Tensor] = {}
        # File-backed source tensors for rank-local mmap.  Unlike cpu_shards,
        # these remain immutable views of the checkpoint and are never pinned
        # or flattened into a model-sized private allocation.
        self.cpu_sources: dict[torch.dtype, list[dict[str, Any]]] = {}
        self.metadata: dict[torch.dtype, list[dict[str, Any]]] = {}

        # Rank-local mmap uses two host staging slots shared by every hook in
        # this worker.  They are assigned by the backend after all block sizes
        # are known, mirroring the shared device-buffer allocation.
        self.cpu_staging_buffers: list[dict[torch.dtype, torch.Tensor] | None] = [None, None]
        self.cpu_staging_events: list[Any | None] = [None, None]

        # Current slot index (0 or 1).  Updated dynamically by the previous
        # hook's prefetch_layer call via _prefetched_slot.  This ensures
        # correct slot tracking for ALL block counts (including odd N).
        self.current_slot = 0
        self._prefetched_slot: int | None = None

        # Backward link to previous hook for fallback (cache-dit skip)
        self._prev_hook: DistributedLayerwiseOffloadHook | None = None

        # Marks the first hook in a shared-buffer group.  When multiple DiT
        # groups share the same 2 GPU buffers, another group may have
        # overwritten this group's slot between forwards.  The first block
        # must sync-prefetch on entry to ensure it loads the correct
        # weights, even if is_materialized sees a non-empty tensor left by
        # the other group.
        self._is_group_first: bool = False

        # Group ID for per-slot contamination tracking.  Shared across all
        # hooks in the same group.  When a hook prefetches into a slot, it
        # stamps the slot with its group ID.  On group entry, the first
        # hook checks whether the slot was last written by its own group
        # (tail hook's async prefetch) — if so, skip the sync-prefetch.
        self._group_id: int = -1
        self._shared_slot_group: list[int] | None = None  # [-1, -1], shared

        # Parameters/buffers of the current and next blocks
        self.block_parameters: dict[str, nn.Parameter] = {}
        self.block_buffers: dict[str, torch.Tensor] = {}
        self.next_block_parameters: dict[str, nn.Parameter] = {}
        self.next_block_buffers: dict[str, torch.Tensor] = {}

        # Per-block synchronization primitive: set after H2D copy completes.
        self._prefetch_done: Any | None = None

        # Shared shard (AllGather input) buffers — assigned by backend.
        self.gpu_shard_buffers: list[dict[torch.dtype, torch.Tensor] | None] = [None, None]

        # Pending async AllGather work (prevent GC before completion).
        self._pending_work: Any | None = None
        self._cached_repoint: list | None = None

    # ------------------------------------------------------------------ #
    #  DTensor helpers (shared with LayerwiseOffloadHook)                 #
    # ------------------------------------------------------------------ #

    def initialize_hook(self, module: nn.Module) -> nn.Module:
        module = super().initialize_hook(module)

        self.block_parameters = dict(module.named_parameters())
        self.block_buffers = dict(module.named_buffers())

        self.next_block_parameters = dict(self.next_block.named_parameters())
        self.next_block_buffers = dict(self.next_block.named_buffers())

        if self.rank_local_mmap:
            self.cpu_sources, self.metadata = self._collect_mmap_sources(
                self.next_block_parameters,
                self.next_block_buffers,
                self.tensor_transforms,
            )
        else:
            # Shard next block's weights and store local shard in pinned CPU memory.
            self.cpu_shards, self.metadata = self._shard_and_pin(
                self.next_block_parameters,
                self.next_block_buffers,
                self.dp_size,
                self.rank,
                self.pin_memory,
                self.tensor_transforms,
            )

        # Allocate device buffers only if not using shared buffers from backend
        if self._owns_buffers:
            self._allocate_device_buffers()

        # Cache parameter re-pointing metadata to avoid per-layer dict lookups.
        self._cached_repoint = []
        for slot in range(2):
            repoint = []
            for dtype, metas in self.metadata.items():
                for m in metas:
                    target = (
                        self.next_block_parameters[m["name"]]
                        if m["name"] in self.next_block_parameters
                        else self.next_block_buffers[m["name"]]
                    )
                    repoint.append(
                        (
                            target,
                            dtype,
                            m["offset"],
                            m["numel"],
                            m["shape"],
                            m["stride"],
                        )
                    )
            self._cached_repoint.append(repoint)

        # Pre-compute AG output sizes (avoid sum() per layer).
        self._ag_output_sizes: dict[torch.dtype, int] = {}
        for dtype, metas in self.metadata.items():
            total_numel = sum(m["numel"] for m in metas)
            if self.rank_local_mmap:
                self._ag_output_sizes[dtype] = total_numel
            else:
                shard_numel = self.cpu_shards[dtype].numel()
                self._ag_output_sizes[dtype] = shard_numel * self.dp_size if self.dp_size > 1 else total_numel

        return module

    @staticmethod
    def _collect_mmap_sources(
        params: dict[str, nn.Parameter],
        bufs: dict[str, torch.Tensor],
        tensor_transforms: dict[int, Any] | None = None,
    ) -> tuple[dict[torch.dtype, list[dict[str, Any]]], dict[torch.dtype, list[dict[str, Any]]]]:
        """Retain file-backed tensors and replace module storage with placeholders.

        The returned sources preserve safetensors mmap storage.  Runtime-layout
        adapters are applied only while packing a bounded staging slot, so no
        full-model anonymous CPU copy is retained by a worker.
        """
        cpu_sources: dict[torch.dtype, list[dict[str, Any]]] = {}
        metadata: dict[torch.dtype, list[dict[str, Any]]] = {}
        offsets: dict[torch.dtype, int] = {}

        for name, target in chain(params.items(), bufs.items()):
            source = target.to_local() if hasattr(target, "to_local") else target
            if source.device.type != "cpu":
                raise ValueError(
                    f"Rank-local mmap storage requires CPU checkpoint views, but {name!r} is on {source.device}."
                )

            dtype = source.dtype
            offset = offsets.get(dtype, 0)
            transform = (tensor_transforms or {}).get(id(target))
            runtime_source = transform(source) if callable(transform) else source
            if runtime_source.dtype != dtype or runtime_source.shape != source.shape:
                raise ValueError(
                    "mmap weight transform changed tensor metadata for "
                    f"{name!r}: expected dtype={dtype}, shape={tuple(source.shape)}, "
                    f"got dtype={runtime_source.dtype}, shape={tuple(runtime_source.shape)}"
                )
            stride = runtime_source.stride()
            storage_numel = (
                0
                if runtime_source.numel() == 0
                else 1 + sum((size - 1) * axis_stride for size, axis_stride in zip(runtime_source.shape, stride))
            )

            cpu_sources.setdefault(dtype, []).append(
                {
                    "name": name,
                    "tensor": source.detach(),
                    "transform": transform,
                }
            )
            metadata.setdefault(dtype, []).append(
                {
                    "name": name,
                    "offset": offset,
                    "numel": storage_numel,
                    "shape": runtime_source.shape,
                    "stride": stride,
                }
            )
            offsets[dtype] = offset + storage_numel

            # The detached source above keeps the mmap storage alive while the
            # module parameter/buffer is rebound to the rotating device slot.
            set_tensor_storage(target, make_offload_placeholder(target))

        return cpu_sources, metadata

    @staticmethod
    def _shard_and_pin(
        params: dict[str, nn.Parameter],
        bufs: dict[str, torch.Tensor],
        dp_size: int,
        rank: int,
        pin_memory: bool,
        tensor_transforms: dict[int, Any] | None = None,
    ) -> tuple[dict[torch.dtype, torch.Tensor], dict[torch.dtype, list[dict[str, Any]]]]:
        """Flatten params+buffers by dtype, split into DP shards, store local shard.

        Each rank stores only ``1/dp_size`` of the total weights. The full
        tensor is reconstructed at runtime via AllGather.
        """
        dtype_grouped: dict[torch.dtype, dict[str, torch.Tensor]] = {}
        dtype_metadata: dict[torch.dtype, list[dict[str, Any]]] = {}

        for name, param_or_buf in chain(params.items(), bufs.items()):
            dtype = param_or_buf.dtype
            if dtype not in dtype_grouped:
                dtype_grouped[dtype] = {}
            dtype_grouped[dtype][name] = param_or_buf

        cpu_shards: dict[torch.dtype, torch.Tensor] = {}

        for dtype, name2weights in dtype_grouped.items():
            # Resolve local tensors (handle DTensor via to_local)
            weights_with_local = []
            for name, t in name2weights.items():
                local_t = t.to_local() if hasattr(t, "to_local") else t
                mmap_transform = (tensor_transforms or {}).get(id(t))
                if callable(mmap_transform):
                    # Some checkpoints use a layout that is converted by the
                    # regular weight loader (for example MiniMax-H3 grouped
                    # QKV).  Apply that conversion one block at a time while
                    # copying the rank-local CPU shard.  Keeping the raw
                    # parameter as an mmap view avoids a private full-model
                    # copy in every worker.
                    local_t = mmap_transform(local_t)
                stride = local_t.stride()
                storage_numel = (
                    0
                    if local_t.numel() == 0
                    else 1 + sum((size - 1) * axis_stride for size, axis_stride in zip(local_t.shape, stride))
                )
                weights_with_local.append((name, t, local_t, storage_numel, stride))

            total_numel = sum(storage_numel for _, _, _, storage_numel, _ in weights_with_local)

            # Equal-sized shards (ceil division) for all_gather_into_tensor
            shard_size = (total_numel + dp_size - 1) // dp_size  # ceil
            shard_start = rank * shard_size
            shard_end = min(shard_start + shard_size, total_numel)

            # Allocate ONLY the shard (1/dp_size), zero-padded to ceil.
            # Avoids materialising the full block on CPU.
            shard = torch.zeros(shard_size, dtype=dtype, device="cpu")

            current_offset = 0
            for (
                name,
                original_tensor,
                local_tensor,
                storage_numel,
                stride,
            ) in weights_with_local:
                if dtype not in dtype_metadata:
                    dtype_metadata[dtype] = []
                # Offsets remain relative to the FULL flattened buffer
                # (needed for correct AllGather reconstruction).
                dtype_metadata[dtype].append(
                    {
                        "name": name,
                        "offset": current_offset,
                        "numel": storage_numel,
                        "shape": local_tensor.shape,
                        "stride": stride,
                    }
                )

                # Copy ONLY the portion within [shard_start, shard_end)
                overlap_start = max(current_offset, shard_start)
                overlap_end = min(current_offset + storage_numel, shard_end)
                if overlap_start < overlap_end:
                    if local_tensor.is_contiguous():
                        flat_storage = local_tensor.flatten()
                    else:
                        # Online FP8 stores Cutlass weights as transposed views
                        # (e.g. stride=(1, K)). Flattening such a tensor in
                        # logical order and later rebuilding it with .view()
                        # changes its layout and makes scaled_mm reject it.
                        # Pack the physical storage order and preserve the
                        # original stride for zero-copy reconstruction.
                        flat_storage = torch.zeros(
                            storage_numel,
                            dtype=dtype,
                            device=local_tensor.device,
                        )
                        physical_view = torch.as_strided(
                            flat_storage,
                            size=local_tensor.shape,
                            stride=stride,
                        )
                        physical_view.copy_(local_tensor)
                    src_start = overlap_start - current_offset
                    src_end = overlap_end - current_offset
                    dst_start = overlap_start - shard_start
                    dst_end = overlap_end - shard_start
                    shard[dst_start:dst_end].copy_(flat_storage[src_start:src_end])

                # Replace original tensor with placeholder (frees CPU storage)
                set_tensor_storage(
                    original_tensor,
                    make_offload_placeholder(original_tensor),
                )
                current_offset += storage_numel

            if pin_memory:
                shard = shard.pin_memory()

            cpu_shards[dtype] = shard

        return cpu_shards, dtype_metadata

    def _allocate_device_buffers(self) -> None:
        """Pre-allocate exactly two device buffers (one per slot)."""
        for slot in range(2):
            gpu_weights: dict[torch.dtype, torch.Tensor] = {}
            for dtype, metas in self.metadata.items():
                total_numel = sum(m["numel"] for m in metas)
                # AllGather output = dp_size * shard_size (padded)
                padded = total_numel
                if self.dp_size > 1:
                    shard_sz = (total_numel + self.dp_size - 1) // self.dp_size
                    padded = shard_sz * self.dp_size
                gpu_weights[dtype] = torch.empty(padded, dtype=dtype, device=self.device)
            self.gpu_buffers[slot] = gpu_weights

    @property
    def is_materialized(self) -> bool:
        """Check whether this block's parameters hold real data on device."""
        for param in self.block_parameters.values():
            return is_materialized_tensor(param)
        return True

    # ------------------------------------------------------------------ #
    #  Prefetch: H2D + AllGather (overlapped on dedicated streams)      #
    # ------------------------------------------------------------------ #

    @staticmethod
    def _pack_mmap_sources(
        cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
        metadata: dict[torch.dtype, list[dict[str, Any]]],
        slot_buffers: dict[torch.dtype, torch.Tensor],
    ) -> dict[torch.dtype, torch.Tensor]:
        staged: dict[torch.dtype, torch.Tensor] = {}
        for dtype, metas in metadata.items():
            total_numel = sum(meta["numel"] for meta in metas)
            destination = slot_buffers[dtype][:total_numel]
            sources = cpu_sources[dtype]
            for source_info, meta in zip(sources, metas, strict=True):
                source = source_info["tensor"]
                transform = source_info["transform"]
                if callable(transform):
                    source = transform(source)
                if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]:
                    raise ValueError(
                        "mmap weight transform changed tensor layout for "
                        f"{source_info['name']!r}: expected dtype={dtype}, "
                        f"shape={tuple(meta['shape'])}, stride={meta['stride']}; "
                        f"got dtype={source.dtype}, shape={tuple(source.shape)}, "
                        f"stride={source.stride()}"
                    )
                start = meta["offset"]
                physical_storage = destination[start : start + meta["numel"]]
                if source.is_contiguous():
                    physical_storage.copy_(source.flatten())
                else:
                    torch.as_strided(
                        physical_storage,
                        size=source.shape,
                        stride=source.stride(),
                    ).copy_(source)
            staged[dtype] = destination
        return staged

    def _stage_mmap_sources(self, slot: int) -> dict[torch.dtype, torch.Tensor]:
        """Pack this block's mmap views into one bounded host staging slot."""
        previous_copy = self.cpu_staging_events[slot]
        if previous_copy is not None:
            synchronize = getattr(previous_copy, "synchronize", None)
            if callable(synchronize):
                synchronize()
            else:
                # Platform events normally expose synchronize().  Retain a
                # correctness fallback for test and non-CUDA platform shims.
                current_omni_platform.synchronize()
            self.cpu_staging_events[slot] = None

        slot_buffers = self.cpu_staging_buffers[slot]
        if slot_buffers is None:
            raise RuntimeError(f"cpu_staging_buffers[{slot}] was not allocated")
        return self._pack_mmap_sources(self.cpu_sources, self.metadata, slot_buffers)

    @torch.compiler.disable
    def prefetch_layer(self, slot: int, non_blocking: bool = True) -> None:
        """Prepare next block's weights into the shared device buffer for *slot*.

        Uses the pre-allocated ``self.gpu_buffers[slot]`` instead of
        allocating fresh tensors every layer.  This enforces the fixed
        double-buffer memory bound (exactly 2 blocks on device).
        """
        self.copy_stream.wait_stream(current_omni_platform.current_stream())

        evt = current_omni_platform.Event()
        gpu_weights = self.gpu_buffers[slot]
        assert gpu_weights is not None, f"gpu_buffers[{slot}] not allocated"

        if self.dp_size <= 1 or self.dp_group is None:
            cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards
            with current_omni_platform.stream(self.copy_stream):
                for dtype, cpu_shard in cpu_weights.items():
                    gw = gpu_weights[dtype]
                    async_copy = non_blocking and cpu_shard.is_pinned()
                    gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy)
                evt.record(self.copy_stream)
            if self.rank_local_mmap:
                # The CPU slot may be overwritten only after this H2D copy has
                # finished.  The shared event protects reuse by another hook.
                self.cpu_staging_events[slot] = evt
        else:
            gpu_shards: dict[torch.dtype, torch.Tensor] = {}
            shard_bufs = self.gpu_shard_buffers[slot]
            assert shard_bufs is not None, f"gpu_shard_buffers[{slot}] not allocated"
            with current_omni_platform.stream(self.copy_stream):
                for dtype, cpu_shard in self.cpu_shards.items():
                    gpu_shard = shard_bufs[dtype][: cpu_shard.numel()]
                    gpu_shard.copy_(cpu_shard, non_blocking=non_blocking)
                    gpu_shards[dtype] = gpu_shard

            self.comm_stream.wait_stream(self.copy_stream)
            with current_omni_platform.stream(self.comm_stream):
                for dtype, local_shard in gpu_shards.items():
                    # Slice the shared (max-sized) output buffer down to this
                    # block's actual AllGather output size. The buffers are
                    # sized to the *largest* block across all groups, so for any
                    # smaller block the full buffer would violate the
                    # all_gather_into_tensor contract
                    # (output.numel() == world_size * input.numel()).
                    # Repoint offsets are relative to the block's flattened
                    # buffer, so a prefix slice is safe.
                    gw = gpu_weights[dtype][: self._ag_output_sizes[dtype]]
                    torch.distributed.all_gather_into_tensor(
                        gw,
                        local_shard,
                        group=self.dp_group,
                    )
                evt.record(self.comm_stream)

        self.ready_events[slot] = evt
        self._prefetch_done = evt
        self._prefetched_slot = slot

        # Stamp the slot with this hook's group ID so that group-first
        # hooks can detect whether another group has overwritten the slot.
        if self._shared_slot_group is not None:
            self._shared_slot_group[slot] = self._group_id

        # Re-point using cached metadata (avoids per-layer dict lookups).
        for target, dtype, offset, numel, shape, stride in self._cached_repoint[slot]:
            set_tensor_storage(
                target,
                torch.as_strided(
                    gpu_weights[dtype][offset : offset + numel],
                    size=shape,
                    stride=stride,
                ),
            )

    def get_weights(self, slot: int) -> dict[torch.dtype, torch.Tensor] | None:
        """Wait for AllGather completion and return full weights for the slot.

        The ready event for this slot was set by the *previous* hook's
        prefetch_layer (which prefetched THIS block's weights into the
        shared buffer).  This hook's own ready_events[slot] is None because
        it never prefetched into this slot itself.  Fall back to the
        previous hook's event to ensure the compute stream waits for the
        AllGather (or H2D) to complete before reading the buffer.
        """
        evt = self.ready_events[slot]
        if evt is None and self._prev_hook is not None:
            evt = self._prev_hook.ready_events[slot]
        if evt is not None:
            current_omni_platform.current_stream().wait_event(evt)
        return self.gpu_buffers[slot]

    # ------------------------------------------------------------------ #
    #  Offload: free device memory for current block                     #
    # ------------------------------------------------------------------ #

    @torch.compiler.disable
    def offload_layer(self) -> None:
        """Free GPU memory for current block by replacing tensors with placeholders."""
        evt = self._prefetch_done
        if evt is not None:
            current_omni_platform.current_stream().wait_event(evt)
        self._prefetch_done = None

        for _, param in self.block_parameters.items():
            set_tensor_storage(param, make_offload_placeholder(param))
        for _, buf in self.block_buffers.items():
            set_tensor_storage(buf, make_offload_placeholder(buf))

    # ------------------------------------------------------------------ #
    #  ModelHook interface                                                #
    # ------------------------------------------------------------------ #

    def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]:
        # Dynamic slot tracking: read current_slot from the previous hook's
        # _prefetched_slot (the slot where this block's data was actually
        # loaded).  This corrects the initial i%2 assignment for odd block
        # counts, where the circular tail→head prefetch would otherwise
        # collide with the head's read slot.
        if self._prev_hook is not None and self._prev_hook._prefetched_slot is not None:
            self.current_slot = self._prev_hook._prefetched_slot

        # Group-first hook: check whether the shared buffer slot was
        # overwritten by another group since our last forward.  If the
        # slot still contains our own data (from the tail hook's async
        # prefetch), skip the sync-prefetch and just wait for the event.
        if self._is_group_first and self._prev_hook is not None:
            slot_contaminated = True
            if self._shared_slot_group is not None:
                slot_contaminated = self._shared_slot_group[self.current_slot] != self._group_id
            if slot_contaminated:
                # Another group (or no group) wrote to our slot — re-fetch
                self._prev_hook.prefetch_layer(self.current_slot, non_blocking=False)
            # Always wait for data to be ready (handles both sync and async paths)
            self._prev_hook.get_weights(self.current_slot)
        elif not self.is_materialized and self._prev_hook is not None:
            # Previous hook was skipped (e.g. by cache-dit), sync-prefetch
            self._prev_hook.prefetch_layer(self.current_slot, non_blocking=False)
            self._prev_hook.get_weights(self.current_slot)

        # Prefetch next layer into the other slot (overlapped with compute).
        # No explicit get_weights() here — offload_layer() in the previous
        # hook's post_forward already enqueued wait_event for the current
        # layer's H2D/AllGather, which is sufficient to ensure compute_stream
        # waits for weights before reading them.  The redundant wait_event
        # caused cascading NPU stream synchronization stalls (~10ms/layer).
        next_slot = 1 - self.current_slot
        self.prefetch_layer(next_slot, non_blocking=True)

        return args, kwargs

    def post_forward(self, module: nn.Module, output: Any) -> Any:
        self.offload_layer()
        return output


# ---------------------------------------------------------------------- #
#  Module-level helpers                                                   #
# ---------------------------------------------------------------------- #


def apply_distributed_block_hook(
    module: nn.Module,
    next_block: nn.Module,
    device: torch.device,
    dp_group: torch.distributed.ProcessGroup | None,
    dp_size: int,
    rank: int,
    copy_stream: Any | None = None,
    comm_stream: Any | None = None,
    pin_memory: bool = True,
    shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] | None = None,
    rank_local_mmap: bool = False,
    tensor_transforms: dict[int, Any] | None = None,
) -> DistributedLayerwiseOffloadHook:
    """Register a DistributedLayerwiseOffloadHook on *module*."""
    registry = HookRegistry.get_or_create(module)
    hook = DistributedLayerwiseOffloadHook(
        next_block=next_block,
        device=device,
        dp_group=dp_group,
        dp_size=dp_size,
        rank=rank,
        copy_stream=copy_stream,
        comm_stream=comm_stream,
        pin_memory=pin_memory,
        shared_buffers=shared_buffers,
        rank_local_mmap=rank_local_mmap,
        tensor_transforms=tensor_transforms,
    )
    registry.register_hook(DistributedLayerwiseOffloadHook._HOOK_NAME, hook)
    return hook


def remove_distributed_block_hook(module: nn.Module) -> None:
    """Remove the distributed layerwise offload hook from *module*."""
    registry: HookRegistry | None = getattr(module, "_hook_registry", None)
    if registry is not None:
        registry.remove_hook(DistributedLayerwiseOffloadHook._HOOK_NAME)
        logger.debug("Removed distributed offload hook from %s", module.__class__.__name__)


class PinnedResidentLayerGroup:
    """Keep selected layers available for stage-scoped device residency.


    TODO(offload): Extract this alongside PinnedModuleStager after the
    distributed shard-and-pin operation becomes a shared storage primitive.
    It currently remains here because it depends on DLO's local-shard layout.
    Unlike ``module.to(device)``/``module.to("cpu")``, this group retains a
    pinned CPU master copy (or mmap source plus bounded staging) and never
    copies generated device weights back to host. Entering the denoise stage
    performs one asynchronous H2D pass;
    leaving it only restores zero-sized placeholders and releases the device
    buffers.  This lets the following VAE stage reuse the same HBM.

    The resident path is intentionally local-shard only.  With tensor
    parallelism, the regular model loader has already produced each rank's TP
    shard, so no DP AllGather is required or desirable here.
    """

    def __init__(
        self,
        blocks: list[nn.Module],
        device: torch.device,
        copy_stream: Any,
        pin_memory: bool,
        rank_local_mmap: bool = False,
        defer_staging: bool = False,
        tensor_transforms: dict[int, Any] | None = None,
    ) -> None:
        self.device = device
        self.copy_stream = copy_stream
        self.loaded = False
        self.rank_local_mmap = rank_local_mmap
        self.pin_memory = pin_memory
        self._states: list[dict[str, Any]] = []
        self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = []

        for block in blocks:
            params = dict(block.named_parameters())
            bufs = dict(block.named_buffers())
            targets: dict[str, torch.Tensor] = {**params, **bufs}
            if rank_local_mmap:
                cpu_sources, metadata = DistributedLayerwiseOffloadHook._collect_mmap_sources(
                    params,
                    bufs,
                    tensor_transforms,
                )
                cpu_shards = {}
            else:
                cpu_shards, metadata = DistributedLayerwiseOffloadHook._shard_and_pin(
                    params,
                    bufs,
                    dp_size=1,
                    rank=0,
                    pin_memory=pin_memory,
                    tensor_transforms=tensor_transforms,
                )
                cpu_sources = {}
            self._states.append(
                {
                    "targets": targets,
                    "cpu_shards": cpu_shards,
                    "cpu_sources": cpu_sources,
                    "metadata": metadata,
                }
            )

        self._cpu_staging_buffers: list[dict[torch.dtype, torch.Tensor]] = []
        self._cpu_staging_events: list[Any | None] = [None, None]
        if rank_local_mmap and not defer_staging:
            max_sizes: dict[torch.dtype, int] = {}
            for state in self._states:
                for dtype, metas in state["metadata"].items():
                    total = sum(meta["numel"] for meta in metas)
                    max_sizes[dtype] = max(max_sizes.get(dtype, 0), total)
            for _ in range(2):
                buffers = {}
                for dtype, total in max_sizes.items():
                    buffer = torch.empty(total, dtype=dtype, device="cpu")
                    if pin_memory:
                        buffer = buffer.pin_memory()
                    buffers[dtype] = buffer
                self._cpu_staging_buffers.append(buffers)

    def load(self) -> None:
        if self.loaded:
            return

        # Allocate on the compute stream so the caching allocator can reuse
        # blocks released by the encoder stage.  Allocating on the copy stream
        # would create a separate stream-local pool and inflate peak HBM.
        gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = []
        for state in self._states:
            block_buffers: dict[torch.dtype, torch.Tensor] = {}
            for dtype, metas in state["metadata"].items():
                total = sum(meta["numel"] for meta in metas)
                block_buffers[dtype] = torch.empty(total, dtype=dtype, device=self.device)
            gpu_buffers.append(block_buffers)

        self.copy_stream.wait_stream(current_omni_platform.current_stream())
        ready = current_omni_platform.Event()
        with current_omni_platform.stream(self.copy_stream):
            for index, (state, block_buffers) in enumerate(zip(self._states, gpu_buffers)):
                if self.rank_local_mmap:
                    slot = index % 2
                    previous_copy = self._cpu_staging_events[slot]
                    if previous_copy is not None:
                        synchronize = getattr(previous_copy, "synchronize", None)
                        if callable(synchronize):
                            synchronize()
                        else:
                            current_omni_platform.synchronize()
                    cpu_weights = DistributedLayerwiseOffloadHook._pack_mmap_sources(
                        state["cpu_sources"],
                        state["metadata"],
                        self._cpu_staging_buffers[slot],
                    )
                else:
                    cpu_weights = state["cpu_shards"]

                for dtype, cpu_weight in cpu_weights.items():
                    block_buffers[dtype].copy_(
                        cpu_weight,
                        non_blocking=cpu_weight.is_pinned(),
                    )
                if self.rank_local_mmap:
                    slot_ready = current_omni_platform.Event()
                    slot_ready.record(self.copy_stream)
                    self._cpu_staging_events[slot] = slot_ready
            ready.record(self.copy_stream)

        for state, block_buffers in zip(self._states, gpu_buffers):
            targets = state["targets"]
            for dtype, metas in state["metadata"].items():
                gpu_buffer = block_buffers[dtype]
                for meta in metas:
                    set_tensor_storage(
                        targets[meta["name"]],
                        torch.as_strided(
                            gpu_buffer[meta["offset"] : meta["offset"] + meta["numel"]],
                            size=meta["shape"],
                            stride=meta["stride"],
                        ),
                    )

        current_omni_platform.current_stream().wait_event(ready)
        self._gpu_buffers = gpu_buffers
        self.loaded = True

    def offload(self) -> None:
        if not self.loaded:
            return

        # Denoise kernels consume these buffers on the current compute stream.
        # Synchronize once at the stage boundary; no D2H copy is necessary
        # because the pinned CPU master weights were never overwritten.
        current_omni_platform.synchronize()
        for state in self._states:
            for target in state["targets"].values():
                set_tensor_storage(target, make_offload_placeholder(target))
        self._gpu_buffers.clear()
        self.loaded = False


# ---------------------------------------------------------------------- #
#  Backend                                                                #
# ---------------------------------------------------------------------- #


class DistributedLayerwiseOffloadBackend(OffloadBackend):
    """Distributed layer-wise (block-level) offloading backend.

    Supports both GPU (CUDA) and NPU (CANN) platforms.
    Device type is determined by the device passed to enable().

    Each rank stores only a shard of each block's weights on host memory.
    Two device slots alternate: one holds current weights, one holds next
    weights. H2D and AllGather run asynchronously on dedicated streams,
    overlapped with computation.
    """

    def __init__(
        self,
        config: OffloadConfig,
        device: torch.device,
        host_weight_plan: HostWeightPlan | None = None,
    ):
        super().__init__(config, device)

        self.copy_stream = current_omni_platform.Stream()
        self.comm_stream = current_omni_platform.Stream()
        self.dp_group: torch.distributed.ProcessGroup | None = None
        self.dp_size = config.dp_size
        self.rank = 0
        self._blocks: list[list[nn.Module]] = []
        self._all_hook_groups: list[list[DistributedLayerwiseOffloadHook]] = []
        self._resident_blocks: list[nn.Module] = []
        self._resident_layer_group: PinnedResidentLayerGroup | None = None
        self._using_mmap = False
        self._using_rank_local_mmap = False
        self.host_weight_plan = host_weight_plan
        self._mmap_transforms_by_tensor_id: dict[int, Any] = {}

    def load_resident_layers(self) -> None:
        """Load the model-declared leading blocks for the denoise stage."""
        if self._resident_layer_group is not None:
            self._resident_layer_group.load()

    def offload_resident_layers(self) -> None:
        """Release leading blocks before VAE decode to bound peak HBM."""
        if self._resident_layer_group is not None:
            self._resident_layer_group.offload()

    def _load_weights_via_mmap(
        self,
        pipeline: nn.Module,
        modules,
        plan: HostWeightPlan,
    ) -> None:
        """Load DiT checkpoint tensors as file-backed safetensors views.

        When the transformer is created on meta device, this method
        replaces meta params with mmap views of the checkpoint files.
        The views point to OS page cache shared across ranks. AllGather mode
        copies only the rank's 1/dp_size shard to a persistent private buffer;
        rank-local mode retains the views and packs one block at a time into
        bounded staging storage.

        Non-DiT modules (VAE, encoders) are NOT affected — they were
        created on CPU with real weights via from_pretrained.
        """
        from safetensors import safe_open

        logger.info(
            "Loading DiT weights via mmap (meta -> shared page cache): %d tensors",
            len(plan.bindings),
        )

        # --- Convert DiT modules to meta device ---
        # The transformer was created normally (with random weights and
        # correct non-persistent buffers from __init__).  We convert it
        # to meta device to release the random weights (they're useless),
        # then load real weights via mmap.
        #
        # Non-persistent buffers (e.g. RoPE inv_freq, timestep freqs) are
        # NOT in the checkpoint — they are computed from formulas in
        # __init__.  Save them before meta conversion and restore after
        # mmap loading, so we don't need model-specific buffer rebuild code.
        #
        # The loader proved topology, coverage, source metadata, and adapter
        # compatibility before it skipped ordinary weight materialization.
        # This method only realizes that exact plan; it does not select or
        # rediscover a checkpoint layout.

        #
        # Important: when DiT modules are nested (e.g. transformer contains
        # transformer.language_model), to_empty("meta") on the parent
        # converts the child's buffers too.  So we must save ALL buffers
        # from ALL DiT modules BEFORE any to_empty call.
        saved_buffers: dict[int, dict[str, torch.Tensor]] = {}
        for dit_module in modules.dits:
            bufs: dict[str, torch.Tensor] = {}
            for name, buf in dit_module.named_buffers():
                # Check if this buffer is non-persistent on its OWNING module
                owner = dit_module
                parts = name.split(".")
                for part in parts[:-1]:
                    owner = getattr(owner, part)
                buf_name = parts[-1]
                is_non_persistent = buf_name in owner._non_persistent_buffers_set
                # Save if non-persistent OR already meta (shouldn't happen
                # before to_empty, but be safe)
                if is_non_persistent:
                    bufs[name] = buf.detach().clone()
            saved_buffers[id(dit_module)] = bufs

        # Now convert all DiT modules to meta (after saving all buffers).
        # Skip modules that are already meta (happens when a parent DiT
        # module contains a child DiT module — to_empty on the parent
        # already converted the child).
        for dit_module in modules.dits:
            if any(p.is_meta for p in dit_module.parameters()):
                logger.info(
                    "%s already on meta device (skipping to_empty, %d buffers saved)",
                    dit_module.__class__.__name__,
                    len(saved_buffers.get(id(dit_module), {})),
                )
                continue
            dit_module.to_empty(device="meta")
            logger.info(
                "Converted %s to meta device (%d non-persistent buffers saved)",
                dit_module.__class__.__name__,
                len(saved_buffers.get(id(dit_module), {})),
            )

        # Cache open file handles
        file_cache: dict[str, Any] = {}
        loaded_names: set[str] = set()

        # Realize the loader's exact runtime-name bindings.  In particular,
        # do not reconstruct names from DLO block discovery: storage planning
        # and transfer topology are separate contracts.
        for runtime_name, binding in plan.bindings.items():
            parent_path, _, leaf_name = runtime_name.rpartition(".")
            try:
                parent = pipeline.get_submodule(parent_path)
            except AttributeError as exc:
                raise RuntimeError(f"Host-weight plan target module {parent_path!r} no longer exists") from exc

            target = parent._parameters.get(leaf_name)
            is_parameter = target is not None
            if target is None:
                target = parent._buffers.get(leaf_name)
            if target is None:
                raise RuntimeError(f"Host-weight plan target tensor {runtime_name!r} no longer exists")

            if binding.file_path not in file_cache:
                file_cache[binding.file_path] = safe_open(
                    binding.file_path,
                    framework="pt",
                    device="cpu",
                )
            tensor = file_cache[binding.file_path].get_tensor(binding.checkpoint_key)
            if is_parameter:
                replacement = torch.nn.Parameter(tensor, requires_grad=target.requires_grad)
                parent._parameters[leaf_name] = replacement
            else:
                replacement = tensor
                parent._buffers[leaf_name] = replacement
            loaded_names.add(runtime_name)

        logger.info("Realized %d loader-planned tensors as mmap views", len(loaded_names))

        # Keep file handles open — _shard_and_pin will read from the mmap views.
        # They will be released after _shard_and_pin completes (when params are
        # replaced with offload placeholders).
        self._mmap_file_cache = file_cache
        # The regular loader runs model-specific post-load transforms after
        # assigning checkpoint tensors. The mmap path bypasses that loader, so
        # preserve the same lifecycle for transforms such as Cosmos3's fp32
        # timestep embedder.
        for dit_name, dit_module in zip(modules.dit_names, modules.dits):
            # Restore non-persistent buffers before post-load hooks and strict
            # validation. They are constructor-derived and intentionally have
            # no checkpoint binding.
            bufs = saved_buffers.get(id(dit_module), {})
            for name, buf in bufs.items():
                parent_path, _, leaf_name = name.rpartition(".")
                parent = dit_module.get_submodule(parent_path)
                restore_device = torch.device("cpu") if self._using_rank_local_mmap else self.device
                parent._buffers[leaf_name] = buf.to(restore_device)

            post_load_weights = getattr(dit_module, "post_load_weights", None)
            if callable(post_load_weights):
                post_load_weights()

            # Call validate_loaded_weights if the model defines it.
            # This preserves the sound-weight and action-weight validation
            # that AutoWeightsLoader.load_weights() triggers in the regular
            # load path (e.g. Cosmos3 checks for missing audio/action weights).
            validate = getattr(dit_module, "validate_loaded_weights", None)
            if callable(validate):
                local_loaded_names = {
                    name.removeprefix(f"{dit_name}.") for name in loaded_names if name.startswith(f"{dit_name}.")
                }
                validate(local_loaded_names)

        # Post-load hooks may rebind parameters (for example while casting a
        # submodule), so associate deferred bounded transforms with the final
        # runtime tensor objects rather than the initial mmap replacements.
        self._mmap_transforms_by_tensor_id.clear()
        for runtime_name, binding in plan.bindings.items():
            if binding.transform is None:
                continue
            parent_path, _, leaf_name = runtime_name.rpartition(".")
            parent = pipeline.get_submodule(parent_path)
            target = parent._parameters.get(leaf_name)
            if target is None:
                target = parent._buffers.get(leaf_name)
            if target is None:
                raise RuntimeError(
                    f"Host-weight transform target {runtime_name!r} no longer exists after post-load processing"
                )
            self._mmap_transforms_by_tensor_id[id(target)] = binding.transform

        remaining_meta: list[str] = []
        for dit_name, dit_module in zip(modules.dit_names, modules.dits):
            remaining_meta.extend(
                f"{dit_name}.{name}" for name, tensor in dit_module.named_parameters() if tensor.is_meta
            )
            for name, tensor in dit_module.named_buffers():
                parent_path, _, leaf_name = name.rpartition(".")
                owner = dit_module.get_submodule(parent_path)
                if leaf_name not in owner._non_persistent_buffers_set and tensor.is_meta:
                    remaining_meta.append(f"{dit_name}.{name}")
        if remaining_meta:
            raise RuntimeError(
                "The prevalidated host-weight plan left "
                f"{len(remaining_meta)} DiT tensors on the meta device "
                f"(first 5: {remaining_meta[:5]})."
            )

    def _init_dp_group(self) -> None:
        """Reuse the process group initialized by parallel_state.

        When DP > 1, uses the DP group (handles strided groups with SP/CFG/PP).
        When DP = 1 but SP > 1 (and dp_size was set to sp_size in OffloadConfig),
        uses the SP group so weights are sharded across SP ranks.
        """
        if self.dp_size <= 1:
            logger.info("Distributed layerwise offload: dp_size=1, running without AllGather")
            self.dp_group = None
            return

        if not torch.distributed.is_initialized():
            raise RuntimeError(
                "torch.distributed is not initialized. "
                "Distributed layerwise offload with dp_size > 1 requires "
                "an initialized process group."
            )

        from vllm_omni.diffusion.distributed.parallel_state import (
            get_data_parallel_world_size,
            get_dp_group,
        )

        # Determine which parallel group to use for sharding.
        # When data_parallel_size > 1, use the DP group.
        # When data_parallel_size == 1 but dp_size > 1 (set from sp_size in
        # OffloadConfig), use the SP group for weight sharding.
        dp_world = get_data_parallel_world_size()
        if dp_world > 1:
            coord = get_dp_group()
        else:
            from vllm_omni.diffusion.distributed.parallel_state import get_sp_group

            coord = get_sp_group()
            logger.info(
                "Distributed layerwise offload: DP=1, using SP group (world_size=%d) for weight sharding",
                coord.world_size,
            )

        self.dp_group = coord.device_group
        self.rank = coord.rank_in_group
        self.dp_size = coord.world_size

        logger.info(
            "Distributed layerwise offload: dp_size=%d, rank_in_group=%d, global_rank=%d, group_ranks=%s",
            self.dp_size,
            self.rank,
            coord.rank,
            coord.ranks,
        )

    def _register_on_demand_hook(self, module: nn.Module, label: str, *, stage_on_demand: bool = False) -> None:
        """Prepare a pipeline-managed stage component or keep it resident.

        Components that expose an explicit stage lifecycle are initially
        offloaded and loaded by their pipeline only around encode/decode.
        Other models retain the conservative resident behavior because a
        generic post-forward hook can disrupt the DiT prefetch streams.
        """
        offload_to_cpu = getattr(module, "offload_to_cpu", None)
        if stage_on_demand and callable(offload_to_cpu):
            offload_to_cpu()
            logger.info("Prepared %s (%s) for pipeline-managed staged offload", label, module.__class__.__name__)
            return
        module.to(self.device)
        logger.info("Moved %s (%s) to GPU (resident)", label, module.__class__.__name__)

    def _try_layerwise_offload_encoder(self, module: nn.Module, name: str, plan: OffloadPlan | None) -> bool:
        """Stream plan-declared encoder blocks on each rank without AllGather."""
        if plan is None or name not in plan.encoder_block_attrs:
            return False
        if getattr(module, "_omni_layerwise_enabled", False):
            return True

        from operator import attrgetter

        from vllm_omni.diffusion.offloader.layerwise_backend import apply_block_hook

        hooks = []
        block_groups = []
        copy_stream = current_omni_platform.Stream()
        for block_path in plan.encoder_block_attrs[name]:
            try:
                blocks = attrgetter(block_path)(module)
            except AttributeError:
                logger.warning("Encoder offload path %s.%s was not found", name, block_path)
                continue
            if not isinstance(blocks, nn.ModuleList) or len(blocks) <= 1:
                logger.warning("Encoder offload path %s.%s is not a streamable block list", name, block_path)
                continue
            group_hooks = [
                apply_block_hook(blocks[-1], blocks[0], self.device, copy_stream, self.config.pin_cpu_memory)
            ]
            group_hooks.extend(
                apply_block_hook(block, blocks[index + 1], self.device, copy_stream, self.config.pin_cpu_memory)
                for index, block in enumerate(blocks[:-1])
            )
            for index, hook in enumerate(group_hooks):
                hook._prev_hook = group_hooks[index - 1]
            hooks.extend(group_hooks)
            block_groups.append(blocks)

        if not hooks:
            return False
        # The component lifecycle uses these generic attributes to keep only
        # non-block encoder state resident during the encode phase.
        module._omni_layerwise_hooks = hooks
        module._omni_layerwise_block_groups = block_groups
        module._omni_layerwise_enabled = True
        logger.info(
            "Enabled rank-local layerwise offload for encoder %s (%d blocks across %d stacks)",
            name,
            sum(len(blocks) for blocks in block_groups),
            len(block_groups),
        )
        return True

    def _try_layerwise_offload_submodule(self, module: nn.Module, name: str, plan: OffloadPlan | None = None) -> bool:
        """Try to apply layerwise offload to a large submodule's blocks.
        Resolution order:
        1. OffloadPlan.offload_submodules (declarative, if plan is provided)
        2. Heuristic search for common block-list attributes

        Returns True if layerwise offload was applied, False otherwise.
        """
        from operator import attrgetter

        blocks = None
        blocks_attr = None

        # 1. Check OffloadPlan first (declarative — no guessing)
        if plan is not None and name in plan.offload_submodules:
            attr_name = plan.offload_submodules[name]
            try:
                candidate = attrgetter(attr_name)(module)
                if isinstance(candidate, nn.ModuleList) and len(candidate) > 1:
                    blocks = candidate
                    blocks_attr = attr_name
            except AttributeError:
                logger.warning(
                    "OffloadPlan declared block attr '%s' for submodule '%s' "
                    "but attribute not found — falling back to heuristic",
                    attr_name,
                    name,
                )

        # 2. Fallback: heuristic search
        if blocks is None:
            for attr_name in ("layers", "blocks", "h", "model.layers"):
                try:
                    candidate = attrgetter(attr_name)(module)
                except AttributeError:
                    continue
                if isinstance(candidate, nn.ModuleList) and len(candidate) > 1:
                    blocks = candidate
                    blocks_attr = attr_name
                    break

        if blocks is None:
            return False

        num_blocks = len(blocks)
        logger.info(
            "Distributed layerwise offload for submodule '%s.%s' (%d blocks, %.0f MB total, dp_size=%d)",
            name,
            blocks_attr,
            num_blocks,
            sum(p.nelement() * p.element_size() for p in module.parameters()) / 1048576,
            self.dp_size,
        )

        # Move non-block parts of the submodule to GPU (small: embeddings, norms)
        for child_name, child in module.named_children():
            if child_name != blocks_attr:
                child.to(self.device)

        # Apply distributed hooks (1/4 sharding + AllGather, same as DiT)
        # Pass shared_buffers=[None,None] to defer per-hook allocation
        # (prevents OOM on large models where N hooks × 2 buffers >> HBM)
        last_block, first_block = blocks[-1], blocks[0]
        last_hook = apply_distributed_block_hook(
            last_block,
            first_block,
            self.device,
            self.dp_group,
            self.dp_size,
            self.rank,
            self.copy_stream,
            self.comm_stream,
            self.config.pin_cpu_memory,
            shared_buffers=[None, None],
            rank_local_mmap=self._using_rank_local_mmap,
            tensor_transforms=self._mmap_transforms_by_tensor_id,
        )
        sub_hooks = [last_hook]
        for i, block in enumerate(blocks[:-1]):
            next_block = blocks[(i + 1) % num_blocks]
            hook = apply_distributed_block_hook(
                block,
                next_block,
                self.device,
                self.dp_group,
                self.dp_size,
                self.rank,
                self.copy_stream,
                self.comm_stream,
                self.config.pin_cpu_memory,
                shared_buffers=[None, None],
                rank_local_mmap=self._using_rank_local_mmap,
                tensor_transforms=self._mmap_transforms_by_tensor_id,
            )
            sub_hooks.append(hook)

        # Wire backward references + slot alternation
        for i in range(len(sub_hooks)):
            sub_hooks[i]._prev_hook = sub_hooks[i - 1]
        # Assign slots in list order: sub_hooks = [last_hook, block0, ..., blockN-2]
        # This ensures last_hook.current_slot != block0_hook.current_slot,
        # so the circular prefetch (last_hook -> block0) writes to a
        # different slot than block0 reads from.  Correct for ALL N.
        for i, hook in enumerate(sub_hooks):
            hook.current_slot = i % 2
        # Mark block0_hook (index 1) as group-first — it must sync-prefetch
        # on entry because another group may have overwritten the shared slot.
        if len(sub_hooks) > 1:
            sub_hooks[1]._is_group_first = True

        # Defer buffer allocation and prefetch to enable() unified allocation
        self._all_hook_groups.append(sub_hooks)
        self._blocks.append(blocks)
        return True

    def _prepare_dit_non_block_modules(
        self,
        dit_module: nn.Module,
        dit_name: str,
        blocks_attr_names: list[str],
        all_dit_modules: set[int],
        plan: OffloadPlan | None,
    ) -> None:
        """Place or hook the DiT parts that are outside its repeated blocks.

        This must run even when every repeated block is resident.  Otherwise
        an all-resident stage skips placement for modules such as H3's token
        refiner and enters the forward pass with CPU or meta tensors.
        """
        _ON_DEMAND_THRESHOLD = _ON_DEMAND_THRESHOLD_MB
        for name, module in dit_module.named_children():
            if name in blocks_attr_names:
                logger.debug("Skipped blocks module %s", name)
                continue

            module_mb = (
                sum(
                    param.nelement() * param.element_size() if not getattr(param, "is_meta", False) else 0
                    for param in module.parameters()
                )
                / 1048576
            )
            explicitly_planned = plan is not None and name in plan.offload_submodules
            if explicitly_planned or module_mb > _ON_DEMAND_THRESHOLD:
                if id(module) in all_dit_modules:
                    logger.info("Submodule '%s' is already a DiT module, skipping layerwise offload", name)
                elif self._try_layerwise_offload_submodule(module, name, plan):
                    pass
                else:
                    self._register_on_demand_hook(module, name)
                continue

            try:
                module.to(self.device)
            except (NotImplementedError, RuntimeError):
                # Non-persistent buffers such as RoPE frequencies do not
                # exist in the checkpoint and must be reconstructed.
                has_meta_buffer = any(getattr(buffer, "is_meta", False) for buffer in module.buffers(recurse=True))
                if has_meta_buffer:
                    saved_params = {
                        param_name: param.data.clone()
                        for param_name, param in module.named_parameters()
                        if not getattr(param, "is_meta", False)
                    }
                    module.to_empty(device=self.device)
                    for submodule in module.modules():
                        if hasattr(submodule, "reset_parameters"):
                            submodule.reset_parameters()
                    for param_name, param in module.named_parameters():
                        if param_name in saved_params:
                            param.data.copy_(saved_params[param_name])
                try:
                    module.to(self.device)
                except Exception:
                    logger.warning("Module %s still has meta params after mmap load", name)

        for param in dit_module._parameters.values():
            if param is not None and not getattr(param, "is_meta", False):
                param.data = param.data.to(self.device, non_blocking=True)
        for buffer in dit_module._buffers.values():
            if buffer is not None:
                buffer.data = buffer.data.to(self.device, non_blocking=True)

    def enable(self, pipeline: nn.Module) -> None:
        if self.enabled:
            logger.warning("DistributedLayerwiseOffloadBackend already enabled")
            return

        self._on_demand_shard_infos: list[dict] = []
        self._on_demand_handles: list[Any] = []

        # Initialize DP group (if not already done by early init)
        if self.dp_group is None and self.dp_size > 1:
            self._init_dp_group()

        modules = ModuleDiscovery.discover(pipeline)
        if not modules.dits:
            if self.host_weight_plan is not None:
                raise RuntimeError(
                    "DLO received a loader-owned host-weight plan, but no DiT modules were discovered to consume it"
                )
            logger.warning("No DiT/transformer modules found, skipping distributed layer-wise offloading")
            return

        # Retrieve optional declarative OffloadPlan from the pipeline.
        # When present, replaces heuristic block discovery.
        plan = get_offload_plan(pipeline)

        if self.config.dlo_resident_layers and (plan is None or not plan.resident_dit_paths):
            logger.warning(
                "dlo_resident_layers=%d was requested, but this model declares no "
                "resident_dit_paths; all blocks will be streamed.",
                self.config.dlo_resident_layers,
            )

        # Storage selection belongs to the loader.  DLO consumes the exact
        # prevalidated plan that caused the loader to skip materialization;
        # without a plan, all weights must already come from the ordinary
        # loader.  The transfer protocol is selected independently below.
        self._using_mmap = self.host_weight_plan is not None
        self._using_rank_local_mmap = self._using_mmap and self.dp_size <= 1
        if self._using_mmap:
            if self.host_weight_plan.backing_kind != "checkpoint_mmap":
                raise ValueError(f"Unsupported DLO host-weight backing: {self.host_weight_plan.backing_kind}")
            self._load_weights_via_mmap(
                pipeline,
                modules,
                self.host_weight_plan,
            )
            if self._using_rank_local_mmap:
                logger.info(
                    "DLO rank-local mmap storage enabled: checkpoint pages are "
                    "node-shared; each worker owns only two bounded host staging slots"
                )
        else:
            remaining_meta = [
                name
                for dit_name, dit_module in zip(modules.dit_names, modules.dits)
                for name, tensor in chain(
                    dit_module.named_parameters(),
                    dit_module.named_buffers(),
                )
                if getattr(tensor, "is_meta", False)
            ]
            if remaining_meta:
                raise RuntimeError(
                    f"DLO received meta tensors without a loader-owned host-weight plan (first 5: {remaining_meta[:5]})"
                )
            logger.info("DLO is using host tensors materialized by the ordinary loader")

        # Keep VAE/encoders on CPU; move to GPU on-demand via hooks.
        # This saves ~4.3 GB HBM per card (VAE 1.3 + encoder 1.1 + sound 1.9)
        # during the DiT forward pass.  They are only needed briefly for
        # text-encoding (before DiT) and VAE-decode (after DiT).
        for enc, enc_name in zip(modules.encoders, modules.encoder_names):
            self._try_layerwise_offload_encoder(enc, enc_name, plan)
            self._register_on_demand_hook(
                enc, "encoder", stage_on_demand=plan is not None and enc_name in plan.on_demand_component_paths
            )
        for vae, vae_name in zip(modules.vaes, modules.vae_names):
            self._register_on_demand_hook(
                vae,
                "vae",
                stage_on_demand=(plan is not None and vae_name in plan.on_demand_component_paths),
            )

        # Move resident modules to GPU (small modules needed every forward)
        for name, module in zip(modules.resident_names, modules.resident_modules):
            try:
                module.to(self.device)
            except Exception as exc:
                logger.debug("Failed to move resident module %s to GPU: %s", name, exc)

        logger.info("Applying distributed layer-wise offloading on %s", modules.dit_names)

        # Collect all DiT module objects to detect submodules that are
        # already handled as a separate DiT module (avoids duplicate hooks).
        all_dit_modules = set(id(m) for m in modules.dits)

        # Apply hooks for each DiT module
        for i, dit_module in enumerate(modules.dits):
            dit_name = modules.dit_names[i]
            logger.info(f"Applying hooks on {dit_name} ({dit_module.__class__.__name__})")

            blocks_attr_names, blocks = get_blocks_from_dit(dit_module)

            if not blocks:
                logger.warning(
                    "Target layers (blocks) not found. Skipping offloading on %s (%s)",
                    dit_name,
                    dit_module.__class__.__name__,
                )
                dit_module.to(self.device)
                continue

            resident_count = 0
            if plan is not None and dit_name in plan.resident_dit_paths:
                resident_count = min(self.config.dlo_resident_layers, len(blocks))
            if resident_count:
                resident_blocks = blocks[:resident_count]
                self._resident_blocks.extend(resident_blocks)
                blocks = blocks[resident_count:]
                logger.info(
                    "Keeping %d leading blocks resident on %s; streaming %d tail blocks",
                    resident_count,
                    dit_name,
                    len(blocks),
                )

            self._prepare_dit_non_block_modules(
                dit_module,
                dit_name,
                blocks_attr_names,
                all_dit_modules,
                plan,
            )

            num_blocks = len(blocks)
            if num_blocks == 0:
                logger.info("All blocks for %s are resident; no streaming hooks required", dit_name)
                continue
            if num_blocks <= 1:
                logger.warning(
                    "#Streaming target layers (blocks) <= 1. Keeping the final block resident on %s (%s)",
                    dit_name,
                    dit_module.__class__.__name__,
                )
                self._resident_blocks.extend(blocks)
                continue

            # Register hooks in a circular sliding window:
            # last block prefetches first block, block i prefetches block (i+1)
            # All hooks share 2 global device buffers (RFC: "exactly two layers on device")
            last_block, first_block = blocks[-1], blocks[0]

            # Pass 1: create hooks with deferred buffer allocation
            # (shared_buffers=[None,None] prevents per-hook OOM on large models)
            last_hook = apply_distributed_block_hook(
                last_block,
                first_block,
                self.device,
                self.dp_group,
                self.dp_size,
                self.rank,
                self.copy_stream,
                self.comm_stream,
                self.config.pin_cpu_memory,
                shared_buffers=[None, None],
                rank_local_mmap=self._using_rank_local_mmap,
                tensor_transforms=self._mmap_transforms_by_tensor_id,
            )

            block_hooks: list[DistributedLayerwiseOffloadHook] = [last_hook]
            for i, block in enumerate(blocks[:-1]):
                next_block = blocks[(i + 1) % num_blocks]
                hook = apply_distributed_block_hook(
                    block,
                    next_block,
                    self.device,
                    self.dp_group,
                    self.dp_size,
                    self.rank,
                    self.copy_stream,
                    self.comm_stream,
                    self.config.pin_cpu_memory,
                    shared_buffers=[None, None],
                    rank_local_mmap=self._using_rank_local_mmap,
                    tensor_transforms=self._mmap_transforms_by_tensor_id,
                )
                block_hooks.append(hook)

            # Wire backward references for cache-dit fallback
            for i in range(len(block_hooks)):
                block_hooks[i]._prev_hook = block_hooks[i - 1]

            # Assign slots in list order: block_hooks = [last_hook, block0, ..., blockN-2]
            # This ensures last_hook.current_slot != block0_hook.current_slot,
            # so the circular prefetch (last_hook -> block0) writes to a
            # different slot than block0 reads from.  Correct for ALL N.
            for i, hook in enumerate(block_hooks):
                hook.current_slot = i % 2

            # Mark block0_hook (index 1) as group-first
            if len(block_hooks) > 1:
                block_hooks[1]._is_group_first = True

            # Defer buffer allocation — collected for unified allocation below
            self._all_hook_groups.append(block_hooks)
            self._blocks.append(blocks)

        if self._resident_blocks:
            self._resident_layer_group = PinnedResidentLayerGroup(
                self._resident_blocks,
                self.device,
                self.copy_stream,
                self.config.pin_cpu_memory,
                rank_local_mmap=self._using_rank_local_mmap,
                defer_staging=bool(self._all_hook_groups),
                tensor_transforms=self._mmap_transforms_by_tensor_id,
            )
            pipeline._dlo_residency_controller = self

        if not self._all_hook_groups:
            self.enabled = bool(self._resident_blocks)
            if self._using_mmap and not self.enabled:
                self._release_mmap_handles()
            return

        # Unified allocation: 2 shared output buffers + 2 shared shard buffers
        # sized to the max block across ALL module groups (gen_layers +
        # language_model).  Groups execute sequentially, so 2 buffers suffice.
        all_hooks: list[DistributedLayerwiseOffloadHook] = []
        for group in self._all_hook_groups:
            all_hooks.extend(group)

        unified_buffers = self._allocate_shared_buffers(all_hooks)
        unified_shard_buffers = None
        if self.dp_size > 1:
            unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks)
        unified_cpu_staging = None
        cpu_staging_events = None
        if self._using_rank_local_mmap:
            unified_cpu_staging = self._allocate_shared_cpu_staging_buffers(
                all_hooks,
                self._resident_layer_group,
            )
            cpu_staging_events = [None, None]
            if self._resident_layer_group is not None:
                # Resident and streamed layers execute in the same stage and
                # reuse the same host slots. Events serialize slot reuse.
                self._resident_layer_group._cpu_staging_buffers = [
                    buffers for buffers in unified_cpu_staging if buffers is not None
                ]
                self._resident_layer_group._cpu_staging_events = cpu_staging_events

        # Shared slot-group tracker: _shared_slot_group[slot] = group_id
        # that last wrote to that slot.  Group-first hooks use this to
        # skip sync-prefetch when the slot still contains their own data.
        shared_slot_group = [-1, -1]

        for group_idx, group in enumerate(self._all_hook_groups):
            for hook in group:
                hook.gpu_buffers = unified_buffers
                hook._owns_buffers = False
                if unified_shard_buffers is not None:
                    hook.gpu_shard_buffers = unified_shard_buffers
                if unified_cpu_staging is not None and cpu_staging_events is not None:
                    hook.cpu_staging_buffers = unified_cpu_staging
                    hook.cpu_staging_events = cpu_staging_events
                hook._group_id = group_idx
                hook._shared_slot_group = shared_slot_group

        # Prefetch first block of the FIRST module group only.
        # Subsequent groups share the same 2 device buffers; prefetching
        # them now would overwrite the first group's data in the shared
        # buffer (both groups default to slot 0).  Instead, subsequent
        # groups' first blocks remain as meta placeholders, and their
        # pre_forward will sync-prefetch on-demand via the is_materialized
        # check — by which point the first group's forward has completed
        # and its buffer slots are free.
        if self._all_hook_groups:
            group = self._all_hook_groups[0]
            first_slot = group[0].current_slot
            group[-1].prefetch_layer(slot=first_slot, non_blocking=False)
            group[-1].get_weights(first_slot)

        total_blocks = sum(len(b) for b in self._blocks)
        logger.info(
            f"Distributed layer-wise offloading enabled on {total_blocks} blocks "
            f"across {len(self._all_hook_groups)} group(s), dp_size={self.dp_size}, "
            f"unified shared_buffers=2"
        )

        self.enabled = True

        if self._using_mmap and not self._using_rank_local_mmap:
            # AllGather mode copied each rank's persistent shard, so the source
            # mappings are no longer needed. Rank-local mode retains them as
            # the node-shared host master until disable().
            self._release_mmap_handles()

        # Assign GPU buffers to sharded on-demand modules (VAE/encoders).
        # Each module gets a dedicated input (shard-sized) and output
        # (full-sized) buffer.  VAE/encoders run before/after DiT, never
        # concurrently, so peak HBM = max(DiT, VAE) not sum.
        for si in self._on_demand_shard_infos:
            out_bufs: dict[torch.dtype, torch.Tensor] = {}
            in_bufs: dict[torch.dtype, torch.Tensor] = {}
            for dtype, shard in si["cpu_shards"].items():
                # AllGather output = dp_size * shard_size
                full_size = shard.numel() * self.dp_size if self.dp_size > 1 else shard.numel()
                out_bufs[dtype] = torch.empty(full_size, dtype=dtype, device=self.device)
                in_bufs[dtype] = torch.empty(shard.shape, dtype=dtype, device=self.device)
            si["gpu_output"] = out_bufs
            si["gpu_input"] = in_bufs
            _mb = sum(t.nelement() * t.element_size() for t in out_bufs.values()) / 1048576
            logger.info("Allocated %.0f MB GPU buffer for sharded on-demand module", _mb)

        self._cleanup_after_loading()

    def _release_mmap_handles(self) -> None:
        """Release safetensors mmap file handles."""
        self._mmap_transforms_by_tensor_id.clear()
        if hasattr(self, "_mmap_file_cache"):
            self._mmap_file_cache.clear()
            del self._mmap_file_cache
            logger.info("Released safetensors mmap file handles")

    def _cleanup_after_loading(self) -> None:
        """Synchronize and release freed device/CPU memory after sharding."""
        current_omni_platform.synchronize()
        current_omni_platform.empty_cache()
        import ctypes as _ctypes
        import gc as _gc

        _gc.collect()
        try:
            _ctypes.CDLL("libc.so.6").malloc_trim(0)
        except Exception:
            pass

    def disable(self) -> None:
        if not self.enabled and not hasattr(self, "_mmap_file_cache"):
            return

        if self._using_rank_local_mmap:
            current_omni_platform.synchronize()

        for blocks in self._blocks:
            for block in blocks:
                remove_distributed_block_hook(block)

        for h in getattr(self, "_on_demand_handles", []):
            h.remove()
        self._on_demand_handles = []
        self._on_demand_shard_infos = []

        self.offload_resident_layers()
        self._blocks.clear()
        self._all_hook_groups.clear()
        self._resident_blocks.clear()
        self._resident_layer_group = None
        self._release_mmap_handles()
        self._using_mmap = False
        self._using_rank_local_mmap = False
        self.enabled = False
        logger.info("Distributed layer-wise offloading disabled")

    @staticmethod
    def _allocate_shared_buffers(
        hooks: list[DistributedLayerwiseOffloadHook],
    ) -> list[dict[torch.dtype, torch.Tensor] | None]:
        """Allocate exactly 2 shared device buffers sized to the largest block.

        All hooks share these 2 buffers. At any time, slot 0 holds the current
        layer's weights and slot 1 holds the next layer's weights (or vice
        versa). This ensures only 2 layers' worth of weights reside on device,
        regardless of the total number of blocks.
        """
        max_sizes: dict[torch.dtype, int] = {}
        for hook in hooks:
            dp = hook.dp_size
            for dtype, metas in hook.metadata.items():
                total = sum(m["numel"] for m in metas)
                # AllGather output = dp * ceil(total/dp) (padded for equal shards)
                if dp > 1:
                    total = ((total + dp - 1) // dp) * dp
                if dtype not in max_sizes or total > max_sizes[dtype]:
                    max_sizes[dtype] = total

        device = hooks[0].device
        shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] = [None, None]
        for slot in range(2):
            gpu_weights: dict[torch.dtype, torch.Tensor] = {}
            for dtype, total_numel in max_sizes.items():
                gpu_weights[dtype] = torch.empty(total_numel, dtype=dtype, device=device)
            shared_buffers[slot] = gpu_weights

        logger.info(
            "Allocated 2 shared device buffers (max block size: %s)",
            {str(k): f"{v * _dtype_size(k) / 1024 / 1024:.1f}MB" for k, v in max_sizes.items()},
        )
        return shared_buffers

    @staticmethod
    def _allocate_shared_cpu_staging_buffers(
        hooks: list[DistributedLayerwiseOffloadHook],
        resident_group: PinnedResidentLayerGroup | None = None,
    ) -> list[dict[torch.dtype, torch.Tensor] | None]:
        """Allocate two bounded host slots for rank-local mmap -> device copies."""
        max_sizes: dict[torch.dtype, int] = {}
        for hook in hooks:
            for dtype, metas in hook.metadata.items():
                total = sum(meta["numel"] for meta in metas)
                max_sizes[dtype] = max(max_sizes.get(dtype, 0), total)
        if resident_group is not None:
            for state in resident_group._states:
                for dtype, metas in state["metadata"].items():
                    total = sum(meta["numel"] for meta in metas)
                    max_sizes[dtype] = max(max_sizes.get(dtype, 0), total)

        pin_memory = hooks[0].pin_memory if hooks else bool(resident_group and resident_group.pin_memory)
        shared_staging: list[dict[torch.dtype, torch.Tensor] | None] = [None, None]
        for slot in range(2):
            buffers: dict[torch.dtype, torch.Tensor] = {}
            for dtype, total_numel in max_sizes.items():
                buffer = torch.empty(total_numel, dtype=dtype, device="cpu")
                if pin_memory:
                    buffer = buffer.pin_memory()
                buffers[dtype] = buffer
            shared_staging[slot] = buffers

        logger.info(
            "Allocated 2 shared host staging buffers for rank-local mmap (max block size: %s, pinned=%s)",
            {str(k): f"{v * _dtype_size(k) / 1024 / 1024:.1f}MB" for k, v in max_sizes.items()},
            pin_memory,
        )
        return shared_staging

    @staticmethod
    def _allocate_shared_shard_buffers(
        hooks: list[DistributedLayerwiseOffloadHook],
    ) -> list[dict[torch.dtype, torch.Tensor] | None]:
        """Allocate 2 shared shard (AllGather input) buffers sized to the
        largest per-rank shard.

        All hooks share these 2 input buffers — the same device address is
        reused for every layer's AllGather so HCCL can reuse its internal
        communication buffers (preventing rank-dependent HBM imbalance).
        Cost: 2 × max_shard_size (~184 MB) instead of 2 × N_hooks × shard.
        """
        max_shard_sizes: dict[torch.dtype, int] = {}
        for hook in hooks:
            for dtype, shard in hook.cpu_shards.items():
                numel = shard.numel()
                if dtype not in max_shard_sizes or numel > max_shard_sizes[dtype]:
                    max_shard_sizes[dtype] = numel

        device = hooks[0].device
        shared_shard_buffers: list[dict[torch.dtype, torch.Tensor] | None] = [None, None]
        for slot in range(2):
            shard_bufs: dict[torch.dtype, torch.Tensor] = {}
            for dtype, numel in max_shard_sizes.items():
                shard_bufs[dtype] = torch.empty(numel, dtype=dtype, device=device)
            shared_shard_buffers[slot] = shard_bufs

        logger.info(
            "Allocated 2 shared shard buffers (max shard size: %s)",
            {str(k): f"{v * _dtype_size(k) / 1024 / 1024:.1f}MB" for k, v in max_shard_sizes.items()},
        )
        return shared_shard_buffers