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

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

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)