Refactor offload

configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_with_fp8_vae_sla.json

@@ -8,7 +8,6 @@
   "enable_cfg": false,
   "cpu_offload": true,
   "offload_granularity": "block",
-  "dit_prepost_resident": true,
   "text_encoder_cpu_offload": true,
   "text_encoder_offload_granularity": "block",
   "vae_cpu_offload": false,

configs/offload/block/wan_t2v_block.json

@@ -17,7 +17,13 @@
     "clip_quantized": true,
     "clip_quant_scheme": "fp8-q8f",
     "cpu_offload": true,
-    "offload_granularity": "block",
+    "offload_plan": {
+        "offload_granularity": "block",
+        "resident_blocks": {
+            "blocks": 20
+        },
+        "use_event_offload": false
+    },
     "t5_cpu_offload": false,
     "vae_cpu_offload": false,
     "clip_cpu_offload": false

configs/platforms/ascend_npu/flux2_dev_t2i_1344x768.json

@@ -13,12 +13,15 @@
     "text_encoder_out_layers": [10, 20, 30],
     "attn_type": "npu_flash_attn",
     "cpu_offload": true,
-    "offload_granularity": "block",
-    "use_event_offload": true,
-    "offload_use_block_slab": true,
-    "offload_resident_double_blocks": "all",
-    "offload_resident_single_blocks": 36,
-    "offload_resident_policy": "interleaved",
+    "offload_plan": {
+        "offload_granularity": "block",
+        "resident_blocks": {
+            "double_blocks": "all",
+            "single_blocks": 36
+        },
+        "use_event_offload": true,
+        "use_block_slab": true
+    },
     "modulate_type": "torch",
     "rope_type": "npu_rope",
     "layer_norm_type": "npu_layer_norm",

configs/platforms/metax/wan_t2v_1_3b_resident.json

@@ -0,0 +1,27 @@
+{
+    "infer_steps": 4,
+    "target_video_length": 81,
+    "text_len": 512,
+    "target_height": 480,
+    "target_width": 832,
+    "self_attn_1_type": "metax_sage_attn2",
+    "cross_attn_1_type": "metax_sage_attn2",
+    "cross_attn_2_type": "metax_sage_attn2",
+    "sample_guide_scale": 6.0,
+    "sample_shift": 8,
+    "enable_cfg": true,
+    "cpu_offload": true,
+    "offload_plan": {
+        "offload_granularity": "block",
+        "resident_blocks": {
+            "blocks": 28
+        },
+        "use_event_offload": true
+    },
+    "t5_cpu_offload": false,
+    "vae_cpu_offload": false,
+    "modulate_type": "torch",
+    "rope_type": "torch_complex_rope",
+    "layer_norm_type": "torch",
+    "rms_norm_type": "torch"
+}

lightx2v/common/modules/weight_module.py

@@ -1,6 +1,18 @@
+from loguru import logger
+
+from lightx2v.common.offload.config import get_offload_granularity, get_offload_plan
 from lightx2v_platform.base.global_var import AI_DEVICE
 
 
