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

[Core] Refactor omni prefix cache into Manager/Controller (#6654)

Browse Source 
Signed-off-by: ZhengWG <zwg0606@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Zheng Wengangamy-why-3459Hongsheng Liu
 authored and GitHub committed 2026-09-19 21:23:59 +08:00
1 parent ea84055507
commit 470ec60848
25 files changed
+5710 -2149
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/worker/test_gpu_ar_model_runner.py
+30 -11
View File
@@ -18,6 +18,7 @@ from vllm.sampling_params import SamplingParams
from vllm.v1.worker import gpu_input_batch
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
from vllm_omni.core.prefix_cache import ModelCachePolicy
from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest
@@ -37,6 +38,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _make_runner(engine_output_type: str | None, downstream_req_ids: set[str]) -> GPUARModelRunner:
runner = object.__new__(GPUARModelRunner)
runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
runner._pooler_payload_include_hidden_flag = True
runner.vllm_config = SimpleNamespace(
model_config=SimpleNamespace(engine_output_type=engine_output_type),
)
@@ -57,6 +60,8 @@ def _mtp_runner(*, async_scheduling: bool, buffers: dict[str, dict]) -> tuple[GP
return torch.tensor([[11, 12, 13]], dtype=torch.long)
runner = object.__new__(GPUARModelRunner)
runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
runner._pooler_payload_include_hidden_flag = True
runner.use_async_scheduling = async_scheduling
runner.model = SimpleNamespace(post_sample_talker_mtp=post_sample_talker_mtp)
runner.requests = {req_id: SimpleNamespace() for req_id in buffers}
@@ -143,6 +148,8 @@ def test_post_sample_talker_mtp_uses_gpu_token_with_async_scheduling() -> None:
def test_post_sample_talker_mtp_skips_shape_validation_without_duplex_rows() -> None:
runner = object.__new__(GPUARModelRunner)
runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
runner._pooler_payload_include_hidden_flag = True
runner.use_async_scheduling = False
runner.model = SimpleNamespace(
post_sample_talker_mtp=lambda **_: pytest.fail("hook must not run"),
@@ -178,6 +185,8 @@ def test_post_sample_talker_mtp_rejects_invalid_selected_token_shape(
error_match: str,
) -> None:
runner = object.__new__(GPUARModelRunner)
runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
runner._pooler_payload_include_hidden_flag = True
runner.use_async_scheduling = False
runner.model = SimpleNamespace(
post_sample_talker_mtp=lambda **_: pytest.fail("hook must not run"),
@@ -265,6 +274,8 @@ def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatc
received = []
runner = object.__new__(GPUARModelRunner)
runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
runner._pooler_payload_include_hidden_flag = True
runner.input_batch = input_batch
runner.model = SimpleNamespace(
prefer_model_sampler=True,
@@ -323,6 +334,8 @@ def test_sparse_mm_req_ids_requires_sparse_audio_marker():
def test_runner_assisted_full_attention_metadata_request_is_opt_in():
runner = object.__new__(GPUARModelRunner)
runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
runner._pooler_payload_include_hidden_flag = True
runner.model = object()
runner.scheduler_config = SimpleNamespace(max_num_seqs=16)
@@ -377,6 +390,8 @@ def test_runner_assisted_full_attention_metadata_request_and_context_hooks():
calls.append(("context", {"enabled": enabled, "num_reqs": num_reqs}))
runner = object.__new__(GPUARModelRunner)
runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
runner._pooler_payload_include_hidden_flag = True
runner.model = Model()
runner.scheduler_config = SimpleNamespace(max_num_seqs=8)
@@ -466,6 +481,8 @@ def test_omni_async_gpu_model_runner_output_reraises_background_exception():
def _make_async_output_runner(engine_output_type: str = "audio"):
runner = object.__new__(GPUARModelRunner)
runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
runner._pooler_payload_include_hidden_flag = True
model_config = SimpleNamespace(
engine_output_type=engine_output_type,
async_chunk=True,
@@ -642,10 +659,11 @@ def test_async_omni_output_guard_requires_safe_conditions():
assert GPUARModelRunner._should_use_async_omni_output(runner)
# Prefix cache no longer forces the sync output path; leftover mm is
# snapshotted at save so the builder can materialize a step late.
runner.omni_prefix_cache = object()
assert not GPUARModelRunner._should_use_async_omni_output(runner)
assert GPUARModelRunner._should_use_async_omni_output(runner)
runner.omni_prefix_cache = None
runner.model.has_postprocess = True
assert not GPUARModelRunner._should_use_async_omni_output(runner)
@@ -656,6 +674,7 @@ def test_async_omni_output_guard_requires_safe_conditions():
def test_build_omni_output_skips_hidden_when_model_opts_out(monkeypatch):
runner = _make_async_output_runner(engine_output_type="latent")
runner.model.omni_pooler_payload_include_hidden = False
runner._pooler_payload_include_hidden_flag = False
monkeypatch.setattr(
GPUARModelRunner,
@@ -699,6 +718,7 @@ def test_build_omni_output_skips_hidden_when_model_opts_out(monkeypatch):
def test_build_omni_output_splits_mm_by_scheduled_tokens_when_hidden_is_tail_only(monkeypatch):
runner = _make_async_output_runner(engine_output_type="latent")
runner.model.omni_pooler_payload_include_hidden = False
runner._pooler_payload_include_hidden_flag = False
runner._async_chunk = False
runner.requests = {"r1": object(), "r2": object(), "r3": object()}
@@ -745,6 +765,7 @@ def test_build_omni_output_splits_mm_by_hidden_len_when_scheduled_is_padded(monk
"""Thinker mm rows align to hidden_states.shape[0], not padded scheduled count."""
runner = _make_async_output_runner(engine_output_type="latent")
runner.model.omni_pooler_payload_include_hidden = False
runner._pooler_payload_include_hidden_flag = False
runner._async_chunk = False
runner.requests = {"r1": object(), "r2": object(), "r3": object()}
@@ -790,6 +811,7 @@ def test_build_omni_output_splits_mm_by_hidden_len_when_scheduled_is_padded(monk
def test_async_snapshot_payload_omits_hidden_when_model_opts_out():
runner = _make_async_output_runner()
runner.model.omni_pooler_payload_include_hidden = False
runner._pooler_payload_include_hidden_flag = False
payload = GPUARModelRunner._build_omni_async_snapshot_payload(
runner,
@@ -819,6 +841,8 @@ def test_runner_assisted_full_attention_metadata_refresh_pads_buffers():
self.commits.append(num_reqs_padded)
runner = object.__new__(GPUARModelRunner)
runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
runner._pooler_payload_include_hidden_flag = True
block_table = BlockTable()
runner.input_batch = SimpleNamespace(
num_computed_tokens_cpu=torch.tensor([10, 20, 99, 99], dtype=torch.int32),
@@ -846,6 +870,8 @@ def test_runner_assisted_full_attention_metadata_refresh_pads_buffers():
@pytest.mark.parametrize("query_start_loc_attr", ["method", "tensor_attr"])
def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monkeypatch, query_start_loc_attr):
runner = object.__new__(GPUARModelRunner)
runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
runner._pooler_payload_include_hidden_flag = True
runner.execute_model_state = ExecuteModelState(
SimpleNamespace(
total_num_scheduled_tokens=3,
@@ -908,12 +934,11 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk
)
monkeypatch.setattr(GPUARModelRunner, "eplb_step", lambda self: None)
monkeypatch.setattr(GPUARModelRunner, "_resolve_pooler_payload_req_ids", lambda self, req_ids: ("audio", req_ids))
monkeypatch.setattr(GPUARModelRunner, "_deferred_prefix_cache_mm_keys", lambda self: set())
monkeypatch.setattr(GPUARModelRunner, "_model_needs_full_prefix_hidden_states", lambda self: False)
monkeypatch.setattr(
GPUARModelRunner,
"_maybe_get_combined_prefix_cache_tensors",
lambda *args, **kwargs: (None, None),
"_prepare_prefix_cache_pooler_payload_sources",
lambda self, **kwargs: (kwargs["staged_hidden_states_cpu"], None, None),
)
monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
@@ -946,8 +971,6 @@ def test_build_omni_output_falls_back_to_mm_cpu_without_prefix_merge(monkeypatch
lambda self, req_ids: ("latent", req_ids),
)
monkeypatch.setattr(GPUARModelRunner, "_model_needs_full_prefix_hidden_states", lambda self: False)
monkeypatch.setattr(GPUARModelRunner, "_deferred_prefix_cache_mm_keys", lambda self: {"codes.audio"})
monkeypatch.setattr(GPUARModelRunner, "_stage_deferred_prefix_cache_mm_outputs", lambda *args, **kwargs: None)
monkeypatch.setattr(
GPUARModelRunner,
"_prepare_prefix_cache_pooler_payload_sources",
@@ -1095,8 +1118,6 @@ def test_build_omni_output_filters_multimodal_by_partial_downstream_batch(monkey
lambda self, req_ids: ("latent", ["r1", "r3"]),
)
monkeypatch.setattr(GPUARModelRunner, "_model_needs_full_prefix_hidden_states", lambda self: False)
monkeypatch.setattr(GPUARModelRunner, "_deferred_prefix_cache_mm_keys", lambda self: {"codes.audio"})
monkeypatch.setattr(GPUARModelRunner, "_stage_deferred_prefix_cache_mm_outputs", lambda *args, **kwargs: None)
monkeypatch.setattr(
GPUARModelRunner,
"_prepare_prefix_cache_pooler_payload_sources",
@@ -1172,8 +1193,6 @@ def test_build_omni_output_uses_combined_prefix_cache_mm_payload_for_partial_dow
lambda self, req_ids: ("latent", ["r1", "r3"]),
)
monkeypatch.setattr(GPUARModelRunner, "_model_needs_full_prefix_hidden_states", lambda self: False)
monkeypatch.setattr(GPUARModelRunner, "_deferred_prefix_cache_mm_keys", lambda self: {"codes.audio"})
monkeypatch.setattr(GPUARModelRunner, "_stage_deferred_prefix_cache_mm_outputs", lambda *args, **kwargs: None)
monkeypatch.setattr(
GPUARModelRunner,
"_prepare_prefix_cache_pooler_payload_sources",