1031 lines
40 KiB
Python
1031 lines
40 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from contextlib import contextmanager
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
|
|
from vllm.v1.cudagraph_dispatcher import CUDAGraphMode
|
|
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
|
|
|
|
from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
|
|
from vllm_omni.worker.gpu_model_runner import (
|
|
OmniGPUModelRunner,
|
|
_filter_mrope_kwargs_for_model,
|
|
)
|
|
from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def _runner_for_talker_graph_init(
|
|
*,
|
|
talker_mtp_graph_safe: bool | None,
|
|
has_separate_talker: bool = True,
|
|
model_stage: str = "talker",
|
|
) -> OmniGPUModelRunner:
|
|
runner = object.__new__(OmniGPUModelRunner)
|
|
runner.model = SimpleNamespace(
|
|
talker=object() if has_separate_talker else None,
|
|
talker_mtp=object(),
|
|
model_stage=model_stage,
|
|
)
|
|
if talker_mtp_graph_safe is not None:
|
|
runner.model.talker_mtp_graph_safe = talker_mtp_graph_safe
|
|
runner.model_config = SimpleNamespace(hf_text_config=SimpleNamespace(hidden_size=4))
|
|
runner.compilation_config = SimpleNamespace(
|
|
cudagraph_mode=CUDAGraphMode.FULL,
|
|
max_cudagraph_capture_size=1,
|
|
)
|
|
runner.vllm_config = object()
|
|
runner.max_num_reqs = 1
|
|
runner.dtype = torch.float32
|
|
runner._make_buffer = lambda *args, **kwargs: SimpleNamespace(args=args, kwargs=kwargs)
|
|
return runner
|
|
|
|
|
|
def test_talker_mtp_skips_graph_when_model_declares_unsafe(monkeypatch):
|
|
runner = _runner_for_talker_graph_init(talker_mtp_graph_safe=False)
|
|
talker_mtp = runner.model.talker_mtp
|
|
monkeypatch.setattr(
|
|
"vllm_omni.worker.gpu_model_runner.current_omni_platform.get_graph_wrapper_cls",
|
|
lambda: pytest.fail("graph wrapper must not be selected"),
|
|
)
|
|
|
|
OmniGPUModelRunner._init_talker_mtp(runner)
|
|
|
|
assert runner.talker_mtp is talker_mtp
|
|
|
|
|
|
@pytest.mark.parametrize("model_stage", ["thinker", "code2wav"])
|
|
def test_non_talker_stage_does_not_use_talker_mtp_graph(monkeypatch, model_stage: str):
|
|
runner = _runner_for_talker_graph_init(
|
|
talker_mtp_graph_safe=None,
|
|
has_separate_talker=False,
|
|
model_stage=model_stage,
|
|
)
|
|
talker_mtp = runner.model.talker_mtp
|
|
monkeypatch.setattr(
|
|
"vllm_omni.worker.gpu_model_runner.current_omni_platform.get_graph_wrapper_cls",
|
|
lambda: pytest.fail("graph wrapper must not be selected"),
|
|
)
|
|
|
|
OmniGPUModelRunner._init_talker_mtp(runner)
|
|
|
|
assert runner.talker_mtp is talker_mtp
|
|
|
|
|
|
@pytest.mark.parametrize("talker_mtp_graph_safe", [None, True])
|
|
def test_talker_mtp_uses_graph_for_legacy_or_explicit_safe_model(monkeypatch, talker_mtp_graph_safe):
|
|
runner = _runner_for_talker_graph_init(talker_mtp_graph_safe=talker_mtp_graph_safe)
|
|
wrapped = object()
|
|
monkeypatch.setattr(
|
|
"vllm_omni.worker.gpu_model_runner.current_omni_platform.get_graph_wrapper_cls",
|
|
lambda: lambda *args, **kwargs: wrapped,
|
|
)
|
|
|
|
OmniGPUModelRunner._init_talker_mtp(runner)
|
|
|
|
assert runner.talker_mtp is wrapped
|
|
|
|
|
|
class DummyBuffer:
|
|
"""A minimal buffer wrapper that exposes the `.gpu` attribute."""
|
|
|
|
def __init__(self, t: torch.Tensor):
|
|
self.gpu = t
|
|
|
|
|
|
class DummyInputBatch:
|
|
"""A minimal input batch that only provides `req_ids`."""
|
|
|
|
def __init__(self, req_ids):
|
|
self.req_ids = req_ids
|
|
self.req_id_to_index = {r: i for i, r in enumerate(req_ids)}
|
|
|
|
|
|
class DummyReqState:
|
|
"""A minimal request state container."""
|
|
|
|
mm_features: list[str]
|
|
additional_information_cpu: dict
|
|
|
|
|
|
def test_model_forward_passes_request_ids_to_decode_metadata(monkeypatch):
|
|
received = {}
|
|
model = SimpleNamespace(
|
|
supports_omni_decode_step_metadata=True,
|
|
update_decode_step_metadata=lambda **kwargs: received.update(kwargs),
|
|
)
|
|
runner = object.__new__(OmniGPUModelRunner)
|
|
runner.model = model
|
|
runner.input_batch = DummyInputBatch(["request-a", "request-b"])
|
|
runner._build_model_kwargs_extra = lambda: {}
|
|
monkeypatch.setattr(GPUModelRunner, "_model_forward", lambda *_args, **_kwargs: torch.zeros(1))
|
|
|
|
OmniGPUModelRunner._model_forward(runner, input_ids=torch.ones(2, dtype=torch.long))
|
|
|
|
assert received["req_ids"] == ["request-a", "request-b"]
|
|
|
|
|
|
class MiMoAudioForConditionalGeneration(torch.nn.Module):
|
|
"""Dummy model whose class name must exactly match the production check."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
# No real forward needed for these tests.
|
|
|
|
|
|
class DummyTalkerMTP(torch.nn.Module):
|
|
"""A fake talker_mtp module for deterministic CPU testing."""
|
|
|
|
def forward(
|
|
self,
|
|
req_input_ids,
|
|
req_embeds,
|
|
last_talker_hidden,
|
|
text_step,
|
|
do_sample=None,
|
|
temperature=None,
|
|
top_k=None,
|
|
top_p=None,
|
|
):
|
|
# Deterministic behavior:
|
|
# - output embeds = input embeds + 1
|
|
# - output codes = [[0], [1], ...]
|
|
bsz = req_embeds.shape[0]
|
|
new_embeds = req_embeds + 1.0
|
|
codes = torch.arange(bsz, dtype=torch.int64).view(bsz, 1)
|
|
return new_embeds, codes
|
|
|
|
|
|
class CaptureTalkerMTP(torch.nn.Module):
|
|
"""A fake talker_mtp module that records sampling kwargs."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.calls = []
|
|
|
|
def forward(
|
|
self,
|
|
req_input_ids,
|
|
req_embeds,
|
|
last_talker_hidden,
|
|
text_step,
|
|
do_sample=None,
|
|
temperature=None,
|
|
top_k=None,
|
|
top_p=None,
|
|
generator=None,
|
|
generators=None,
|
|
):
|
|
self.calls.append(
|
|
{
|
|
"batch_size": int(req_embeds.shape[0]),
|
|
"do_sample": do_sample,
|
|
"temperature": temperature,
|
|
"top_k": top_k,
|
|
"top_p": top_p,
|
|
"generator": generator,
|
|
"generators": generators,
|
|
}
|
|
)
|
|
codes = torch.zeros((req_embeds.shape[0], 1), dtype=torch.int64)
|
|
return req_embeds, codes
|
|
|
|
|
|
class StrictMRoPEModel:
|
|
def get_mrope_input_positions(self, input_tokens, mm_features):
|
|
raise NotImplementedError
|
|
|
|
|
|
class FlexibleMRoPEModel:
|
|
def get_mrope_input_positions(self, input_tokens, mm_features=None, **kwargs):
|
|
raise NotImplementedError
|
|
|
|
|
|
@contextmanager
|
|
def _noop_forward_context(*args, **kwargs):
|
|
"""A no-op context manager to replace vLLM forward context in CPU tests."""
|
|
yield
|
|
|
|
|
|
def test_filter_mrope_kwargs_for_strict_model_signature():
|
|
kwargs = {
|
|
"mm_features": ["audio"],
|
|
"hf_config": object(),
|
|
"image_grid_thw": [],
|
|
}
|
|
|
|
assert _filter_mrope_kwargs_for_model(StrictMRoPEModel(), kwargs) == {
|
|
"mm_features": ["audio"],
|
|
}
|
|
|
|
|
|
def test_filter_mrope_kwargs_preserves_flexible_model_kwargs():
|
|
kwargs = {
|
|
"mm_features": ["video"],
|
|
"hf_config": object(),
|
|
"video_grid_thw": [[1, 2, 3]],
|
|
}
|
|
|
|
assert _filter_mrope_kwargs_for_model(FlexibleMRoPEModel(), kwargs) is kwargs
|
|
|
|
|
|
def _make_runner(req_ids=("r1", "r2"), hidden_size=4):
|
|
# Create an instance without calling OmniGPUModelRunner.__init__
|
|
runner = object.__new__(OmniGPUModelRunner)
|
|
|
|
# Minimal attributes used by OmniGPUModelRunner._talker_mtp_forward
|
|
runner.input_batch = DummyInputBatch(list(req_ids))
|
|
runner.requests = {rid: DummyReqState() for rid in req_ids}
|
|
runner.model_intermediate_buffer = {}
|
|
|
|
# query_start_loc.cpu[req_index] is used to locate the token position
|
|
# in the flattened `inputs_embeds`.
|
|
runner.query_start_loc = type("QSL", (), {})()
|
|
# Map: r1 -> offset 0, r2 -> offset 3
|
|
runner.query_start_loc.cpu = torch.tensor([0, 3], dtype=torch.int32)
|
|
|
|
bsz = len(req_ids)
|
|
runner.talker_mtp_input_ids = DummyBuffer(torch.zeros((bsz,), dtype=torch.int64))
|
|
runner.talker_mtp_inputs_embeds = DummyBuffer(torch.zeros((bsz, hidden_size), dtype=torch.float32))
|
|
runner.last_talker_hidden = DummyBuffer(torch.zeros((bsz, hidden_size), dtype=torch.float32))
|
|
runner.text_step = DummyBuffer(torch.zeros((bsz, hidden_size), dtype=torch.float32))
|
|
|
|
runner.talker_mtp = DummyTalkerMTP()
|
|
runner.model = SimpleNamespace(talker_mtp_output_key=("codes", "audio"))
|
|
runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace())
|
|
|
|
# Provide a minimal implementation that returns the expected 4-tuple.
|
|
def _determine_batch_execution_and_padding(**kwargs):
|
|
return None, object(), None, None, None
|
|
|
|
runner._determine_batch_execution_and_padding = _determine_batch_execution_and_padding
|
|
|
|
# Use the real merge method from OmniGPUModelRunner.
|
|
return runner
|
|
|
|
|
|
def _make_runner_for_mimo(req_id="r_mimo"):
|
|
"""Create a minimal runner with MiMoAudio-like model and request state."""
|
|
runner = object.__new__(OmniGPUModelRunner)
|
|
runner.model = MiMoAudioForConditionalGeneration()
|
|
|
|
# Minimal vllm_config / model_config used by helper.
|
|
class _DummyModelConfig:
|
|
async_chunk = False
|
|
|
|
class _DummyVllmConfig:
|
|
model_config = _DummyModelConfig()
|
|
|
|
runner.vllm_config = _DummyVllmConfig()
|
|
|
|
# Attach a single request state with mm_features and additional_information_cpu.
|
|
req_state = DummyReqState()
|
|
req_state.mm_features = ["mm_feature_obj"]
|
|
req_state.additional_information_cpu = {"some_key": "some_value"}
|
|
|
|
runner.requests = {req_id: req_state}
|
|
|
|
return runner
|
|
|
|
|
|
def test_talker_mtp_forward_cpu_updates_inputs_and_info(monkeypatch):
|
|
# `_talker_mtp_forward` dispatches through the active platform.
|
|
import vllm_omni.worker.gpu_model_runner as mod # Must be the same module that defines OmniGPUModelRunner
|
|
|
|
monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
|
|
|
|
runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
|
|
|
|
def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
|
|
batch_desc = SimpleNamespace(num_tokens=int(num_tokens))
|
|
return (False, batch_desc, None, None, None)
|
|
|
|
monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner)))
|
|
|
|
# Initialize per-request embeds (batch-major inside talker_mtp_inputs_embeds)
|
|
runner.talker_mtp_inputs_embeds.gpu[0] = torch.tensor([1.0, 2.0, 3.0, 4.0])
|
|
runner.talker_mtp_inputs_embeds.gpu[1] = torch.tensor([10.0, 20.0, 30.0, 40.0])
|
|
|
|
# Flattened `inputs_embeds`: offsets 0 and 3 will be overwritten
|
|
inputs_embeds = torch.zeros((6, 4), dtype=torch.float32)
|
|
|
|
# Call the original implementation from OmniGPUModelRunner (no re-implementation)
|
|
OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds)
|
|
|
|
# Validate embeds were written back (+1)
|
|
assert torch.allclose(inputs_embeds[0], torch.tensor([2.0, 3.0, 4.0, 5.0]))
|
|
assert torch.allclose(inputs_embeds[3], torch.tensor([11.0, 21.0, 31.0, 41.0]))
|
|
|
|
# Validate per-request additional_information_cpu was updated
|
|
info_r1 = runner.requests["r1"].additional_information_cpu
|
|
info_r2 = runner.requests["r2"].additional_information_cpu
|
|
assert int(info_r1["codes"]["audio"][0, 0]) == 0
|
|
assert int(info_r2["codes"]["audio"][0, 0]) == 1
|
|
|
|
|
|
def test_talker_mtp_forward_cpu_empty_batch_noop(monkeypatch):
|
|
import vllm_omni.worker.gpu_model_runner as mod
|
|
|
|
monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
|
|
|
|
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
|
|
|
inputs_embeds = torch.randn((2, 4))
|
|
before = inputs_embeds.clone()
|
|
|
|
OmniGPUModelRunner._talker_mtp_forward(runner, [], inputs_embeds)
|
|
|
|
# Ensure no changes were made
|
|
assert torch.allclose(inputs_embeds, before)
|
|
|
|
|
|
def test_talker_mtp_forward_ignores_default_sampling_seed_without_request_marker(monkeypatch):
|
|
import vllm_omni.worker.gpu_model_runner as mod
|
|
|
|
monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
|
|
|
|
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
|
runner.requests["r1"].sampling_params = SimpleNamespace(seed=42)
|
|
runner.talker_mtp = CaptureTalkerMTP()
|
|
runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))
|
|
|
|
def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
|
|
batch_desc = SimpleNamespace(num_tokens=int(num_tokens))
|
|
return (False, batch_desc, None, None, None)
|
|
|
|
monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner)))
|
|
|
|
inputs_embeds = torch.zeros((2, 4), dtype=torch.float32)
|
|
OmniGPUModelRunner._talker_mtp_forward(runner, ["r1"], inputs_embeds)
|
|
|
|
assert runner.talker_mtp.calls[0]["generator"] is None
|
|
|
|
|
|
def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker(monkeypatch):
|
|
import vllm_omni.worker.gpu_model_runner as mod
|
|
|
|
monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
|
|
|
|
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
|
runner.requests["r1"].sampling_params = SimpleNamespace(
|
|
seed=42,
|
|
extra_args={"tts_local_seed": 42},
|
|
)
|
|
runner.talker_mtp = CaptureTalkerMTP()
|
|
runner.vllm_config = SimpleNamespace(
|
|
model_config=SimpleNamespace(
|
|
subtalker_sampling_params={
|
|
"do_sample": False,
|
|
"temperature": 0.2,
|
|
"top_k": 9,
|
|
"top_p": 0.55,
|
|
}
|
|
)
|
|
)
|
|
|
|
def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
|
|
batch_desc = SimpleNamespace(num_tokens=int(num_tokens))
|
|
return (False, batch_desc, None, None, None)
|
|
|
|
monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner)))
|
|
|
|
inputs_embeds = torch.zeros((2, 4), dtype=torch.float32)
|
|
OmniGPUModelRunner._talker_mtp_forward(runner, ["r1"], inputs_embeds)
|
|
|
|
assert runner.talker_mtp.calls == [
|
|
{
|
|
"batch_size": 1,
|
|
"do_sample": False,
|
|
"temperature": 0.2,
|
|
"top_k": 9,
|
|
"top_p": 0.55,
|
|
"generator": runner.talker_mtp.calls[0]["generator"],
|
|
"generators": None,
|
|
}
|
|
]
|
|
assert runner.talker_mtp.calls[0]["generator"] is not None
|
|
|
|
|
|
def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch):
|
|
import vllm_omni.worker.gpu_model_runner as mod
|
|
|
|
monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
|
|
|
|
runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
|
|
runner.requests["r1"].sampling_params = SimpleNamespace(
|
|
seed=11,
|
|
extra_args={"tts_local_seed": 11},
|
|
)
|
|
runner.requests["r2"].sampling_params = SimpleNamespace(
|
|
seed=22,
|
|
extra_args={"tts_local_seed": 22},
|
|
)
|
|
runner.talker_mtp = CaptureTalkerMTP()
|
|
runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))
|
|
|
|
def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
|
|
batch_desc = SimpleNamespace(num_tokens=int(num_tokens))
|
|
return (False, batch_desc, None, None, None)
|
|
|
|
monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner)))
|
|
|
|
runner.talker_mtp_input_ids.gpu[:] = torch.tensor([101, 202], dtype=torch.int64)
|
|
runner.talker_mtp_inputs_embeds.gpu[0] = torch.tensor([1.0, 2.0, 3.0, 4.0])
|
|
runner.talker_mtp_inputs_embeds.gpu[1] = torch.tensor([10.0, 20.0, 30.0, 40.0])
|
|
saved_input_ids = runner.talker_mtp_input_ids.gpu.clone()
|
|
saved_embeds = runner.talker_mtp_inputs_embeds.gpu.clone()
|
|
|
|
inputs_embeds = torch.zeros((6, 4), dtype=torch.float32)
|
|
OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds)
|
|
|
|
assert [call["batch_size"] for call in runner.talker_mtp.calls] == [1, 1]
|
|
assert all(call["generator"] is not None for call in runner.talker_mtp.calls)
|
|
assert runner.talker_mtp.calls[0]["generator"] is not runner.talker_mtp.calls[1]["generator"]
|
|
assert torch.equal(runner.talker_mtp_input_ids.gpu, saved_input_ids)
|
|
assert torch.equal(runner.talker_mtp_inputs_embeds.gpu, saved_embeds)
|
|
|
|
|
|
def test_talker_mtp_forward_batches_seeded_requests_for_opted_in_models(monkeypatch):
|
|
"""Models with talker_mtp_accepts_per_row_generators get one batched call (#4883)."""
|
|
import vllm_omni.worker.gpu_model_runner as mod
|
|
|
|
monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
|
|
|
|
runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
|
|
runner.requests["r1"].sampling_params = SimpleNamespace(
|
|
seed=11,
|
|
extra_args={"tts_local_seed": 11},
|
|
)
|
|
runner.requests["r2"].sampling_params = SimpleNamespace(
|
|
seed=22,
|
|
extra_args={"tts_local_seed": 22},
|
|
)
|
|
runner.talker_mtp = CaptureTalkerMTP()
|
|
runner.model = SimpleNamespace(
|
|
talker_mtp_output_key=("codes", "audio"),
|
|
talker_mtp_accepts_per_row_generators=True,
|
|
)
|
|
runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))
|
|
|
|
def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
|
|
batch_desc = SimpleNamespace(num_tokens=int(num_tokens))
|
|
return (False, batch_desc, None, None, None)
|
|
|
|
monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner)))
|
|
|
|
inputs_embeds = torch.zeros((6, 4), dtype=torch.float32)
|
|
OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds)
|
|
|
|
# One batched call with distinct per-row generators, not two scalar calls.
|
|
assert [call["batch_size"] for call in runner.talker_mtp.calls] == [2]
|
|
row_generators = runner.talker_mtp.calls[0]["generators"]
|
|
assert runner.talker_mtp.calls[0]["generator"] is None
|
|
assert len(row_generators) == 2
|
|
assert all(generator is not None for generator in row_generators)
|
|
assert row_generators[0] is not row_generators[1]
|
|
|
|
# The per-request generator stream persists across steps...
|
|
OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds)
|
|
assert runner.talker_mtp.calls[1]["generators"][0] is row_generators[0]
|
|
assert runner.talker_mtp.calls[1]["generators"][1] is row_generators[1]
|
|
|
|
# ...and is evicted once its request finishes.
|
|
del runner.requests["r2"]
|
|
OmniGPUModelRunner._talker_mtp_forward(runner, ["r1"], inputs_embeds)
|
|
assert set(runner._talker_mtp_generators) == {"r1"}
|
|
assert runner.talker_mtp.calls[2]["generator"] is row_generators[0]
|
|
|
|
|
|
def test_update_intermediate_buffer_writes_to_buffer_and_setattr(monkeypatch):
|
|
"""Validate that _update_intermediate_buffer writes to model_intermediate_buffer
|
|
(forward path) and mirrors to additional_information_cpu setattr (backward compat)."""
|
|
import vllm_omni.worker.gpu_model_runner as mod
|
|
|
|
monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
|
|
|
|
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
|
|
|
update = {"my_tensor": torch.tensor([1.0, 2.0]), "my_list": [3, 4]}
|
|
OmniGPUModelRunner._update_intermediate_buffer(runner, "r1", update)
|
|
|
|
# Forward: buffer is populated
|
|
assert "r1" in runner.model_intermediate_buffer
|
|
buf = runner.model_intermediate_buffer["r1"]
|
|
assert torch.allclose(buf["my_tensor"], torch.tensor([1.0, 2.0]))
|
|
assert buf["my_list"] == [3, 4]
|
|
|
|
# Backward compat: setattr is also populated
|
|
info_cpu = runner.requests["r1"].additional_information_cpu
|
|
assert torch.allclose(info_cpu["my_tensor"], torch.tensor([1.0, 2.0]))
|
|
assert info_cpu["my_list"] == [3, 4]
|
|
|
|
|
|
def test_update_intermediate_buffer_accumulates():
|
|
"""Validate that successive merges accumulate keys in the buffer."""
|
|
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
|
|
|
OmniGPUModelRunner._update_intermediate_buffer(runner, "r1", {"a": torch.tensor([1.0])})
|
|
OmniGPUModelRunner._update_intermediate_buffer(runner, "r1", {"b": torch.tensor([2.0])})
|
|
|
|
buf = runner.model_intermediate_buffer["r1"]
|
|
assert "a" in buf and "b" in buf
|
|
assert torch.allclose(buf["a"], torch.tensor([1.0]))
|
|
assert torch.allclose(buf["b"], torch.tensor([2.0]))
|
|
|
|
|
|
def test_update_additional_information_deserializes_new_request_payload():
|
|
from vllm_omni.engine.serialization import serialize_additional_information
|
|
|
|
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
|
conditioning = {
|
|
"tts_token_ids": torch.tensor([1, 2]),
|
|
"tts_hidden_states": torch.ones(2, 4),
|
|
}
|
|
scheduler_output = SimpleNamespace(
|
|
scheduled_new_reqs=[
|
|
SimpleNamespace(
|
|
req_id="r1",
|
|
additional_information=serialize_additional_information(conditioning),
|
|
)
|
|
],
|
|
scheduled_cached_reqs=SimpleNamespace(),
|
|
)
|
|
|
|
OmniGPUModelRunner._update_additional_information(runner, scheduler_output)
|
|
|
|
assert torch.equal(runner.model_intermediate_buffer["r1"]["tts_token_ids"], conditioning["tts_token_ids"])
|
|
assert torch.equal(
|
|
runner.model_intermediate_buffer["r1"]["tts_hidden_states"],
|
|
conditioning["tts_hidden_states"],
|
|
)
|
|
|
|
|
|
def test_streaming_new_request_marker_replaces_terminal_chunk_snapshot():
|
|
from vllm_omni.engine.serialization import serialize_additional_information
|
|
|
|
runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
|
|
runner.model.replace_runtime_additional_information = True
|
|
terminal = {
|
|
"codes": {"audio": torch.tensor([1, 2])},
|
|
"meta": {"cache_epoch": 0, "chunk_seq": 2, "last_chunk": True},
|
|
}
|
|
peer = {
|
|
"codes": {"audio": torch.tensor([9])},
|
|
"meta": {"cache_epoch": 3, "chunk_seq": 1, "last_chunk": False},
|
|
}
|
|
runner.model_intermediate_buffer.update(r1=terminal, r2=peer)
|
|
marker = {
|
|
"meta": {
|
|
"finished": False,
|
|
"is_segment_finished": True,
|
|
"request_finished": False,
|
|
"replace_runtime_additional_information": True,
|
|
}
|
|
}
|
|
new_req = SimpleNamespace(
|
|
req_id="r1",
|
|
model_intermediate_buffer=marker,
|
|
additional_information=serialize_additional_information(terminal),
|
|
)
|
|
|
|
OmniGPUModelRunner._update_streaming_input_additional_info(runner, new_req, "r1")
|
|
OmniGPUModelRunner._update_additional_information(
|
|
runner,
|
|
SimpleNamespace(
|
|
scheduled_new_reqs=[new_req],
|
|
scheduled_cached_reqs=SimpleNamespace(),
|
|
),
|
|
)
|
|
|
|
info = runner.model_intermediate_buffer["r1"]
|
|
assert "codes" not in info
|
|
assert info["meta"] == {
|
|
**marker["meta"],
|
|
"num_processed_tokens": 0,
|
|
"resumable": True,
|
|
}
|
|
assert runner.requests["r1"].additional_information_cpu == info
|
|
assert runner.model_intermediate_buffer["r2"] == peer
|
|
|
|
|
|
def test_cached_empty_marker_replaces_terminal_chunk_snapshot():
|
|
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
|
runner.model.replace_runtime_additional_information = True
|
|
runner.model_intermediate_buffer["r1"] = {
|
|
"codes": {"audio": torch.tensor([1, 2])},
|
|
"meta": {"cache_epoch": 0, "chunk_seq": 2, "last_chunk": True},
|
|
}
|
|
marker = {
|
|
"meta": {
|
|
"is_segment_finished": torch.tensor(True, dtype=torch.bool),
|
|
"replace_runtime_additional_information": True,
|
|
}
|
|
}
|
|
|
|
OmniGPUModelRunner._update_additional_information(
|
|
runner,
|
|
SimpleNamespace(
|
|
scheduled_new_reqs=[],
|
|
scheduled_cached_reqs=SimpleNamespace(additional_information={"r1": marker}),
|
|
),
|
|
)
|
|
|
|
assert runner.model_intermediate_buffer["r1"] == marker
|
|
assert runner.requests["r1"].additional_information_cpu == marker
|
|
|
|
|
|
def test_update_intermediate_buffer_skips_empty_update():
|
|
"""Validate that an empty update dict is a no-op."""
|
|
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
|
|
|
OmniGPUModelRunner._update_intermediate_buffer(runner, "r1", {})
|
|
|
|
assert "r1" not in runner.model_intermediate_buffer
|
|
|
|
|
|
def test_update_intermediate_buffer_skips_unknown_req_id():
|
|
"""Validate that merge is a no-op when req_id is not in self.requests."""
|
|
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
|
|
|
OmniGPUModelRunner._update_intermediate_buffer(runner, "unknown_req", {"key": torch.tensor([1.0])})
|
|
|
|
assert "unknown_req" not in runner.model_intermediate_buffer
|
|
|
|
|
|
def test_streaming_input_update_merges_model_intermediate_buffer():
|
|
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
|
runner.model_intermediate_buffer["r1"] = {
|
|
"duplex": {
|
|
"session_id": "sid",
|
|
"seq": 1,
|
|
}
|
|
}
|
|
runner.requests["r1"].additional_information_cpu = runner.model_intermediate_buffer["r1"]
|
|
new_req_data = SimpleNamespace(
|
|
model_intermediate_buffer={
|
|
"duplex": {
|
|
"session_id": "sid",
|
|
"seq": 2,
|
|
"payload": {"type": "audio"},
|
|
}
|
|
},
|
|
additional_information=None,
|
|
)
|
|
|
|
OmniGPUModelRunner._update_streaming_input_additional_info(runner, new_req_data, "r1")
|
|
|
|
info = runner.model_intermediate_buffer["r1"]
|
|
assert info["duplex"]["session_id"] == "sid"
|
|
assert info["duplex"]["seq"] == 2
|
|
assert info["duplex"]["payload"] == {"type": "audio"}
|
|
assert runner.requests["r1"].additional_information_cpu is info
|
|
|
|
|
|
def _make_full_payload_accumulation_runner(
|
|
model_arch="Qwen3OmniMoeForConditionalGeneration",
|
|
model_stage="talker",
|
|
async_chunk=False,
|
|
final_output=False,
|
|
custom_process_next_stage_input_func="module.full_payload",
|
|
):
|
|
runner = object.__new__(OmniConnectorModelRunnerMixin)
|
|
runner.model_config = SimpleNamespace(
|
|
model_arch=model_arch,
|
|
model_stage=model_stage,
|
|
async_chunk=async_chunk,
|
|
final_output=final_output,
|
|
custom_process_next_stage_input_func=custom_process_next_stage_input_func,
|
|
)
|
|
runner._custom_process_func = object()
|
|
runner._pending_full_payload_send = {}
|
|
runner._stage_id = 1
|
|
runner._omni_connector = object()
|
|
return runner
|
|
|
|
|
|
def test_accumulate_full_payload_output_preserves_aligned_all_zero_qwen3_omni_codec_rows():
|
|
runner = _make_full_payload_accumulation_runner()
|
|
request = SimpleNamespace(output_token_ids=[0, 1])
|
|
codes = torch.zeros((2, 3), dtype=torch.long)
|
|
|
|
OmniConnectorModelRunnerMixin.accumulate_full_payload_output(runner, "r1", {"codes.audio": codes}, request)
|
|
|
|
stored, _ = OmniConnectorModelRunnerMixin._materialize_full_payload_entry(runner._pending_full_payload_send["r1"])
|
|
assert torch.equal(stored["codes.audio"], codes)
|
|
|
|
|
|
def test_accumulate_full_payload_output_keeps_misaligned_all_zero_qwen3_omni_codec_rows():
|
|
# After removing the sender-side zero filter, the accumulator keeps every
|
|
# codec row including misaligned all-zero rows. The downstream consumer
|
|
# (_extract_qwen3_full_payload_codec_rows) is the authoritative crop and
|
|
# filters by output_token_ids.
|
|
runner = _make_full_payload_accumulation_runner()
|
|
request = SimpleNamespace(output_token_ids=[0, 1])
|
|
codes = torch.zeros((1, 3), dtype=torch.long)
|
|
|
|
OmniConnectorModelRunnerMixin.accumulate_full_payload_output(runner, "r1", {"codes.audio": codes}, request)
|
|
|
|
stored, _ = OmniConnectorModelRunnerMixin._materialize_full_payload_entry(runner._pending_full_payload_send["r1"])
|
|
assert "codes.audio" in stored
|
|
assert torch.equal(stored["codes.audio"], codes)
|
|
|
|
|
|
def test_accumulate_full_payload_output_preserves_incremental_aligned_all_zero_qwen3_omni_codec_rows():
|
|
runner = _make_full_payload_accumulation_runner()
|
|
request = SimpleNamespace(output_token_ids=[0, 1])
|
|
runner._pending_full_payload_send["r1"] = (
|
|
{"codes.audio": torch.ones((1, 3), dtype=torch.long)},
|
|
request,
|
|
)
|
|
codes = torch.zeros((1, 3), dtype=torch.long)
|
|
|
|
OmniConnectorModelRunnerMixin.accumulate_full_payload_output(runner, "r1", {"codes.audio": codes}, request)
|
|
|
|
stored, _ = OmniConnectorModelRunnerMixin._materialize_full_payload_entry(runner._pending_full_payload_send["r1"])
|
|
assert stored["codes.audio"].shape == (2, 3)
|
|
assert torch.equal(stored["codes.audio"][1], torch.zeros(3, dtype=torch.long))
|
|
|
|
|
|
def test_accumulate_full_payload_output_keeps_all_zero_qwen3_omni_prefill_placeholder():
|
|
# Prefill placeholder rows (output_token_ids empty) are no longer dropped
|
|
# at the sender. The consumer-side crop trims them off using
|
|
# output_token_ids, so the end-to-end semantics are unchanged.
|
|
runner = _make_full_payload_accumulation_runner()
|
|
request = SimpleNamespace(output_token_ids=[])
|
|
codes = torch.zeros((2, 3), dtype=torch.long)
|
|
|
|
OmniConnectorModelRunnerMixin.accumulate_full_payload_output(runner, "r1", {"codes.audio": codes}, request)
|
|
|
|
stored, _ = OmniConnectorModelRunnerMixin._materialize_full_payload_entry(runner._pending_full_payload_send["r1"])
|
|
assert "codes.audio" in stored
|
|
assert torch.equal(stored["codes.audio"], codes)
|
|
|
|
|
|
def test_full_payload_output_accumulation_hook_matrix():
|
|
assert _make_full_payload_accumulation_runner(model_stage="thinker")._should_accumulate_full_payload_output()
|
|
assert _make_full_payload_accumulation_runner(model_stage="talker")._should_accumulate_full_payload_output()
|
|
assert not _make_full_payload_accumulation_runner(
|
|
model_stage="code2wav", final_output=True
|
|
)._should_accumulate_full_payload_output()
|
|
assert not _make_full_payload_accumulation_runner(
|
|
model_stage="token2audio",
|
|
custom_process_next_stage_input_func=None,
|
|
)._should_accumulate_full_payload_output()
|
|
assert not _make_full_payload_accumulation_runner(
|
|
model_stage="talker", async_chunk=True
|
|
)._should_accumulate_full_payload_output()
|
|
for model_arch in (
|
|
"Qwen3TTSForConditionalGeneration",
|
|
"Qwen2_5OmniForConditionalGeneration",
|
|
):
|
|
runner = _make_full_payload_accumulation_runner(model_arch=model_arch)
|
|
runner._custom_process_func = None
|
|
assert not runner._should_accumulate_full_payload_output()
|
|
|
|
|
|
def _make_request_end_payload_runner(*, enabled=True, prefix_cache=None):
|
|
runner = object.__new__(GPUARModelRunner)
|
|
runner.model = SimpleNamespace(omni_payload_at_request_end=enabled)
|
|
runner.omni_prefix_cache = prefix_cache
|
|
runner.model_config = SimpleNamespace(
|
|
model_arch="IndexTTS25TalkerForConditionalGeneration",
|
|
model_stage="indextts2_5_talker",
|
|
async_chunk=False,
|
|
final_output=False,
|
|
custom_process_next_stage_input_func="module.full_payload",
|
|
)
|
|
runner._custom_process_func = object()
|
|
runner._pending_full_payload_send = {}
|
|
runner._stage_id = 0
|
|
runner._omni_connector = object()
|
|
return runner
|
|
|
|
|
|
def test_request_end_payload_d2h_gate_requires_opt_in_and_no_prefix_cache():
|
|
assert _make_request_end_payload_runner()._should_defer_full_payload_d2h()
|
|
assert not _make_request_end_payload_runner(enabled=False)._should_defer_full_payload_d2h()
|
|
assert not _make_request_end_payload_runner(prefix_cache=object())._should_defer_full_payload_d2h()
|
|
|
|
|
|
def test_request_end_payload_suppresses_per_step_multimodal_outputs():
|
|
runner = _make_request_end_payload_runner()
|
|
|
|
def unexpected_build(_payload):
|
|
raise AssertionError("request-end payloads must stay inside the GPU accumulator")
|
|
|
|
runner._build_multimodal_outputs = unexpected_build
|
|
pooler_inter = [{"codes.mel": torch.tensor([[7]])}]
|
|
|
|
inter_stage, client = runner._build_omni_step_outputs(
|
|
pooler_inter,
|
|
pooler_inter,
|
|
defer_full_payload_d2h=True,
|
|
)
|
|
|
|
assert inter_stage is None
|
|
assert client is None
|
|
|
|
|
|
def test_sync_local_stage_payloads_retains_payload_until_request_is_active():
|
|
runner = object.__new__(OmniGPUModelRunner)
|
|
payload = {"codes": {"audio": [1, 2, 3]}}
|
|
runner._local_stage_payload_cache = {"late": payload}
|
|
runner._full_payload_pending_broadcast_req_ids = set()
|
|
runner.requests = {}
|
|
runner.model_intermediate_buffer = {}
|
|
|
|
OmniGPUModelRunner._sync_local_stage_payloads(runner)
|
|
|
|
assert runner._local_stage_payload_cache == {"late": payload}
|
|
assert runner.model_intermediate_buffer == {}
|
|
|
|
runner.requests = {"late": DummyReqState()}
|
|
OmniGPUModelRunner._sync_local_stage_payloads(runner)
|
|
|
|
assert runner._local_stage_payload_cache == {}
|
|
assert runner.model_intermediate_buffer["late"] == payload
|
|
assert runner.requests["late"].additional_information_cpu == payload
|
|
|
|
|
|
def test_maybe_attach_mimo_audio_req_infos_enriches_dict():
|
|
runner = _make_runner_for_mimo()
|
|
req_id = "r_mimo"
|
|
req_state = runner.requests[req_id]
|
|
|
|
# Existing req_infos should be copied and enriched, not mutated in place.
|
|
original_req_infos = {"existing": 1}
|
|
enriched = OmniGPUModelRunner._maybe_attach_mimo_audio_req_infos(runner, req_state, original_req_infos, req_id)
|
|
|
|
assert enriched is not original_req_infos
|
|
assert enriched["existing"] == 1
|
|
# mm_features should be filled from req_state when missing
|
|
assert enriched["mm_features"] == req_state.mm_features
|
|
# req_id should always be attached
|
|
assert enriched["req_id"] == req_id
|
|
|
|
|
|
def test_maybe_attach_mimo_audio_req_infos_no_req_state_returns_input():
|
|
runner = _make_runner_for_mimo()
|
|
req_id = "missing"
|
|
req_state = None
|
|
req_infos = {"k": "v"}
|
|
|
|
result = OmniGPUModelRunner._maybe_attach_mimo_audio_req_infos(runner, req_state, req_infos, req_id)
|
|
|
|
# When no req_state, helper should be a no-op.
|
|
assert result is req_infos
|
|
|
|
|
|
def _make_phase_runner(monkeypatch, rows, *, batched_decode):
|
|
"""Exercise real preprocessing and MTP routing with CPU model buffers.
|
|
|
|
Each row is (request id, prompt length, computed tokens, scheduled tokens).
|
|
Only model computation and the device forward context are replaced.
|
|
"""
|
|
import vllm_omni.worker.gpu_model_runner as mod
|
|
|
|
monkeypatch.setattr(mod, "get_pp_group", lambda: SimpleNamespace(is_first_rank=True))
|
|
monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
|
|
runner = _make_runner(req_ids=tuple(row[0] for row in rows))
|
|
total = sum(row[3] for row in rows)
|
|
runner.supports_mm_inputs = False
|
|
runner.enable_prompt_embeds = False
|
|
runner.uses_mrope = False
|
|
runner.uses_xdrope_dim = 0
|
|
runner.has_talker_mtp = True
|
|
runner.model_config = SimpleNamespace(is_encoder_decoder=False)
|
|
runner.vllm_config.model_config.async_chunk = True
|
|
runner._init_model_kwargs = dict
|
|
runner.input_ids = DummyBuffer(torch.arange(total, dtype=torch.int64))
|
|
runner.inputs_embeds = DummyBuffer(torch.zeros(total, 4))
|
|
runner.positions = torch.arange(total)
|
|
runner.input_batch.num_computed_tokens_cpu = [row[2] for row in rows]
|
|
offsets = [0]
|
|
for rid, prompt_len, _, scheduled in rows:
|
|
runner.requests[rid].prompt_token_ids = list(range(prompt_len))
|
|
offsets.append(offsets[-1] + scheduled)
|
|
# Runner-owned metadata must override stale model/intermediate state.
|
|
runner.model_intermediate_buffer[rid] = {
|
|
"_omni_is_prefill": "stale",
|
|
"_omni_prompt_len": -1,
|
|
"_omni_num_computed_tokens": -1,
|
|
}
|
|
runner.query_start_loc.cpu = torch.tensor(offsets, dtype=torch.int32)
|
|
calls = []
|
|
batch_calls = []
|
|
|
|
def preprocess(input_ids, input_embeds, **info):
|
|
calls.append(("normal", dict(info)))
|
|
embeds = input_ids.float().view(-1, 1).expand(-1, 4).clone()
|
|
updates = {}
|
|
if not info["_omni_is_prefill"]:
|
|
updates["mtp_inputs"] = (torch.zeros_like(embeds), torch.zeros_like(embeds))
|
|
return input_ids, embeds, updates
|
|
|
|
def preprocess_decode_batch(input_ids, req_infos):
|
|
batch_calls.append([info["request_id"] for info in req_infos])
|
|
calls.extend(("batch", dict(info)) for info in req_infos)
|
|
embeds = input_ids.float().view(-1, 1).expand(-1, 4).clone()
|
|
return input_ids, embeds, torch.zeros_like(embeds), torch.zeros_like(embeds), [{} for _ in req_infos]
|
|
|
|
runner.model = SimpleNamespace(
|
|
has_preprocess=True,
|
|
preprocess=preprocess,
|
|
talker_mtp_output_key=("codes", "audio"),
|
|
)
|
|
if batched_decode:
|
|
runner.model.preprocess_decode_batch = preprocess_decode_batch
|
|
output = SchedulerOutput(
|
|
total_num_scheduled_tokens=total,
|
|
num_scheduled_tokens={row[0]: row[3] for row in rows},
|
|
scheduled_new_reqs=[],
|
|
scheduled_cached_reqs=CachedRequestData.make_empty(),
|
|
scheduled_spec_decode_tokens={},
|
|
scheduled_encoder_inputs={},
|
|
num_common_prefix_blocks=[],
|
|
finished_req_ids=set(),
|
|
free_encoder_mm_hashes=[],
|
|
)
|
|
return runner, output, calls, batch_calls
|
|
|
|
|
|
@pytest.mark.parametrize("batched_decode", [False, True], ids=["normal", "batch"])
|
|
@pytest.mark.parametrize("order", ["interleaved", "reversed", "adjacent"])
|
|
def test_preprocess_phase_contract_mixed_batch(monkeypatch, batched_decode, order):
|
|
rows = [
|
|
("decode-before", 5, 5, 1),
|
|
("tail", 5, 4, 1),
|
|
("decode-after", 5, 6, 1),
|
|
("single-prompt", 1, 0, 1),
|
|
("cached-prefill", 8, 4, 3),
|
|
("multi-token-decode", 5, 5, 2),
|
|
]
|
|
if order == "reversed":
|
|
rows.reverse()
|
|
elif order == "adjacent":
|
|
rows = [rows[1], rows[0], rows[2], *rows[3:]]
|
|
runner, scheduled, calls, batch_calls = _make_phase_runner(monkeypatch, rows, batched_decode=batched_decode)
|
|
_, embeds, *_ = runner._preprocess(scheduled, scheduled.total_num_scheduled_tokens)
|
|
|
|
if batched_decode:
|
|
decode_order = [rid for rid, _, _, _ in rows if rid.startswith("decode-")]
|
|
expected_batches = [decode_order] if order == "adjacent" else [[rid] for rid in decode_order]
|
|
assert batch_calls == expected_batches
|
|
else:
|
|
assert batch_calls == []
|
|
assert len(calls) == len(rows)
|
|
offset = 0
|
|
code_index = 0
|
|
for (route, info), (rid, prompt_len, computed, span) in zip(calls, rows, strict=True):
|
|
assert info["request_id"] == rid
|
|
assert type(info["_omni_prompt_len"]) is int
|
|
assert type(info["_omni_num_computed_tokens"]) is int
|
|
assert type(info["_omni_is_prefill"]) is bool
|
|
assert info["_omni_prompt_len"] == prompt_len
|
|
assert info["_omni_num_computed_tokens"] == computed
|
|
assert info["_omni_is_prefill"] == (computed < prompt_len)
|
|
mtp_eligible = rid in {"decode-before", "decode-after"}
|
|
assert route == ("batch" if batched_decode and mtp_eligible else "normal")
|
|
expected = torch.arange(offset, offset + span).float().view(-1, 1).expand(-1, 4)
|
|
torch.testing.assert_close(embeds[offset : offset + span], expected + int(mtp_eligible))
|
|
assert ("codes" in runner.model_intermediate_buffer[rid]) == mtp_eligible
|
|
if mtp_eligible:
|
|
torch.testing.assert_close(
|
|
runner.model_intermediate_buffer[rid]["codes"]["audio"],
|
|
torch.tensor([[code_index]], dtype=torch.int64),
|
|
rtol=0,
|
|
atol=0,
|
|
)
|
|
code_index += 1
|
|
offset += span
|
|
|
|
|
|
@pytest.mark.parametrize("batched_decode", [False, True], ids=["normal", "batch"])
|
|
def test_preprocess_one_token_chunked_prefill_tail_then_decode(monkeypatch, batched_decode):
|
|
# Replay a five-token prompt split at a four-token scheduling budget.
|
|
runner, scheduled, calls, _ = _make_phase_runner(monkeypatch, [("r", 5, 0, 4)], batched_decode=batched_decode)
|
|
runner._preprocess(scheduled, 4)
|
|
assert calls[-1][1]["_omni_is_prefill"] is True
|
|
assert "codes" not in runner.model_intermediate_buffer["r"]
|
|
|
|
for computed, is_prefill in [(4, True), (5, False)]:
|
|
runner.input_batch.num_computed_tokens_cpu = [computed]
|
|
scheduled.total_num_scheduled_tokens = 1
|
|
scheduled.num_scheduled_tokens = {"r": 1}
|
|
runner.query_start_loc.cpu = torch.tensor([0, 1], dtype=torch.int32)
|
|
runner.input_ids.gpu[0] = 42
|
|
_, embeds, *_ = runner._preprocess(scheduled, 1)
|
|
route, info = calls[-1]
|
|
assert info["_omni_is_prefill"] is is_prefill
|
|
assert info["_omni_num_computed_tokens"] == computed
|
|
assert info["_omni_prompt_len"] == 5
|
|
assert route == ("batch" if batched_decode and not is_prefill else "normal")
|
|
torch.testing.assert_close(embeds, torch.full((1, 4), 42.0 if is_prefill else 43.0))
|
|
assert ("codes" in runner.model_intermediate_buffer["r"]) == (not is_prefill)
|