96 lines
3.7 KiB
Python
96 lines
3.7 KiB
Python
|
# SPDX-License-Identifier: Apache-2.0
|
||
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||
|
|
|
||
|
|
"""Unit tests for OmniIntermediateBuffer lifecycle.
|
||
|
|
|
||
|
|
Retained contracts: wire decode with malformed-payload rejection, slot
|
||
|
|
lifecycle with index migration on reuse, gather in batch order, update merge
|
||
|
|
semantics (cpu/gpu-resident keys), single owned snapshot for
|
||
|
|
update_gpu_tensor_rows, and decode-span merge until ack.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from types import SimpleNamespace
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from vllm_omni.engine import (
|
||
|
|
AdditionalInformationEntry,
|
||
|
|
AdditionalInformationPayload,
|
||
|
|
PromptEmbedsPayload,
|
||
|
|
)
|
||
|
|
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]
|
||
|
|
|
||
|
|
|
||
|
|
def _make_new_req_data(req_id="r1", **kwargs):
|
||
|
|
kwargs.setdefault("mm_features", [])
|
||
|
|
return SimpleNamespace(req_id=req_id, **kwargs)
|
||
|
|
|
||
|
|
|
||
|
|
def test_wire_decode_passthrough_and_malformed_rejection():
|
||
|
|
tensor = torch.randn(3, 4)
|
||
|
|
assert _resolve_prompt_embeds(tensor) is not None
|
||
|
|
assert _resolve_prompt_embeds(None) is None
|
||
|
|
with pytest.raises(ValueError, match="Failed to decode prompt_embeds payload"):
|
||
|
|
_resolve_prompt_embeds(PromptEmbedsPayload(data=b"\x00", shape=[2, 2], dtype="float32"))
|
||
|
|
|
||
|
|
assert _resolve_additional_information({"key": "val"}) == {"key": "val"}
|
||
|
|
assert _resolve_additional_information(None) == {}
|
||
|
|
payload = AdditionalInformationPayload(
|
||
|
|
entries={
|
||
|
|
"embed.decode": AdditionalInformationEntry(tensor_data=b"\x00", tensor_shape=[2], tensor_dtype="float32")
|
||
|
|
}
|
||
|
|
)
|
||
|
|
with pytest.raises(ValueError, match="Failed to decode additional_information payload"):
|
||
|
|
_resolve_additional_information(payload)
|
||
|
|
|
||
|
|
|
||
|
|
def test_slot_lifecycle_fields_reuse_and_remove_idempotent():
|
||
|
|
buf = OmniIntermediateBuffer(max_num_reqs=4)
|
||
|
|
embeds = torch.randn(2, 8)
|
||
|
|
features = [torch.randn(3, 4)]
|
||
|
|
params = SimpleNamespace(extra_args={"tts_local_seed": 1234})
|
||
|
|
buf.add_request(0, _make_new_req_data("r0", prompt_embeds=embeds, mm_features=features, sampling_params=params))
|
||
|
|
assert torch.equal(buf.buffers[0]["prompt_embeds_cpu"], embeds)
|
||
|
|
assert buf.buffers[0]["mm_features"] is features
|
||
|
|
assert buf.buffers[0]["sampling_params"] is params
|
||
|
|
|
||
|
|
buf.add_request(2, _make_new_req_data("r2"))
|
||
|
|
buf.add_request(3, _make_new_req_data("r3"))
|
||
|
|
batch = SimpleNamespace(idx_mapping_np=[2, 0, 3])
|
||
|
|
assert [g["req_id"] for g in buf.gather(batch)] == ["r2", "r0", "r3"]
|
||
|
|
|
||
|
|
# Slot reuse: req_id mapping migrates; duplicate removal is a no-op.
|
||
|
|
buf.add_request(0, _make_new_req_data("new"))
|
||
|
|
assert buf.req_id_to_index["new"] == 0
|
||
|
|
buf.remove_request(0)
|
||
|
|
buf.remove_request(0)
|
||
|
|
assert "new" not in buf.req_id_to_index
|
||
|
|
assert buf.buffers[0] == {}
|
||
|
|
|
||
|
|
|
||
|
|
def test_update_merge_semantics(monkeypatch):
|
||
|
|
buf = OmniIntermediateBuffer(max_num_reqs=1)
|
||
|
|
buf.add_request(0, _make_new_req_data("r0"))
|
||
|
|
|
||
|
|
buf.update(0, {"hidden": torch.randn(4)})
|
||
|
|
assert buf.buffers[0]["hidden"].device == torch.device("cpu") # default migrates to cpu
|
||
|
|
buf.update(0, {})
|
||
|
|
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]
|
||
|
|
|
||
|
|
# GPU-resident keys keep values (including list items) off the cpu path.
|
||
|
|
tensor = torch.randn(2)
|
||
|
|
monkeypatch.setattr(torch.Tensor, "cpu", lambda *_a, **_k: pytest.fail("GPU-resident value must not move to CPU"))
|
||
|
|
buf.update(0, {"items": [tensor], "kv": torch.randn(4)}, gpu_resident_keys={"items", "kv"})
|
||
|
|
stored = buf.buffers[0]["items"][0]
|
||
|
|
assert stored is not tensor
|
||
|
|
assert torch.equal(stored, tensor)
|
||
|
|
assert "kv" in buf.buffers[0]
|