[Feature] End-to-end LoRA support for BAGEL (#2494)

Signed-off-by: Zhengyuan Su <su.zhengyuan@u.nus.edu>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

tests/e2e/offline_inference/test_bagel_lora.py

@@ -0,0 +1,198 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+End-to-end test for BAGEL LoRA support (Stage 1 / DiT).
+
+Validates that LoRA adapters are correctly loaded, applied with controllable
+scale, and cleanly deactivated.  Uses a synthetic rank-1 adapter targeting the
+first decoder layer's QKV projection.
+
+Assertions:
+  (a) LoRA at scale=1.0 visibly changes the output  (diff > 0.5)
+  (b) scale=2.0 produces a larger delta than scale=1.0  (linearity)
+  (c) The delta is bounded  (diff < 80, not corrupted)
+  (d) Deactivating LoRA exactly restores the baseline  (diff == 0)
+"""
+
+import json
+import os
+
+from vllm_omni.inputs.data import OmniSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from safetensors.torch import save_file
+
+from tests.conftest import modify_stage_config
+from tests.utils import hardware_test
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.lora.request import LoRARequest
+from vllm_omni.lora.utils import stable_lora_int_id
+
+MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
+BAGEL_STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
+DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>"
+
+
+# ---------------------------------------------------------------------------
+# Helpers (reused from test_bagel_text2img.py patterns)
+# ---------------------------------------------------------------------------
+
+
+def _resolve_stage_config(config_path: str, run_level: str) -> str:
+    if run_level == "advanced_model":
+        return modify_stage_config(
+            config_path,
+            deletes={
+                "stage_args": {
+                    0: ["engine_args.load_format"],
+                    1: ["engine_args.load_format"],
+                }
+            },
+        )
+    return config_path
+
+
+def _configure_sampling_params(omni: Omni, num_inference_steps: int = 10) -> list[OmniSamplingParams]:
+    params_list = omni.default_sampling_params_list
+    if len(params_list) > 1:
+        params_list[1].num_inference_steps = num_inference_steps
+        params_list[1].extra_args = {
+            "cfg_text_scale": 4.0,
+            "cfg_img_scale": 1.5,
+        }
+    return params_list
+
+
+def _extract_generated_image(omni_outputs: list[OmniRequestOutput]) -> Image.Image | None:
+    for req_output in omni_outputs:
+        if req_output.images:
+            return req_output.images[0]
+    return None
+
+
+def _generate_bagel_image(omni: Omni) -> Image.Image:
+    params_list = _configure_sampling_params(omni)
+    params_list[1].lora_request = None
+    outputs = list(
+        omni.generate(
+            prompts=[{"prompt": DEFAULT_PROMPT, "modalities": ["image"]}],
+            sampling_params_list=params_list,
+        )
+    )
+    img = _extract_generated_image(outputs)
+    assert img is not None, "No image generated"
+    return img
+
+
+def _generate_bagel_image_with_lora(
+    omni: Omni,
+    lora_request: LoRARequest,
+    lora_scale: float = 1.0,
+) -> Image.Image:
+    params_list = _configure_sampling_params(omni)
+    params_list[1].lora_request = lora_request
+    params_list[1].lora_scale = lora_scale
+    outputs = list(
+        omni.generate(
+            prompts=[{"prompt": DEFAULT_PROMPT, "modalities": ["image"]}],
+            sampling_params_list=params_list,
+        )
+    )
+    img = _extract_generated_image(outputs)
+    assert img is not None, "No image generated with LoRA"
+    return img
+
+
+# BAGEL uses GQA: hidden_size=3584, 28 Q heads, 4 KV heads, head_dim=128
+# QKV packed dim = 28*128 + 4*128 + 4*128 = 3584 + 512 + 512 = 4608
+_LORA_DIM = 3584
+_LORA_QKV_DIM = 4608
+_LORA_MODULE = "bagel.language_model.model.layers.0.self_attn.qkv_proj"
+_LORA_RANK = 4
+
+
+def _make_file_lora_request(adapter_dir: Path) -> LoRARequest:
+    """Write synthetic adapter to disk and return a file-backed LoRARequest."""
+    adapter_dir.mkdir(parents=True, exist_ok=True)
+    gen = torch.Generator().manual_seed(42)
+    lora_a = torch.randn((_LORA_RANK, _LORA_DIM), dtype=torch.float32, generator=gen) * 0.1
+    lora_b = torch.randn((_LORA_QKV_DIM, _LORA_RANK), dtype=torch.float32, generator=gen) * 0.5
+    save_file(
+        {
+            f"base_model.model.{_LORA_MODULE}.lora_A.weight": lora_a,
+            f"base_model.model.{_LORA_MODULE}.lora_B.weight": lora_b,
+        },
+        str(adapter_dir / "adapter_model.safetensors"),
+    )
+    (adapter_dir / "adapter_config.json").write_text(
+        json.dumps({"r": _LORA_RANK, "lora_alpha": _LORA_RANK, "target_modules": [_LORA_MODULE]}),
+        encoding="utf-8",
+    )
+    lora_dir = str(adapter_dir)
+    return LoRARequest(lora_name="test_file", lora_int_id=stable_lora_int_id(lora_dir), lora_path=lora_dir)
+
+
+# ---------------------------------------------------------------------------
+# Test
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
+def test_bagel_lora_scale_and_deactivation(run_level, tmp_path):
+    """Validate LoRA effect, bounded perturbation, and clean deactivation."""
+    config_path = _resolve_stage_config(BAGEL_STAGE_CONFIG, run_level)
+    omni = Omni(model=MODEL, stage_configs_path=config_path, stage_init_timeout=300)
+    try:
+        lora_request = _make_file_lora_request(tmp_path / "bagel_lora")
+
+        # 1) Baseline (no LoRA)
+        baseline = _generate_bagel_image(omni)
+
+        # 2) LoRA with scale=1.0
+        img_1x = _generate_bagel_image_with_lora(omni, lora_request, lora_scale=1.0)
+
+        # 3) LoRA with scale=2.0
+        img_2x = _generate_bagel_image_with_lora(omni, lora_request, lora_scale=2.0)
+
+        # 4) No LoRA again (deactivation)
+        restored = _generate_bagel_image(omni)
+
+        baseline_arr = np.array(baseline, dtype=np.int16)
+        img_1x_arr = np.array(img_1x, dtype=np.int16)
+        img_2x_arr = np.array(img_2x, dtype=np.int16)
+        restored_arr = np.array(restored, dtype=np.int16)
+
+        diff_1x = np.abs(baseline_arr - img_1x_arr).mean()
+        diff_2x = np.abs(baseline_arr - img_2x_arr).mean()
+        diff_restored = np.abs(baseline_arr - restored_arr).mean()
+
+        # (a) Adapter has visible effect at both scales
+        assert diff_1x > 0.5, f"LoRA scale=1.0 had no visible effect: diff={diff_1x}"
+        assert diff_2x > 0.5, f"LoRA scale=2.0 had no visible effect: diff={diff_2x}"
+
+        # (b) Different scales produce different outputs
+        assert not np.isclose(diff_1x, diff_2x, atol=1.0), (
+            f"LoRA scale has no effect: diff_1x={diff_1x:.2f}, diff_2x={diff_2x:.2f}"
+        )
+
+        # (c) Output is not corrupted
+        assert diff_1x < 80, f"LoRA output looks corrupted: diff_1x={diff_1x}"
+        assert diff_2x < 80, f"LoRA output looks corrupted: diff_2x={diff_2x}"
+
+        # (d) Deactivation fully restores base model
+        assert diff_restored == 0.0, f"Base model not restored after LoRA deactivation: diff={diff_restored}"
+    finally:
+        omni.close()

tests/engine/test_cross_stage_lora.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for cross-stage LoRA routing in the orchestrator."""
+
+from __future__ import annotations
+
+import pytest
+from vllm.lora.request import LoRARequest
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.engine.orchestrator import build_engine_core_request_from_tokens
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestBuildEngineCoreRequestLoRA:
+    """Verify build_engine_core_request_from_tokens passes LoRA from params."""
+
+    def test_lora_extracted_from_diffusion_params(self):
+        lr = LoRARequest(lora_name="test", lora_int_id=1, lora_path="/tmp/fake")
+        params = OmniDiffusionSamplingParams(lora_request=lr)
+
+        # OmniDiffusionSamplingParams is not a SamplingParams, so
+        # build_engine_core_request_from_tokens takes the pooling path.
+        # We only care that lora_request is extracted via getattr.
+        request = build_engine_core_request_from_tokens(
+            request_id="req-1",
+            prompt={"prompt_token_ids": [1, 2, 3]},
+            params=params,
+            model_config=None,
+        )
+        assert request.lora_request is lr
+
+    def test_no_lora_on_sampling_params(self):
+        params = SamplingParams(max_tokens=10)
+
+        request = build_engine_core_request_from_tokens(
+            request_id="req-2",
+            prompt={"prompt_token_ids": [1, 2, 3]},
+            params=params,
+            model_config=None,
+        )
+        assert request.lora_request is None

vllm_omni/engine/orchestrator.py

@@ -79,7 +79,7 @@ def build_engine_core_request_from_tokens(
         sampling_params=sampling_params,
         pooling_params=pooling_params,
         arrival_time=arrival_time,
-        lora_request=None,
+        lora_request=getattr(params, "lora_request", None),
         cache_salt=None,
         data_parallel_rank=None,
         prompt_embeds=prompt_embeds,