From c588208cc6132b08f5066420468e047ca581fbdc Mon Sep 17 00:00:00 2001 From: Nguyen Kha Nhat Long Date: Mon, 10 Aug 2026 11:27:51 +0800 Subject: [PATCH] [bugfix] Support lora request for non diffusion model (#5374) Signed-off-by: knlnguyen1802 Co-authored-by: Samit <285365963@qq.com> Co-authored-by: Gao Han Co-authored-by: Hongsheng Liu --- tests/entrypoints/test_async_omni.py | 37 +++++++++++++++++++- vllm_omni/engine/async_omni_engine.py | 4 +++ vllm_omni/entrypoints/async_omni.py | 7 ++++ vllm_omni/entrypoints/openai/serving_chat.py | 1 + 4 files changed, 48 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index 7448cc5c0..6120d5ea4 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/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", diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 5a8642d59..fc80cfa93 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/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, ) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 1437cc680..f0af77699 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/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, ) diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 305f970f8..2e105d9f6 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/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)