+def resolve_resident_block_indices(resident_count, blocks_num):
+    count = blocks_num if resident_count == "all" else resident_count
+    if not 0 <= count <= blocks_num:
+        raise ValueError(f"resident block count must be between 0 and {blocks_num}, got {count}")
+    if count == 0:
+        return frozenset()
+    return frozenset((index * blocks_num) // count for index in range(count))
+
+
 class WeightModule:
     def __init__(self):
         self._modules = {}
@@ -187,10 +199,98 @@ def to_cuda_async(self, non_blocking=True):
                 if module is not None and hasattr(module, "to_cuda"):
                     module.to_cuda(non_blocking=True)
 
+    def release_device_weights(self):
+        self.to_cpu()
+
+    def release_non_block_weights(self):
+        block_groups = tuple(self.get_offload_block_groups().values())
+        excluded_modules = {id(blocks) for blocks in block_groups}
+        for blocks in block_groups:
+            for name in ("offload_cuda_buffers", "offload_cpu_buffers"):
+                buffers = getattr(blocks, name, None)
+                if buffers is not None:
+                    excluded_modules.add(id(buffers))
+
+        for module in self._modules.values():
+            if module is not None and id(module) not in excluded_modules and hasattr(module, "to_cpu"):
+                module.to_cpu()
+        for name, parameter in self._parameters.items():
+            if parameter is None or id(parameter) in excluded_modules:
+                continue
+            if hasattr(parameter, "cpu"):
+                self._parameters[name] = parameter.to("cpu")
+                setattr(self, name, self._parameters[name])
+            elif hasattr(parameter, "to_cpu"):
+                parameter.to_cpu()
+
+    def register_offload_block_group(self, config, name, blocks):
+        granularity = get_offload_granularity(config)
+        if not config.get("cpu_offload", False):
+            resident_count = "all"
+        elif granularity == "block":
+            resident_count = get_offload_plan(config).get("resident_blocks", {}).get(name, 0)
+        else:
+            resident_count = 0
+
+        resident_indices = resolve_resident_block_indices(resident_count, len(blocks))
+        blocks.offload_group_name = name
+        blocks.resident_block_indices = resident_indices
+        blocks.offload_block_indices = tuple(index for index in range(len(blocks)) if index not in resident_indices)
+
+        if not hasattr(self, "_offload_block_groups"):
+            self._offload_block_groups = {}
+        self._offload_block_groups[name] = blocks
+        slot_count = min(2, len(blocks.offload_block_indices)) if granularity == "block" and config.get("cpu_offload", False) else 0
+        if config.get("cpu_offload", False) and granularity == "block":
+            logger.info(
+                "Block offload group '{}': resident={}/{}, indices={}, staging_slots={}",
+                name,
+                len(resident_indices),
+                len(blocks),
+                tuple(sorted(resident_indices)),
+                slot_count,
+            )
+        return slot_count
+
+    def register_offload_block_buffers(self, name, cuda_buffers, cpu_buffers=None):
+        blocks = self._offload_block_groups[name]
+        blocks.offload_cuda_buffers = cuda_buffers
+        blocks.offload_cpu_buffers = cpu_buffers
+
+    def validate_offload_block_groups(self, config):
+        if not config.get("cpu_offload", False) or get_offload_granularity(config) != "block":
+            return
+
+        plan = get_offload_plan(config)
+        configured_groups = set(plan.get("resident_blocks", {}))
+        registered_groups = set(self.get_offload_block_groups())
+        unknown_groups = configured_groups - registered_groups
+        if unknown_groups:
+            raise ValueError(f"Unknown resident block groups {sorted(unknown_groups)}; available groups are {sorted(registered_groups)}")
+
+    def get_offload_block_groups(self):
+        return getattr(self, "_offload_block_groups", {})
+
+    def resident_blocks_to_cuda(self, non_blocking=True):
+        for blocks in self.get_offload_block_groups().values():
+            for block_index in blocks.resident_block_indices:
+                blocks[block_index].to_cuda(non_blocking=non_blocking)
+
+    def release_resident_blocks(self):
+        for blocks in self.get_offload_block_groups().values():
+            for block_index in blocks.resident_block_indices:
+                blocks[block_index].release_device_weights()
+
 
 class WeightModuleList(WeightModule):
     def __init__(self, modules=None):
         super().__init__()
+        # Transformer block lists receive these values when registered for offload.
+        self.offload_group_name = None
+        self.resident_block_indices = frozenset()
+        self.offload_block_indices = ()
+        self.offload_cuda_buffers = None
+        self.offload_cpu_buffers = None
         self._list = []
         if modules is not None:
             for idx, module in enumerate(modules):

lightx2v/common/offload/config.py

@@ -0,0 +1,29 @@
+def get_offload_plan(config):
+    plan = config.get("offload_plan")
+    if plan is not None:
+        return plan
+
+    return {
+        "offload_granularity": config.get("offload_granularity", "block"),
+        "use_event_offload": config.get("use_event_offload", False),
+        "resident_blocks": {},
+    }
+
+
+def get_offload_granularity(config):
+    return get_offload_plan(config).get("offload_granularity", "block")
+
+
+def use_event_offload(config):
+    return get_offload_plan(config).get("use_event_offload", False)
+
+
+def normalize_offload_plan(config):
+    if "offload_plan" in config:
+        plan = dict(config["offload_plan"])
+    else:
+        plan = get_offload_plan(config)
+    plan.setdefault("offload_granularity", "block")
+    plan.setdefault("use_event_offload", False)
+    plan.setdefault("resident_blocks", {})
+    config["offload_plan"] = plan

lightx2v/common/offload/event_manager.py

@@ -10,9 +10,9 @@
 class EventSlotWeightAsyncStreamManager(WeightAsyncStreamManager):
     """Weight offload with reusable buffers protected by device events."""
 
-    _EVENT_SLOT_COUNT = 2
+    uses_events = True
 
-    def __init__(self, offload_granularity, load_stream=None, compute_stream=None):
+    def __init__(self, offload_granularity, slot_count=2, load_stream=None, compute_stream=None):
         if offload_granularity != "block":
             raise ValueError("Event-slot weight offload only supports block granularity")
 
@@ -22,17 +22,24 @@ def __init__(self, offload_granularity, load_stream=None, compute_stream=None):
         if compute_stream is not None:
             self.compute_stream = compute_stream
 
-        self.device_module = torch_device_module
-        self._ready_events = [torch_device_module.Event() for _ in range(self._EVENT_SLOT_COUNT)]
-        self._free_events = [torch_device_module.Event() for _ in range(self._EVENT_SLOT_COUNT)]
+        self._slot_count = slot_count
+        self._ready_events = []
+        self._free_events = []
+        self.reset_slots()
+
+    def init_cuda_buffer(self, blocks_cuda_buffer=None, phases_cuda_buffer=None):
+        super().init_cuda_buffer(blocks_cuda_buffer, phases_cuda_buffer)
+        self._slot_count = len(self.cuda_buffers)
+        self._ready_events = [torch_device_module.Event() for _ in range(self.slot_count)]
+        self._free_events = [torch_device_module.Event() for _ in range(self.slot_count)]
         self.reset_slots()
 
     @property
     def slot_count(self):
-        return self._EVENT_SLOT_COUNT
+        return self._slot_count
 
     def reset_slots(self):
-        """Reset slot bookkeeping; synchronize pending device work first."""
+        """Reset slot bookkeeping after pending device work has completed."""
         self._slot_pending = [False] * self.slot_count
         self._slot_ready_waited = [False] * self.slot_count
         self._slot_free_recorded = [False] * self.slot_count
@@ -45,7 +52,7 @@ def _validate_slot(self, slot_idx):
         if len(self.cuda_buffers) < self.slot_count:
             raise RuntimeError(f"Event-slot weight offload requires {self.slot_count} device buffers")
 
-    def _load_block_to_buffer(self, target_buffer, block_idx, blocks, adapter_block_idx):
+    def _load_block_to_buffer(self, target_buffer, block_idx, blocks, adapter_block_idx, state_dict_transform):
         block_slab = getattr(self, "block_slabs", {}).get(block_idx)
         if block_slab is not None:
             copy_block_slab_(
@@ -67,9 +74,12 @@ def _load_block_to_buffer(self, target_buffer, block_idx, blocks, adapter_block_
             if blocks is None:
                 raise ValueError("blocks must be provided when CPU buffers have not been initialized")
             source = blocks[block_idx]
-        target_buffer.load_state_dict(source.state_dict(), block_idx, adapter_block_idx)
+        state_dict = source.state_dict()
+        if state_dict_transform is not None:
+            state_dict = state_dict_transform(block_idx, state_dict)
+        target_buffer.load_state_dict(state_dict, block_idx, adapter_block_idx)
 
-    def prefetch_to_slot(self, slot_idx, block_idx, blocks=None, adapter_block_idx=None):
+    def prefetch_to_slot(self, slot_idx, block_idx, blocks=None, adapter_block_idx=None, state_dict_transform=None):
         """Enqueue one block copy into a fixed staging slot."""
         self._validate_slot(slot_idx)
         if self._slot_pending[slot_idx]:
@@ -83,6 +93,7 @@ def prefetch_to_slot(self, slot_idx, block_idx, blocks=None, adapter_block_idx=N
                 block_idx,
                 blocks,
                 adapter_block_idx,
+                state_dict_transform,
             )
             self._ready_events[slot_idx].record(self.cuda_load_stream)
 

lightx2v/common/offload/manager.py

@@ -12,10 +12,14 @@
 
 
 class WeightAsyncStreamManager(object):
+    uses_events = False
+
     def __init__(self, offload_granularity):
         self.offload_granularity = offload_granularity
         self.init_stream = torch_device_module.Stream(priority=0)
         self.need_init_first_buffer = True
+        self.loaded_state_dict_transform = None
+        self.loaded_first_adapter_block_index = None
         self.lazy_load = False
         torch_version = parse(torch.__version__.split("+")[0])
         # Legacy name: this is the active device backend's weight-loading stream, not a CUDA-only stream.
@@ -60,27 +64,34 @@ def _sync(self):
         else:
             self.init_stream.synchronize()
 
-    def init_first_buffer(self, blocks, adapter_block_idx=None):
+    def init_first_buffer(self, blocks, adapter_block_idx=None, block_idx=0, state_dict_transform=None):
         with torch_device_module.stream(self.init_stream):
             if hasattr(self, "cpu_buffers"):
                 if self.offload_granularity == "block":
-                    self.cuda_buffers[0].load_state_dict(self.cpu_buffers[0].state_dict(), 0, adapter_block_idx)
+                    state_dict = self.cpu_buffers[0].state_dict()
                 else:
-                    self.cuda_buffers[0].load_state_dict(self.cpu_buffers[0][0].state_dict(), 0, adapter_block_idx)
+                    state_dict = self.cpu_buffers[0][0].state_dict()
             else:
                 if self.offload_granularity == "block":
-                    self.cuda_buffers[0].load_state_dict(blocks[0].state_dict(), 0, adapter_block_idx)
+                    state_dict = blocks[block_idx].state_dict()
                 else:
-                    self.cuda_buffers[0].load_state_dict(blocks[0].compute_phases[0].state_dict(), 0, adapter_block_idx)
+                    state_dict = blocks[block_idx].compute_phases[0].state_dict()
+            if state_dict_transform is not None:
+                state_dict = state_dict_transform(block_idx, state_dict)
+            self.cuda_buffers[0].load_state_dict(state_dict, block_idx, adapter_block_idx)
         self._sync()
         self.need_init_first_buffer = False
 
     def prefetch_weights(self, block_idx, blocks, adapter_block_idx=None):
+        self.prefetch_weights_to_buffer(1, block_idx, blocks, adapter_block_idx)
+
+    def prefetch_weights_to_buffer(self, buffer_idx, block_idx, blocks, adapter_block_idx=None, state_dict_transform=None):
         with torch_device_module.stream(self.cuda_load_stream):
-            if hasattr(self, "cpu_buffers"):
-                self.cuda_buffers[1].load_state_dict(self.cpu_buffers[0].state_dict(), block_idx, adapter_block_idx)
-            else:
-                self.cuda_buffers[1].load_state_dict(blocks[block_idx].state_dict(), block_idx, adapter_block_idx)
+            source = self.cpu_buffers[0] if hasattr(self, "cpu_buffers") else blocks[block_idx]
+            state_dict = source.state_dict()
+            if state_dict_transform is not None:
+                state_dict = state_dict_transform(block_idx, state_dict)
+            self.cuda_buffers[buffer_idx].load_state_dict(state_dict, block_idx, adapter_block_idx)
 
     def prefetch_phase(self, block_idx, phase_idx, blocks, adapter_block_idx=None):
         with torch_device_module.stream(self.cuda_load_stream):
@@ -89,17 +100,26 @@ def prefetch_phase(self, block_idx, phase_idx, blocks, adapter_block_idx=None):
             else:
                 self.cuda_buffers[phase_idx].load_state_dict(blocks[block_idx].compute_phases[phase_idx].state_dict(), block_idx, adapter_block_idx)
 
-    def swap_blocks(self):
+    def synchronize_block_streams(self):
         if AI_DEVICE == "xpu":
             torch_device_module.synchronize()
         else:
             self.cuda_load_stream.synchronize()
             self.compute_stream.synchronize()
+
+    def swap_blocks(self):
+        self.synchronize_block_streams()
         self.cuda_buffers[0], self.cuda_buffers[1] = (
             self.cuda_buffers[1],
             self.cuda_buffers[0],
         )
 
+    def wait_for_block_compute(self):
+        if AI_DEVICE == "xpu":
+            torch_device_module.synchronize()
+        else:
+            self.compute_stream.synchronize()
+
     def swap_phases(self):
         if AI_DEVICE == "xpu":
             torch_device_module.synchronize()