Taichang ZhouandClaude Opus 4.8
9678ae58cf
Squashed migration rebased onto current main. Collapses the 46-commit dev/migrate-MR-v2 history (10 merge commits + superseded pre-rebase work) into a single commit, so the diff reflects only the genuine MR-V2 feature (worker_v2/ runners + model_states, omni schedulers, stage engine pipeline, pooling_output serialization bridge). Content already merged to main separately (step_audio2, tts_adapters, npu/_310p, etc.) drops out. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
169 lines
5.0 KiB
Python
169 lines
5.0 KiB
Python
"""Unit tests for OmniIntermediateBuffer lifecycle."""
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.worker_v2.model_states.intermediate_buffer import (
|
|
OmniIntermediateBuffer,
|
|
_resolve_additional_information,
|
|
_resolve_prompt_embeds,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
class _DummyInputBatch:
|
|
def __init__(self, idx_mapping):
|
|
self.idx_mapping_np = idx_mapping
|
|
|
|
|
|
def _make_new_req_data(req_id="r1", prompt_embeds=None, mm_features=None, additional_information=None):
|
|
ns = SimpleNamespace(
|
|
req_id=req_id,
|
|
mm_features=mm_features or [],
|
|
)
|
|
if prompt_embeds is not None:
|
|
ns.prompt_embeds = prompt_embeds
|
|
if additional_information is not None:
|
|
ns.additional_information = additional_information
|
|
return ns
|
|
|
|
|
|
# ---------------------------------------------------------------
|
|
# _resolve_prompt_embeds
|
|
# ---------------------------------------------------------------
|
|
|
|
|
|
def test_resolve_prompt_embeds_tensor():
|
|
t = torch.randn(3, 4, device="cpu")
|
|
result = _resolve_prompt_embeds(t)
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.shape == (3, 4)
|
|
assert result.is_contiguous()
|
|
|
|
|
|
def test_resolve_prompt_embeds_none():
|
|
assert _resolve_prompt_embeds(None) is None
|
|
|
|
|
|
# ---------------------------------------------------------------
|
|
# _resolve_additional_information
|
|
# ---------------------------------------------------------------
|
|
|
|
|
|
def test_resolve_additional_information_dict_passthrough():
|
|
d = {"key": "val"}
|
|
assert _resolve_additional_information(d) == d
|
|
|
|
|
|
def test_resolve_additional_information_none():
|
|
assert _resolve_additional_information(None) == {}
|
|
|
|
|
|
# ---------------------------------------------------------------
|
|
# OmniIntermediateBuffer core lifecycle
|
|
# ---------------------------------------------------------------
|
|
|
|
|
|
def test_add_and_remove():
|
|
buf = OmniIntermediateBuffer(max_num_reqs=4)
|
|
req = _make_new_req_data(req_id="r1")
|
|
buf.add_request(0, req)
|
|
assert buf.buffers[0]["req_id"] == "r1"
|
|
|
|
buf.remove_request(0)
|
|
assert buf.buffers[0] == {}
|
|
|
|
|
|
def test_add_with_prompt_embeds():
|
|
pe = torch.randn(2, 8)
|
|
req = _make_new_req_data(req_id="r1", prompt_embeds=pe)
|
|
buf = OmniIntermediateBuffer(max_num_reqs=4)
|
|
buf.add_request(0, req)
|
|
assert "prompt_embeds_cpu" in buf.buffers[0]
|
|
assert buf.buffers[0]["prompt_embeds_cpu"].shape == (2, 8)
|
|
|
|
|
|
def test_add_with_mm_features():
|
|
features = [torch.randn(3, 4)]
|
|
req = _make_new_req_data(req_id="r1", mm_features=features)
|
|
buf = OmniIntermediateBuffer(max_num_reqs=4)
|
|
buf.add_request(1, req)
|
|
assert buf.buffers[1]["mm_features"] is features
|
|
|
|
|
|
def test_add_with_additional_information():
|
|
info = {"audio_embeds": torch.randn(5)}
|
|
req = _make_new_req_data(req_id="r1", additional_information=info)
|
|
buf = OmniIntermediateBuffer(max_num_reqs=4)
|
|
buf.add_request(0, req)
|
|
assert torch.allclose(buf.buffers[0]["audio_embeds"], info["audio_embeds"])
|
|
|
|
|
|
# ---------------------------------------------------------------
|
|
# gather
|
|
# ---------------------------------------------------------------
|
|
|
|
|
|
def test_gather_returns_batch_order():
|
|
buf = OmniIntermediateBuffer(max_num_reqs=4)
|
|
buf.add_request(0, _make_new_req_data(req_id="r0"))
|
|
buf.add_request(2, _make_new_req_data(req_id="r2"))
|
|
buf.add_request(3, _make_new_req_data(req_id="r3"))
|
|
|
|
batch = _DummyInputBatch(idx_mapping=[2, 0, 3])
|
|
gathered = buf.gather(batch)
|
|
assert [g["req_id"] for g in gathered] == ["r2", "r0", "r3"]
|
|
|
|
|
|
# ---------------------------------------------------------------
|
|
# update
|
|
# ---------------------------------------------------------------
|
|
|
|
|
|
def test_update_merges_tensors_to_cpu():
|
|
buf = OmniIntermediateBuffer(max_num_reqs=2)
|
|
buf.add_request(0, _make_new_req_data(req_id="r0"))
|
|
|
|
buf.update(0, {"hidden": torch.randn(4)})
|
|
assert "hidden" in buf.buffers[0]
|
|
assert buf.buffers[0]["hidden"].device == torch.device("cpu")
|
|
|
|
|
|
def test_update_gpu_resident_keys():
|
|
buf = OmniIntermediateBuffer(max_num_reqs=2)
|
|
buf.add_request(0, _make_new_req_data(req_id="r0"))
|
|
|
|
t = torch.randn(4)
|
|
buf.update(0, {"kv": t}, gpu_resident_keys={"kv"})
|
|
assert "kv" in buf.buffers[0]
|
|
|
|
|
|
def test_update_list_values():
|
|
buf = OmniIntermediateBuffer(max_num_reqs=2)
|
|
buf.add_request(0, _make_new_req_data(req_id="r0"))
|
|
|
|
items = [torch.randn(2), "text"]
|
|
buf.update(0, {"items": items})
|
|
stored = buf.buffers[0]["items"]
|
|
assert isinstance(stored[0], torch.Tensor)
|
|
assert stored[1] == "text"
|
|
|
|
|
|
def test_update_empty_is_noop():
|
|
buf = OmniIntermediateBuffer(max_num_reqs=2)
|
|
buf.add_request(0, _make_new_req_data(req_id="r0"))
|
|
before = dict(buf.buffers[0])
|
|
buf.update(0, {})
|
|
assert buf.buffers[0] == before
|
|
|
|
|
|
def test_update_accumulates():
|
|
buf = OmniIntermediateBuffer(max_num_reqs=2)
|
|
buf.add_request(0, _make_new_req_data(req_id="r0"))
|
|
buf.update(0, {"a": torch.tensor([1.0])})
|
|
buf.update(0, {"b": torch.tensor([2.0])})
|
|
assert "a" in buf.buffers[0] and "b" in buf.buffers[0]
|