anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[FEAT] support multi-stage deployment (#2396)

Browse Source 
Signed-off-by: ZhengWG <zwg0606@gmail.com>
Signed-off-by: Zheng Wengang <zwg0606@gmail.com>
Signed-off-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Signed-off-by: yinpe <11810305@mail.sustech.edu.cn>
Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
Co-authored-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Co-authored-by: yinpe <11810305@mail.sustech.edu.cn>
Co-authored-by: yinpeiqi <yinpeiqi809@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Chenguang Zheng <645327136@qq.com>
This commit is contained in:
Zheng WengangPeiqi YinyinpeyinpeiqiHongsheng LiuGao HanChenguang Zheng
 authored and GitHub committed 2026-05-06 19:33:18 +08:00
1 parent e969d2e99f
commit 1e5f288a91
55 files changed
+4684 -3407
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py
+4 -8
View File
@@ -44,16 +44,12 @@ def _transfer_manager(
def test_text2flow_supports_batched_source_outputs():
stage_list = [
SimpleNamespace(
engine_outputs=[
_source_output("req-0", [10, 11], [1, 2, 3], {"speech_token": torch.tensor([[1, 2]])}),
_source_output("req-1", [20, 21], [4, 5], {"speech_token": torch.tensor([[3, 4]])}),
]
)
source_outputs = [
_source_output("req-0", [10, 11], [1, 2, 3], {"speech_token": torch.tensor([[1, 2]])}),
_source_output("req-1", [20, 21], [4, 5], {"speech_token": torch.tensor([[3, 4]])}),
]
outputs = text2flow(stage_list=stage_list, engine_input_source=[0], prompt=None)
outputs = text2flow(source_outputs=source_outputs, prompt=None)
assert len(outputs) == 2
assert outputs[0]["prompt_token_ids"] == [1, 2, 3]
tests/model_executor/stage_input_processors/test_glm_image.py
+16 -30
View File
@@ -32,11 +32,6 @@ def _source_output(token_ids: list[int], mm_output: dict | None = None):
)
def _stage_with_outputs(outputs):
"""Create a stage list entry with engine_outputs."""
return SimpleNamespace(engine_outputs=outputs)
# =============================================================================
# Tests for _has_source_image
# =============================================================================
@@ -285,11 +280,11 @@ class TestAr2Diffusion:
"""Test basic text-to-image pipeline: AR -> Diffusion."""
# 1024x1024 t2i: small(256) + large(1024) + EOS
token_ids = list(range(256)) + list(range(1024)) + [16385]
stage_list = [_stage_with_outputs([_source_output(token_ids)])]
source_outputs = [_source_output(token_ids)]
prompt = {"prompt": "a cat", "mm_processor_kwargs": {"target_h": 1024, "target_w": 1024}}
result = ar2diffusion(stage_list, [0], prompt=[prompt])
result = ar2diffusion(source_outputs, prompt=[prompt])
assert len(result) == 1
assert result[0]["prompt"] == "a cat"
assert result[0]["height"] == 1024
@@ -300,7 +295,7 @@ class TestAr2Diffusion:
"""Test image-to-image with prior_token_image_ids from AR model."""
token_ids = list(range(1024)) + [16385]
mm_output = {"ids": {"prior_image": torch.tensor([1, 2, 3])}}
stage_list = [_stage_with_outputs([_source_output(token_ids, mm_output)])]
source_outputs = [_source_output(token_ids, mm_output)]
from PIL import Image
@@ -311,14 +306,14 @@ class TestAr2Diffusion:
"multi_modal_data": {"image": img},
}
result = ar2diffusion(stage_list, [0], prompt=[prompt])
result = ar2diffusion(source_outputs, prompt=[prompt])
assert len(result) == 1
assert result[0]["extra"]["prior_token_image_ids"] is not None
def test_i2i_detected_via_modalities(self):
"""Test i2i mode detected via modalities field."""
token_ids = list(range(1024)) + [16385]
stage_list = [_stage_with_outputs([_source_output(token_ids)])]
source_outputs = [_source_output(token_ids)]
prompt = {
"prompt": "edit this",
@@ -326,35 +321,26 @@ class TestAr2Diffusion:
"modalities": ["img2img"],
}
result = ar2diffusion(stage_list, [0], prompt=[prompt])
result = ar2diffusion(source_outputs, prompt=[prompt])
assert len(result) == 1
def test_empty_engine_input_source_raises(self):
with pytest.raises(ValueError, match="engine_input_source cannot be empty"):
ar2diffusion([], [], prompt={})
def test_invalid_stage_id_raises(self):
with pytest.raises(IndexError, match="Invalid stage_id"):
ar2diffusion([_stage_with_outputs(None)], [5], prompt={})
def test_no_outputs_raises(self):
with pytest.raises(RuntimeError, match="has no outputs yet"):
ar2diffusion([SimpleNamespace(engine_outputs=None)], [0], prompt={})
def test_empty_source_outputs_returns_empty_list(self):
assert ar2diffusion([], prompt={}) == []
def test_default_dimensions(self):
"""When no height/width in prompt, defaults to 1024x1024."""
token_ids = list(range(256)) + list(range(1024)) + [16385]
stage_list = [_stage_with_outputs([_source_output(token_ids)])]
source_outputs = [_source_output(token_ids)]
prompt = {"prompt": "test"}
result = ar2diffusion(stage_list, [0], prompt=[prompt])
result = ar2diffusion(source_outputs, prompt=[prompt])
assert result[0]["height"] == 1024
assert result[0]["width"] == 1024
def test_requires_multimodal_data_with_pil_image(self):
"""Test that pil_image is included when requires_multimodal_data=True."""
token_ids = list(range(256)) + list(range(1024)) + [16385]
stage_list = [_stage_with_outputs([_source_output(token_ids)])]
source_outputs = [_source_output(token_ids)]
from PIL import Image
@@ -364,13 +350,13 @@ class TestAr2Diffusion:
"multi_modal_data": {"image": img},
}
result = ar2diffusion(stage_list, [0], prompt=[prompt], requires_multimodal_data=True)
result = ar2diffusion(source_outputs, prompt=[prompt], requires_multimodal_data=True)
assert result[0]["pil_image"] is img
def test_extra_params_passed_through(self):
"""Test that seed, num_inference_steps, guidance_scale, negative_prompt are passed."""
token_ids = list(range(256)) + list(range(1024)) + [16385]
stage_list = [_stage_with_outputs([_source_output(token_ids)])]
source_outputs = [_source_output(token_ids)]
prompt = {
"prompt": "test",
@@ -380,7 +366,7 @@ class TestAr2Diffusion:
"negative_prompt": "blurry",
}
result = ar2diffusion(stage_list, [0], prompt=[prompt])
result = ar2diffusion(source_outputs, prompt=[prompt])
assert result[0]["seed"] == 42
assert result[0]["num_inference_steps"] == 50
assert result[0]["guidance_scale"] == 7.5
@@ -390,14 +376,14 @@ class TestAr2Diffusion:
"""Test processing multiple requests in a batch."""
tokens1 = list(range(256)) + list(range(1024)) + [16385]
tokens2 = list(range(256)) + list(range(1024)) + [16385]
stage_list = [_stage_with_outputs([_source_output(tokens1), _source_output(tokens2)])]
source_outputs = [_source_output(tokens1), _source_output(tokens2)]
prompts = [
{"prompt": "first", "mm_processor_kwargs": {"target_h": 1024, "target_w": 1024}},
{"prompt": "second", "mm_processor_kwargs": {"target_h": 512, "target_w": 512}},
]
result = ar2diffusion(stage_list, [0], prompt=prompts)
result = ar2diffusion(source_outputs, prompt=prompts)
assert len(result) == 2
assert result[0]["prompt"] == "first"
assert result[1]["prompt"] == "second"
tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py
+7 -8
View File
@@ -17,12 +17,11 @@ from vllm_omni.model_executor.stage_input_processors.mimo_audio import (
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _make_stage_list(codec_codes: torch.Tensor, request_id: str = "req-0"):
"""Build a minimal stage_list[0] with a single talker output carrying codec_codes."""
def _make_source_outputs(codec_codes: torch.Tensor, request_id: str = "req-0"):
output = SimpleNamespace(multimodal_output={"codes": {"audio": codec_codes}})
"""Build minimal source_outputs with one talker output carrying codec codes."""
talker_output = SimpleNamespace(outputs=[output], request_id=request_id)
stage0 = SimpleNamespace(engine_outputs=[talker_output])
return [stage0]
return [talker_output]
def test_llm2code2wav_truncates_when_flat_exceeds_max(caplog):
@@ -33,7 +32,7 @@ def test_llm2code2wav_truncates_when_flat_exceeds_max(caplog):
frames = (MAX_CODE2WAV_TOKENS // 36) + 100
codec_codes = torch.ones(frames, 1, 8, 4, dtype=torch.long)
stage_list = _make_stage_list(codec_codes, request_id="req-long")
source_outputs = _make_source_outputs(codec_codes, request_id="req-long")
# Attach caplog's handler directly to the module logger so the warning is
# captured regardless of propagation (vllm's logger configuration can
@@ -43,7 +42,7 @@ def test_llm2code2wav_truncates_when_flat_exceeds_max(caplog):
prev_level = target_logger.level
target_logger.setLevel(logging.WARNING)
try:
prompts = llm2code2wav(stage_list, engine_input_source=[0])
prompts = llm2code2wav(source_outputs)
finally:
target_logger.removeHandler(caplog.handler)
target_logger.setLevel(prev_level)
@@ -58,9 +57,9 @@ def test_llm2code2wav_truncates_when_flat_exceeds_max(caplog):
def test_llm2code2wav_short_sequence_unchanged():
"""Short codec sequences are returned without truncation."""
codec_codes = torch.ones(4, 1, 8, 4, dtype=torch.long)
stage_list = _make_stage_list(codec_codes, request_id="req-short")
source_outputs = _make_source_outputs(codec_codes, request_id="req-short")
prompts = llm2code2wav(stage_list, engine_input_source=[0])
prompts = llm2code2wav(source_outputs)
assert len(prompts) == 1
# 4 frames + 1 pad row, flattened col-major → well below the cap
tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py
+2 -2
View File
@@ -342,7 +342,7 @@ def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
)
prompts = talker2code2wav(stage_list=[stage], engine_input_source=[0])
prompts = talker2code2wav(stage.engine_outputs)
assert len(prompts) == 1
prompt = prompts[0]
@@ -388,7 +388,7 @@ def test_non_async_processor_filters_out_of_range_codec_values():
engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
)
prompts = talker2code2wav(stage_list=[stage], engine_input_source=[0])
prompts = talker2code2wav(stage.engine_outputs)
assert len(prompts) == 1
prompt = prompts[0]
tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py
+4 -16
View File
@@ -38,7 +38,7 @@ class TestGenerator2Tokenizer:
audio = torch.tensor([10, 20, 30, 40])
stage = _make_stage([_make_output([audio])])
result = generator2tokenizer([stage], engine_input_source=[0])
result = generator2tokenizer(stage.engine_outputs)
assert len(result) == 1
assert result[0]["prompt_token_ids"] == [10, 20, 30, 40]
@@ -49,27 +49,15 @@ class TestGenerator2Tokenizer:
audio = torch.tensor([[1, 2, 3], [4, 5, 6]]) # (2, 3)
stage = _make_stage([_make_output([audio])])
result = generator2tokenizer([stage], engine_input_source=[0])
result = generator2tokenizer(stage.engine_outputs)
assert result[0]["prompt_token_ids"] == [1, 2, 3, 4, 5, 6]
def test_empty_engine_input_source_raises(self):
"""Empty engine_input_source raises ValueError."""
stage = _make_stage([_make_output([torch.tensor([1])])])
with pytest.raises(ValueError, match="cannot be empty"):
generator2tokenizer([stage], engine_input_source=[])
def test_invalid_stage_id_raises(self):
"""Stage id beyond stage_list length raises IndexError."""
stage = _make_stage([_make_output([torch.tensor([1])])])
with pytest.raises(IndexError, match="Invalid stage_id"):
generator2tokenizer([stage], engine_input_source=[5])
def test_no_outputs_yet_raises(self):
"""Stage with engine_outputs=None raises RuntimeError."""
stage = _make_stage(engine_outputs=None)
with pytest.raises(RuntimeError, match="no outputs yet"):
generator2tokenizer([stage], engine_input_source=[0])
with pytest.raises(TypeError):
generator2tokenizer(stage.engine_outputs)
# ---- Helpers for generator2tokenizer_async_chunk (streaming) ----