[Feature] Support MiniMax-H3 Turbo LoRA with DLO (#6550)

Signed-off-by: mglyn <1203789601@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

recipes/MiniMaxAI/MiniMax-H3.md

@@ -839,7 +839,7 @@ hf download lightx2v/Minimax-h3-Turbo "${TURBO_FILE}" --local-dir "${TURBO_DIR}"
 export TURBO_LORA="${TURBO_DIR}/${TURBO_FILE}"
 ```
 
-Start from a non-offloaded FL2VA server command and add
+Start from a non-offloaded or DLO FL2VA server command and add
 `--task-type fl2va --lora-backend peft --lora-path "${TURBO_LORA}"`.
 `--lora-path` preloads the adapter; each request still activates it and uses
 the published sampling settings:
@@ -853,9 +853,10 @@ the published sampling settings:
 
 For FL2VA, change `task` and add `input_reference` as shown above. The 8-step,
 ComfyUI, Ref2VA, and v1.1 artifacts are not supported. This integration is
-dynamic-only and does not support prefusion, DLO, or LoRA composition.
-It also rejects model-level, layerwise, and distributed layerwise offload;
-the legacy dynamic LoRA tensors do not participate in those weight lifecycles.
+dynamic-only and does not support prefusion or LoRA composition. DLO is
+supported by keeping the request-switchable LoRA A/B buffers resident on the
+accelerator while DLO streams only the base blocks; budget for this additional
+fixed HBM usage. Model-level and standard layerwise offload remain unsupported.
 The five requested sigma points produce the four denoiser evaluations expected
 by the Turbo artifact.
 

tests/diffusion/lora/test_base_linear.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -89,6 +89,29 @@ def test_diffusion_base_linear_apply_multi_slice():
     assert torch.allclose(out, expected)
 
 
+def test_diffusion_lora_buffer_device_persists_across_reallocation(monkeypatch):
+    from vllm.lora.layers.base_linear import BaseLinearLayerWithLoRA
+
+    layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
+    torch.nn.Module.__init__(layer)
+    layer.base_layer = torch.nn.Module()
+    layer.n_slices = 1
+    layer.lora_a_stacked = (torch.zeros(2),)
+    layer.lora_b_stacked = (torch.zeros(3),)
+
+    layer._set_diffusion_lora_buffer_device(torch.device("meta"))
+    assert all(tensor.is_meta for tensor in (*layer.lora_a_stacked, *layer.lora_b_stacked))
+
+    def recreate_on_cpu(self, max_loras, lora_config, model_config):
+        self.lora_a_stacked = (torch.zeros(4),)
+        self.lora_b_stacked = (torch.zeros(5),)
+
+    monkeypatch.setattr(BaseLinearLayerWithLoRA, "create_lora_weights", recreate_on_cpu)
+    layer.create_lora_weights(max_loras=1, lora_config=object())
+
+    assert all(tensor.is_meta for tensor in (*layer.lora_a_stacked, *layer.lora_b_stacked))
+
+
 def test_diffusion_base_linear_reset_lora_disables_fast_path(monkeypatch):
     # Verify that after reset_lora(), apply() skips LoRA matmuls even if the
     # LoRA tensors are still allocated and non-empty.

tests/diffusion/lora/test_lora_manager.py

@@ -3,6 +3,8 @@
 
 from __future__ import annotations
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 from vllm.lora.lora_weights import LoRALayerWeights
@@ -143,6 +145,37 @@ def test_lora_manager_replace_layers_does_not_rewrap_base_layer(monkeypatch):
     assert replace_calls == ["foo"]
 
 
+def test_lora_manager_keeps_dlo_lora_buffers_on_compute_device(monkeypatch):
+    import vllm_omni.diffusion.lora.manager as manager_mod
+
+    class _DLOLoRALayer(DummyBaseLayerWithLoRA):
+        lora_buffer_device = None
+
+        def _set_diffusion_lora_buffer_device(self, device):
+            self.lora_buffer_device = device
+
+    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DLOLoRALayer)
+    monkeypatch.setattr(
+        manager_mod,
+        "from_layer_diffusion",
+        lambda *, layer, **_kwargs: _DLOLoRALayer(layer),
+    )
+
+    pipeline = torch.nn.Module()
+    pipeline.od_config = SimpleNamespace(enable_distributed_layerwise_offload=True)
+    pipeline.transformer = torch.nn.Module()
+    pipeline.transformer.foo = _FakeLinearBase()
+    manager = DiffusionLoRAManager(
+        pipeline=pipeline,
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+    )
+
+    manager._replace_layers_with_lora(type("_PH", (), {"r": 1})())
+
+    assert pipeline.transformer.foo.lora_buffer_device == torch.device("cpu")
+
+
 def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch):
     import vllm_omni.diffusion.lora.manager as manager_mod
 

tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py

@@ -197,7 +197,6 @@ def test_h3_turbo_rejects_wrong_alpha_and_ref2va(tmp_path):
     [
         "model-level CPU offload (--enable-cpu-offload)",
         "layerwise offload (--enable-layerwise-offload)",
-        "distributed layerwise offload (--enable-distributed-layerwise-offload)",
     ],
 )
 def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode):
@@ -214,6 +213,38 @@ def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode):
         )
 
 
+def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import (
+        pipeline_minimax_h3 as pipeline_module,
+    )
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.partition = "fl2va"
+    pipeline._turbo_lora_adapter_ids = set()
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=True,
+    )
+    captured = {}
+
+    def load_turbo(**kwargs):
+        captured.update(kwargs)
+        return object(), object()
+
+    monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", load_turbo)
+    loaded = pipeline._load_diffusion_lora_adapter(
+        lora_request=_request("turbo"),
+        lora_path="turbo",
+        dtype=torch.bfloat16,
+    )
+
+    assert loaded is not None
+    assert captured["unsupported_offload_mode"] is None
+
+
 def test_h3_turbo_accepts_only_the_declared_v1_artifact(tmp_path):
     unsupported = tmp_path / "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors"
     _write_tiny_turbo(unsupported)

vllm_omni/diffusion/lora/layers/base_linear.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -22,6 +22,19 @@ class DiffusionBaseLinearLayerWithLoRA(BaseLinearLayerWithLoRA):
     is inherited from vLLM's BaseLinearLayerWithLoRA.
     """
 
+    lora_a_stacked: tuple[torch.Tensor, ...]
+    lora_b_stacked: tuple[torch.Tensor, ...]
+
+    def _move_lora_buffers(self, device: torch.device) -> None:
+        self.lora_a_stacked = tuple(tensor.to(device=device) for tensor in self.lora_a_stacked)
+        self.lora_b_stacked = tuple(tensor.to(device=device) for tensor in self.lora_b_stacked)
+
+    def _set_diffusion_lora_buffer_device(self, device: torch.device) -> None:
+        """Keep dynamic LoRA sidecars resident on the compute device."""
+
+        self._diffusion_lora_buffer_device = device
+        self._move_lora_buffers(device)
+
     def create_lora_weights(
         self,
         max_loras: int,
@@ -29,6 +42,9 @@ class DiffusionBaseLinearLayerWithLoRA(BaseLinearLayerWithLoRA):
         model_config=None,
     ) -> None:
         super().create_lora_weights(max_loras, lora_config, model_config)
+        buffer_device = getattr(self, "_diffusion_lora_buffer_device", None)
+        if buffer_device is not None:
+            self._move_lora_buffers(buffer_device)
         # Keep a direct reference for attribute forwarding: `base_layer` is a
         # registered submodule (stored under `_modules`), so direct access via
         # `object.__getattribute__` will not find it. We stash a ref in

vllm_omni/diffusion/lora/manager.py

@@ -68,6 +68,12 @@ class DiffusionLoRAManager:
         self.pipeline = pipeline
         self.device = device
         self.dtype = dtype
+        od_config = getattr(pipeline, "od_config", None)
+        # DLO owns the base-weight lifecycle. Keep request-switchable LoRA
+        # sidecars resident instead of rebuilding DLO host shards per request.
+        self._resident_lora_device = (
+            device if getattr(od_config, "enable_distributed_layerwise_offload", False) is True else None
+        )
 
         # Cache supported/expected module suffixes once, before any layer
         # replacement happens. After LoRA layers are injected, the original
@@ -450,6 +456,13 @@ class DiffusionLoRAManager:
                 )
 
                 if lora_layer is not module and isinstance(lora_layer, BaseLayerWithLoRA):
+                    if self._resident_lora_device is not None:
+                        set_buffer_device = getattr(lora_layer, "_set_diffusion_lora_buffer_device", None)
+                        if not callable(set_buffer_device):
+                            raise RuntimeError(
+                                f"{type(lora_layer).__name__} cannot keep dynamic LoRA buffers resident for DLO"
+                            )
+                        set_buffer_device(self._resident_lora_device)
                     replace_submodule(component, module_name, lora_layer)
                     self._lora_modules[full_module_name] = lora_layer
                     logger.debug("Replaced layer: %s -> %s", full_module_name, type(lora_layer).__name__)

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -724,8 +724,6 @@ class MiniMaxH3Pipeline(
             offload_modes.append("model-level CPU offload (--enable-cpu-offload)")
         if getattr(od_config, "enable_layerwise_offload", False):
             offload_modes.append("layerwise offload (--enable-layerwise-offload)")
-        if getattr(od_config, "enable_distributed_layerwise_offload", False):
-            offload_modes.append("distributed layerwise offload (--enable-distributed-layerwise-offload)")
         loaded = load_minimax_h3_turbo_lora(
             partition=self.partition,
             lora_request=lora_request,