[bugfix] Support lora request for non diffusion model (#5374)

Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Co-authored-by: Samit <285365963@qq.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/entrypoints/test_async_omni.py

@@ -3,6 +3,7 @@ import re
 from types import SimpleNamespace
 
 import pytest
+from vllm.lora.request import LoRARequest
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 
 from tests.helpers.mark import hardware_test
@@ -21,10 +22,13 @@ async def _noop(*args, **kw):
     pass
 
 
-def get_fake_add_request(submitted_request_ids):
+def get_fake_add_request(submitted_request_ids, submitted_lora_requests=None):
     async def fake_add_request_async(*, request_id, prompt, sampling_params_list, final_stage_id, **kwargs):
+        lora_request = kwargs.get("lora_request")
         del prompt, sampling_params_list, final_stage_id, kwargs
         submitted_request_ids.append(request_id)
+        if submitted_lora_requests is not None:
+            submitted_lora_requests.append(lora_request)
 
     return fake_add_request_async
 
@@ -92,6 +96,37 @@ def test_generate_submits_randomized_id_to_engine():
     asyncio.run(run())
 
 
+@pytest.mark.cpu
+def test_generate_forwards_lora_request_to_engine():
+    """Ensure the lora_request passed to generate() reaches add_request_async.
+
+    Regression test for https://github.com/vllm-project/vllm-omni/issues/5369:
+    AsyncOmni.generate() accepted lora_request but silently dropped it when
+    submitting the request, so generation always used the base model.
+    """
+
+    async def run():
+        submitted_ids = []
+        submitted_loras = []
+        omni = get_async_omni_instance(fake_add_request=get_fake_add_request(submitted_ids, submitted_loras))
+
+        lora = LoRARequest(lora_name="test", lora_int_id=1, lora_path="/tmp/fake")
+        async for _ in omni.generate(
+            prompt={"prompt": "test"},
+            request_id="lora-req",
+            sampling_params_list=[SimpleNamespace()],
+            output_modalities=["text"],
+            lora_request=lora,
+        ):
+            pass
+
+        assert len(submitted_ids) == 1
+        assert len(submitted_loras) == 1
+        assert submitted_loras[0] is lora
+
+    asyncio.run(run())
+
+
 @pytest.mark.cpu
 @pytest.mark.parametrize(
     "req_ids,cancel_prefix,expected_cancel_count",

vllm_omni/engine/async_omni_engine.py

@@ -1405,6 +1405,7 @@ class AsyncOmniEngine:
         final_stage_id: int = 0,
         final_output_stage_ids: Sequence[int] | None = None,
         arrival_time: float | None = None,
+        lora_request: Any = None,
         *,
         resumable: bool = True,
     ) -> None:
@@ -1417,6 +1418,7 @@ class AsyncOmniEngine:
             final_stage_id=final_stage_id,
             final_output_stage_ids=final_output_stage_ids,
             arrival_time=arrival_time,
+            lora_request=lora_request,
             resumable=resumable,
             message_type="streaming_update",
         )
@@ -1431,6 +1433,7 @@ class AsyncOmniEngine:
         final_stage_id: int = 0,
         final_output_stage_ids: Sequence[int] | None = None,
         arrival_time: float | None = None,
+        lora_request: Any = None,
         *,
         resumable: bool = True,
     ) -> None:
@@ -1443,6 +1446,7 @@ class AsyncOmniEngine:
             final_stage_id=final_stage_id,
             final_output_stage_ids=final_output_stage_ids,
             arrival_time=arrival_time,
+            lora_request=lora_request,
             resumable=resumable,
         )
 

vllm_omni/entrypoints/async_omni.py

@@ -599,6 +599,7 @@ class AsyncOmni(EngineClient, OmniBase):
                     final_stage_id=final_stage_id_for_e2e,
                     final_output_stage_ids=final_output_stage_ids,
                     arrival_time=wall_start_ts,
+                    lora_request=lora_request,
                 )
             else:
                 await self.engine.add_request_async(
@@ -608,6 +609,7 @@ class AsyncOmni(EngineClient, OmniBase):
                     final_stage_id=final_stage_id_for_e2e,
                     final_output_stage_ids=final_output_stage_ids,
                     arrival_time=wall_start_ts,
+                    lora_request=lora_request,
                 )
             submit_ts = time.time()
             req_state.metrics.stage_first_ts[0] = submit_ts
@@ -659,6 +661,7 @@ class AsyncOmni(EngineClient, OmniBase):
         final_stage_id: int,
         final_output_stage_ids: Sequence[int],
         arrival_time: float,
+        lora_request: Any = None,
     ) -> asyncio.Task:
         """Submit a streaming input generator as incremental stage-0 updates."""
         if not sampling_params_list:
@@ -696,6 +699,7 @@ class AsyncOmni(EngineClient, OmniBase):
                             final_stage_id=final_stage_id,
                             final_output_stage_ids=final_output_stage_ids,
                             arrival_time=arrival_time,
+                            lora_request=lora_request,
                             resumable=True,
                         )
                         has_submitted_first_chunk = True
@@ -708,6 +712,7 @@ class AsyncOmni(EngineClient, OmniBase):
                             final_stage_id=final_stage_id,
                             final_output_stage_ids=final_output_stage_ids,
                             arrival_time=arrival_time,
+                            lora_request=lora_request,
                             resumable=True,
                         )
             except (asyncio.CancelledError, GeneratorExit):
@@ -739,6 +744,7 @@ class AsyncOmni(EngineClient, OmniBase):
                             final_stage_id=final_stage_id,
                             final_output_stage_ids=final_output_stage_ids,
                             arrival_time=arrival_time,
+                            lora_request=lora_request,
                             resumable=False,
                         )
                     else:
@@ -750,6 +756,7 @@ class AsyncOmni(EngineClient, OmniBase):
                             final_stage_id=final_stage_id,
                             final_output_stage_ids=final_output_stage_ids,
                             arrival_time=arrival_time,
+                            lora_request=lora_request,
                             resumable=False,
                         )
 

vllm_omni/entrypoints/openai/serving_chat.py

@@ -786,6 +786,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     sampling_params_list=sampling_params_list,
                     output_modalities=output_modalities,
                     arrival_time=request_timestamp,
+                    lora_request=lora_request,
                 )
 
                 generators.append(generator)