wehub-resource-sync
724 lines
28 KiB
Python
724 lines
28 KiB
Python
from collections.abc import Sequence
|
|
from types import SimpleNamespace
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.outputs import OmniModelRunnerOutput
|
|
from vllm_omni.worker.gpu_ar_model_runner import (
|
|
ExecuteModelState,
|
|
GPUARModelRunner,
|
|
OmniAsyncGPUModelRunnerOutput,
|
|
)
|
|
from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest
|
|
|
|
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.vllm_config = SimpleNamespace(
|
|
model_config=SimpleNamespace(engine_output_type=engine_output_type),
|
|
)
|
|
runner._request_needs_downstream_stage_payload = lambda rid: rid in downstream_req_ids
|
|
return runner
|
|
|
|
|
|
def test_resolve_pooler_payload_req_ids_audio_terminal_stage_keeps_payload():
|
|
runner = _make_runner(engine_output_type="audio", downstream_req_ids=set())
|
|
|
|
engine_output_type, payload_req_ids = GPUARModelRunner._resolve_pooler_payload_req_ids(runner, ["r1", "r2"])
|
|
|
|
assert engine_output_type == "audio"
|
|
assert payload_req_ids == ["r1", "r2"]
|
|
|
|
|
|
def test_resolve_pooler_payload_req_ids_text_terminal_stage_drops_payload():
|
|
runner = _make_runner(engine_output_type="text", downstream_req_ids=set())
|
|
|
|
engine_output_type, payload_req_ids = GPUARModelRunner._resolve_pooler_payload_req_ids(runner, ["r1", "r2"])
|
|
|
|
assert engine_output_type == "text"
|
|
assert payload_req_ids == []
|
|
|
|
|
|
def test_resolve_pooler_payload_req_ids_downstream_stage_uses_filtered_requests():
|
|
runner = _make_runner(engine_output_type="latent", downstream_req_ids={"r2"})
|
|
|
|
engine_output_type, payload_req_ids = GPUARModelRunner._resolve_pooler_payload_req_ids(runner, ["r1", "r2", "r3"])
|
|
|
|
assert engine_output_type == "latent"
|
|
assert payload_req_ids == ["r2"]
|
|
|
|
|
|
def test_sparse_mm_req_ids_requires_sparse_audio_marker():
|
|
assert GPUARModelRunner._sparse_mm_req_ids({"meta": {"req_id": ["r1"]}}) is None
|
|
assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"]}) is None
|
|
|
|
assert GPUARModelRunner._sparse_mm_req_ids({"meta": {"req_id": ["r1"], "sparse_audio": ["1"]}}) == ["r1"]
|
|
assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"], "meta.sparse_audio": ["1"]}) == ["r1"]
|
|
|
|
|
|
def test_runner_assisted_full_attention_metadata_request_is_opt_in():
|
|
runner = object.__new__(GPUARModelRunner)
|
|
runner.model = object()
|
|
runner.scheduler_config = SimpleNamespace(max_num_seqs=16)
|
|
|
|
request = runner._get_runner_assisted_full_attention_metadata_request(
|
|
req_ids=["r1", "r2"],
|
|
num_reqs=2,
|
|
num_reqs_padded=4,
|
|
num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32),
|
|
num_computed_tokens_cpu=np.array([5, 6], dtype=np.int32),
|
|
max_num_scheduled_tokens=1,
|
|
)
|
|
|
|
assert request is None
|
|
|
|
|
|
def test_runner_assisted_full_attention_metadata_request_and_context_hooks():
|
|
calls = []
|
|
|
|
class Model:
|
|
def get_runner_assisted_full_attention_metadata_request(
|
|
self,
|
|
*,
|
|
req_ids: Sequence[str],
|
|
num_reqs: int,
|
|
num_scheduled_tokens: Sequence[int],
|
|
num_computed_tokens: Sequence[int],
|
|
max_num_scheduled_tokens: int,
|
|
) -> RunnerAssistedFullAttentionMetadataRequest:
|
|
calls.append(
|
|
(
|
|
"request",
|
|
{
|
|
"req_ids": list(req_ids),
|
|
"num_reqs": num_reqs,
|
|
"num_scheduled_tokens": [int(n) for n in num_scheduled_tokens],
|
|
"num_computed_tokens": [int(n) for n in num_computed_tokens],
|
|
"max_num_scheduled_tokens": max_num_scheduled_tokens,
|
|
},
|
|
)
|
|
)
|
|
return RunnerAssistedFullAttentionMetadataRequest(
|
|
num_reqs_padded=12,
|
|
for_cudagraph_capture=True,
|
|
)
|
|
|
|
def set_runner_assisted_full_attention_metadata_context(
|
|
self,
|
|
*,
|
|
enabled: bool,
|
|
num_reqs: int = 0,
|
|
) -> None:
|
|
calls.append(("context", {"enabled": enabled, "num_reqs": num_reqs}))
|
|
|
|
runner = object.__new__(GPUARModelRunner)
|
|
runner.model = Model()
|
|
runner.scheduler_config = SimpleNamespace(max_num_seqs=8)
|
|
|
|
request = runner._get_runner_assisted_full_attention_metadata_request(
|
|
req_ids=["r1", "r2"],
|
|
num_reqs=2,
|
|
num_reqs_padded=4,
|
|
num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32),
|
|
num_computed_tokens_cpu=np.array([5, 6], dtype=np.int32),
|
|
max_num_scheduled_tokens=1,
|
|
)
|
|
context_enabled = runner._set_runner_assisted_full_attention_metadata_context(
|
|
enabled=True,
|
|
num_reqs=2,
|
|
)
|
|
context_disabled = runner._set_runner_assisted_full_attention_metadata_context(enabled=False)
|
|
|
|
assert request == RunnerAssistedFullAttentionMetadataRequest(
|
|
num_reqs_padded=8,
|
|
for_cudagraph_capture=True,
|
|
)
|
|
assert context_enabled
|
|
assert context_disabled
|
|
assert calls == [
|
|
(
|
|
"request",
|
|
{
|
|
"req_ids": ["r1", "r2"],
|
|
"num_reqs": 2,
|
|
"num_scheduled_tokens": [1, 1],
|
|
"num_computed_tokens": [5, 6],
|
|
"max_num_scheduled_tokens": 1,
|
|
},
|
|
),
|
|
("context", {"enabled": True, "num_reqs": 2}),
|
|
("context", {"enabled": False, "num_reqs": 0}),
|
|
]
|
|
|
|
|
|
def test_omni_async_gpu_model_runner_output_builds_lazily_once():
|
|
async_output = object.__new__(OmniAsyncGPUModelRunnerOutput)
|
|
calls = []
|
|
sync_calls = []
|
|
|
|
def builder():
|
|
calls.append("build")
|
|
return OmniModelRunnerOutput(req_ids=["r1"], req_id_to_index={"r1": 0})
|
|
|
|
async_output._model_runner_output = None
|
|
async_output._model_runner_output_builder = builder
|
|
async_output._invalid_req_indices = []
|
|
async_output.sampled_token_ids_cpu = torch.tensor([[7]], dtype=torch.long)
|
|
async_output.async_copy_ready_event = SimpleNamespace(synchronize=lambda: sync_calls.append("sync"))
|
|
async_output._sampled_token_ids = torch.tensor([[7]], dtype=torch.long)
|
|
async_output._logprobs_tensors = None
|
|
async_output._logprobs_tensors_cpu = None
|
|
async_output._routed_experts = None
|
|
async_output._routed_experts_cpu = None
|
|
async_output.vocab_size = 10
|
|
|
|
output = async_output.get_output()
|
|
|
|
assert calls == ["build"]
|
|
assert sync_calls == ["sync"]
|
|
assert async_output._model_runner_output_builder is None
|
|
assert output.req_ids == ["r1"]
|
|
assert output.sampled_token_ids == [[7]]
|
|
|
|
|
|
def test_omni_async_gpu_model_runner_output_reraises_background_exception():
|
|
async_output = object.__new__(OmniAsyncGPUModelRunnerOutput)
|
|
joined = []
|
|
|
|
class FakeThread:
|
|
def join(self):
|
|
joined.append("join")
|
|
|
|
async_output._background_thread = FakeThread()
|
|
async_output._background_exception = RuntimeError("background failed")
|
|
|
|
with pytest.raises(RuntimeError, match="background failed"):
|
|
async_output.get_output()
|
|
|
|
assert joined == ["join"]
|
|
assert async_output._background_thread is None
|
|
|
|
|
|
def _make_async_output_runner(engine_output_type: str = "audio"):
|
|
runner = object.__new__(GPUARModelRunner)
|
|
model_config = SimpleNamespace(
|
|
engine_output_type=engine_output_type,
|
|
async_chunk=True,
|
|
enable_return_routed_experts=False,
|
|
)
|
|
runner.vllm_config = SimpleNamespace(model_config=model_config)
|
|
runner.model_config = model_config
|
|
runner._async_chunk = True
|
|
runner.omni_prefix_cache = None
|
|
runner.requests = {"r1": object(), "r2": object()}
|
|
runner.supports_mm_inputs = False
|
|
runner.routed_experts_initialized = False
|
|
runner.model = SimpleNamespace(has_postprocess=False)
|
|
runner.model_intermediate_buffer = {}
|
|
runner.input_batch = SimpleNamespace(
|
|
req_ids=["mutated"],
|
|
req_id_to_index={"mutated": 0},
|
|
)
|
|
return runner
|
|
|
|
|
|
def test_build_omni_output_uses_snapshots_and_connector_after_accumulation(monkeypatch):
|
|
runner = _make_async_output_runner()
|
|
events = []
|
|
|
|
monkeypatch.setattr(
|
|
GPUARModelRunner,
|
|
"_resolve_pooler_payload_req_ids",
|
|
lambda self, req_ids: ("audio", req_ids),
|
|
)
|
|
monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: True)
|
|
monkeypatch.setattr(
|
|
GPUARModelRunner,
|
|
"accumulate_full_payload_output",
|
|
lambda self, rid, payload, request: events.append(f"accumulate:{rid}"),
|
|
)
|
|
monkeypatch.setattr(
|
|
GPUARModelRunner,
|
|
"get_omni_connector_output",
|
|
lambda self: events.append("connector") or "connector-output",
|
|
)
|
|
|
|
output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
|
|
runner,
|
|
scheduler_output=SimpleNamespace(
|
|
total_num_scheduled_tokens=3,
|
|
num_scheduled_tokens={"r1": 1, "r2": 2},
|
|
),
|
|
hidden_states=torch.tensor([[1.0], [2.0], [3.0]]),
|
|
staged_hidden_states_cpu=None,
|
|
multimodal_outputs={"foo": torch.tensor([10.0, 20.0, 30.0])},
|
|
req_ids_output_copy=["r1", "r2"],
|
|
req_id_to_index_output_copy={"r1": 0, "r2": 1},
|
|
valid_sampled_token_ids=[[101], [102]],
|
|
logprobs_lists=None,
|
|
prompt_logprobs_dict={},
|
|
num_nans_in_logits=None,
|
|
kv_connector_output=None,
|
|
ec_connector_output=None,
|
|
cudagraph_stats=None,
|
|
kv_extracted_req_ids=["r2"],
|
|
num_scheduled_tokens_np=torch.tensor([1, 2], dtype=torch.int32).numpy(),
|
|
query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
|
|
)
|
|
|
|
assert output.req_ids == ["r1", "r2"]
|
|
assert output.inter_stage_outputs is not None
|
|
assert torch.equal(output.inter_stage_outputs[0]["hidden"], torch.tensor([[1.0]]))
|
|
assert torch.equal(output.inter_stage_outputs[1]["hidden"], torch.tensor([[2.0], [3.0]]))
|
|
assert output.multimodal_outputs is None
|
|
assert output.kv_extracted_req_ids == ["r2"]
|
|
assert output.omni_connector_output == "connector-output"
|
|
assert events == ["accumulate:r1", "accumulate:r2", "connector"]
|
|
|
|
|
|
def test_build_omni_output_copies_hidden_for_partial_downstream_batch(monkeypatch):
|
|
runner = _make_async_output_runner(engine_output_type="latent")
|
|
|
|
monkeypatch.setattr(
|
|
GPUARModelRunner,
|
|
"_resolve_pooler_payload_req_ids",
|
|
lambda self, req_ids: ("latent", ["r2"]),
|
|
)
|
|
monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
|
|
monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
|
|
monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
|
|
|
|
output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
|
|
runner,
|
|
scheduler_output=SimpleNamespace(
|
|
total_num_scheduled_tokens=6,
|
|
num_scheduled_tokens={"r1": 1, "r2": 2, "r3": 3},
|
|
),
|
|
hidden_states=torch.tensor([[1.0], [2.0], [3.0], [4.0], [5.0], [6.0]]),
|
|
staged_hidden_states_cpu=None,
|
|
multimodal_outputs={},
|
|
req_ids_output_copy=["r1", "r2", "r3"],
|
|
req_id_to_index_output_copy={"r1": 0, "r2": 1, "r3": 2},
|
|
valid_sampled_token_ids=[[], [], []],
|
|
logprobs_lists=None,
|
|
prompt_logprobs_dict={},
|
|
num_nans_in_logits=None,
|
|
kv_connector_output=None,
|
|
ec_connector_output=None,
|
|
cudagraph_stats=None,
|
|
kv_extracted_req_ids=None,
|
|
num_scheduled_tokens_np=np.array([1, 2, 3], dtype=np.int32),
|
|
query_start_loc_cpu=torch.tensor([0, 1, 3], dtype=torch.long),
|
|
)
|
|
|
|
assert output.inter_stage_outputs is not None
|
|
assert output.inter_stage_outputs[0] is None
|
|
assert torch.equal(output.inter_stage_outputs[1]["hidden"], torch.tensor([[2.0], [3.0]]))
|
|
assert output.inter_stage_outputs[2] is None
|
|
assert output.multimodal_outputs is None
|
|
|
|
|
|
def test_process_additional_information_uses_snapshot_request_order(monkeypatch):
|
|
runner = _make_async_output_runner()
|
|
seen = []
|
|
|
|
class PostprocessModel:
|
|
has_postprocess = True
|
|
|
|
def postprocess(self, hidden_states, **kwargs):
|
|
seen.append(hidden_states.clone())
|
|
return {}
|
|
|
|
runner.model = PostprocessModel()
|
|
runner.model_intermediate_buffer = {"r1": {}, "r2": {}}
|
|
|
|
monkeypatch.setattr(
|
|
GPUARModelRunner,
|
|
"_resolve_pooler_payload_req_ids",
|
|
lambda self, req_ids: ("audio", req_ids),
|
|
)
|
|
monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
|
|
monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
|
|
monkeypatch.setattr(GPUARModelRunner, "_update_intermediate_buffer", lambda *args, **kwargs: None)
|
|
|
|
GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
|
|
runner,
|
|
scheduler_output=SimpleNamespace(
|
|
total_num_scheduled_tokens=3,
|
|
num_scheduled_tokens={"r1": 1, "r2": 2},
|
|
),
|
|
hidden_states=torch.tensor([[1.0], [2.0], [3.0]]),
|
|
staged_hidden_states_cpu=None,
|
|
multimodal_outputs={},
|
|
req_ids_output_copy=["r1", "r2"],
|
|
req_id_to_index_output_copy={"r1": 0, "r2": 1},
|
|
valid_sampled_token_ids=[[], []],
|
|
logprobs_lists=None,
|
|
prompt_logprobs_dict={},
|
|
num_nans_in_logits=None,
|
|
kv_connector_output=None,
|
|
ec_connector_output=None,
|
|
cudagraph_stats=None,
|
|
kv_extracted_req_ids=None,
|
|
num_scheduled_tokens_np=torch.tensor([1, 2], dtype=torch.int32).numpy(),
|
|
query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
|
|
)
|
|
|
|
assert len(seen) == 2
|
|
assert torch.equal(seen[0], torch.tensor([[1.0]]))
|
|
assert torch.equal(seen[1], torch.tensor([[2.0], [3.0]]))
|
|
|
|
|
|
def test_async_omni_output_guard_requires_safe_conditions():
|
|
runner = _make_async_output_runner()
|
|
runner.use_async_scheduling = True
|
|
runner.speculative_config = None
|
|
runner.model.use_async_omni_output = True
|
|
|
|
assert GPUARModelRunner._should_use_async_omni_output(runner)
|
|
|
|
runner.omni_prefix_cache = object()
|
|
assert not 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)
|
|
|
|
runner.model.eager_omni_postprocess_before_async_output = True
|
|
assert GPUARModelRunner._should_use_async_omni_output(runner)
|
|
|
|
|
|
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
|
|
|
|
monkeypatch.setattr(
|
|
GPUARModelRunner,
|
|
"_resolve_pooler_payload_req_ids",
|
|
lambda self, req_ids: ("latent", req_ids),
|
|
)
|
|
monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
|
|
monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
|
|
monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
|
|
|
|
output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
|
|
runner,
|
|
scheduler_output=SimpleNamespace(
|
|
total_num_scheduled_tokens=2,
|
|
num_scheduled_tokens={"r1": 2},
|
|
),
|
|
hidden_states=torch.tensor([[1.0], [2.0]]),
|
|
staged_hidden_states_cpu=None,
|
|
multimodal_outputs={"codes": {"audio": torch.tensor([[7, 8], [9, 10]], dtype=torch.long)}},
|
|
req_ids_output_copy=["r1"],
|
|
req_id_to_index_output_copy={"r1": 0},
|
|
valid_sampled_token_ids=[[101]],
|
|
logprobs_lists=None,
|
|
prompt_logprobs_dict={},
|
|
num_nans_in_logits=None,
|
|
kv_connector_output=None,
|
|
ec_connector_output=None,
|
|
cudagraph_stats=None,
|
|
kv_extracted_req_ids=None,
|
|
num_scheduled_tokens_np=np.array([2], dtype=np.int32),
|
|
query_start_loc_cpu=torch.tensor([0], dtype=torch.long),
|
|
)
|
|
|
|
assert output.inter_stage_outputs is not None
|
|
assert len(output.inter_stage_outputs) == 1
|
|
assert "hidden" not in output.inter_stage_outputs[0]
|
|
assert torch.equal(output.inter_stage_outputs[0]["codes.audio"], torch.tensor([[7, 8], [9, 10]], dtype=torch.long))
|
|
assert output.multimodal_outputs is None
|
|
|
|
|
|
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._async_chunk = False
|
|
runner.requests = {"r1": object(), "r2": object(), "r3": object()}
|
|
|
|
monkeypatch.setattr(
|
|
GPUARModelRunner,
|
|
"_resolve_pooler_payload_req_ids",
|
|
lambda self, req_ids: ("latent", req_ids),
|
|
)
|
|
monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
|
|
monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
|
|
monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
|
|
|
|
codes = torch.arange(48, dtype=torch.long).reshape(3, 16)
|
|
output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
|
|
runner,
|
|
scheduler_output=SimpleNamespace(
|
|
total_num_scheduled_tokens=3,
|
|
num_scheduled_tokens={"r1": 1, "r2": 1, "r3": 1},
|
|
),
|
|
hidden_states=torch.tensor([[1.0]]),
|
|
staged_hidden_states_cpu=None,
|
|
multimodal_outputs={"codes": {"audio": codes}},
|
|
req_ids_output_copy=["r1", "r2", "r3"],
|
|
req_id_to_index_output_copy={"r1": 0, "r2": 1, "r3": 2},
|
|
valid_sampled_token_ids=[[101], [102], [103]],
|
|
logprobs_lists=None,
|
|
prompt_logprobs_dict={},
|
|
num_nans_in_logits=None,
|
|
kv_connector_output=None,
|
|
ec_connector_output=None,
|
|
cudagraph_stats=None,
|
|
kv_extracted_req_ids=None,
|
|
num_scheduled_tokens_np=np.array([1, 1, 1], dtype=np.int32),
|
|
query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.long),
|
|
)
|
|
|
|
assert output.inter_stage_outputs is not None
|
|
assert torch.equal(output.inter_stage_outputs[0]["codes.audio"], codes[0:1])
|
|
assert torch.equal(output.inter_stage_outputs[1]["codes.audio"], codes[1:2])
|
|
assert torch.equal(output.inter_stage_outputs[2]["codes.audio"], codes[2:3])
|
|
|
|
|
|
def test_build_omni_output_splits_mm_by_hidden_len_when_scheduled_is_padded(monkeypatch):
|
|
"""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._async_chunk = False
|
|
runner.requests = {"r1": object(), "r2": object(), "r3": object()}
|
|
|
|
monkeypatch.setattr(
|
|
GPUARModelRunner,
|
|
"_resolve_pooler_payload_req_ids",
|
|
lambda self, req_ids: ("latent", req_ids),
|
|
)
|
|
monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
|
|
monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
|
|
monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
|
|
|
|
layers = torch.arange(96, dtype=torch.long).reshape(3, 32)
|
|
output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
|
|
runner,
|
|
scheduler_output=SimpleNamespace(
|
|
total_num_scheduled_tokens=5,
|
|
num_scheduled_tokens={"r1": 1, "r2": 1, "r3": 1},
|
|
),
|
|
hidden_states=torch.randn(3, 4),
|
|
staged_hidden_states_cpu=None,
|
|
multimodal_outputs={"hidden_states": {"layers": {0: layers}}},
|
|
req_ids_output_copy=["r1", "r2", "r3"],
|
|
req_id_to_index_output_copy={"r1": 0, "r2": 1, "r3": 2},
|
|
valid_sampled_token_ids=[[101], [102], [103]],
|
|
logprobs_lists=None,
|
|
prompt_logprobs_dict={},
|
|
num_nans_in_logits=None,
|
|
kv_connector_output=None,
|
|
ec_connector_output=None,
|
|
cudagraph_stats=None,
|
|
kv_extracted_req_ids=None,
|
|
num_scheduled_tokens_np=np.array([1, 1, 1], dtype=np.int32),
|
|
query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.long),
|
|
)
|
|
|
|
assert output.inter_stage_outputs is not None
|
|
assert torch.equal(output.inter_stage_outputs[0]["hidden_states.layer_0"], layers[0:1])
|
|
assert torch.equal(output.inter_stage_outputs[1]["hidden_states.layer_0"], layers[1:2])
|
|
assert torch.equal(output.inter_stage_outputs[2]["hidden_states.layer_0"], layers[2:3])
|
|
|
|
|
|
def test_async_snapshot_payload_omits_hidden_when_model_opts_out():
|
|
runner = _make_async_output_runner()
|
|
runner.model.omni_pooler_payload_include_hidden = False
|
|
|
|
payload = GPUARModelRunner._build_omni_async_snapshot_payload(
|
|
runner,
|
|
hidden_states=torch.tensor([[1.0], [2.0]]),
|
|
staged_hidden_states_cpu=torch.tensor([[3.0]]),
|
|
multimodal_outputs={"codes": {"audio": torch.tensor([[1]], dtype=torch.long)}},
|
|
)
|
|
|
|
assert set(payload.keys()) == {"multimodal_outputs"}
|
|
assert payload["multimodal_outputs"]["codes"]["audio"].tolist() == [[1]]
|
|
|
|
|
|
def test_runner_assisted_full_attention_metadata_refresh_pads_buffers():
|
|
class QueryStartLoc:
|
|
def __init__(self):
|
|
self.np = np.full(5, -1, dtype=np.int32)
|
|
self.copied = False
|
|
|
|
def copy_to_gpu(self):
|
|
self.copied = True
|
|
|
|
class BlockTable:
|
|
def __init__(self):
|
|
self.commits = []
|
|
|
|
def commit_block_table(self, num_reqs_padded):
|
|
self.commits.append(num_reqs_padded)
|
|
|
|
runner = object.__new__(GPUARModelRunner)
|
|
block_table = BlockTable()
|
|
runner.input_batch = SimpleNamespace(
|
|
num_computed_tokens_cpu=torch.tensor([10, 20, 99, 99], dtype=torch.int32),
|
|
block_table=block_table,
|
|
)
|
|
runner.optimistic_seq_lens_cpu = torch.zeros(4, dtype=torch.int32)
|
|
runner.seq_lens = torch.empty(4, dtype=torch.int32)
|
|
runner.query_pos = SimpleNamespace(np=np.empty(4, dtype=np.int32))
|
|
runner.query_start_loc = QueryStartLoc()
|
|
runner._get_cumsum_and_arange = lambda scheduled, _query_pos: np.cumsum(scheduled, dtype=np.int32)
|
|
|
|
runner._refresh_runner_assisted_full_attention_metadata_buffers(
|
|
num_reqs=2,
|
|
num_reqs_padded=4,
|
|
num_scheduled_tokens_np=np.array([1, 2], dtype=np.int32),
|
|
)
|
|
|
|
assert runner.optimistic_seq_lens_cpu.tolist() == [11, 22, 0, 0]
|
|
assert runner.seq_lens.tolist() == [11, 22, 0, 0]
|
|
assert runner.query_start_loc.np.tolist() == [0, 1, 3, 3, 3]
|
|
assert runner.query_start_loc.copied
|
|
assert block_table.commits == [4]
|
|
|
|
|
|
@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.execute_model_state = ExecuteModelState(
|
|
SimpleNamespace(
|
|
total_num_scheduled_tokens=3,
|
|
num_scheduled_tokens={"r1": 1, "r2": 2},
|
|
),
|
|
None,
|
|
None,
|
|
None,
|
|
torch.zeros((3, 2), dtype=torch.float32),
|
|
torch.tensor([[1.0, 10.0], [2.0, 20.0], [3.0, 30.0]]),
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
{},
|
|
None,
|
|
)
|
|
runner.kv_connector_output = None
|
|
runner.input_batch = SimpleNamespace(
|
|
req_ids=["r1", "r2"],
|
|
req_id_to_index={"r1": 0, "r2": 1},
|
|
sampling_metadata=SimpleNamespace(no_penalties=True),
|
|
vocab_size=10,
|
|
num_tokens_no_spec=None,
|
|
)
|
|
query_start_loc = torch.tensor([0, 1], dtype=torch.long)
|
|
if query_start_loc_attr == "method":
|
|
runner.query_start_loc = query_start_loc
|
|
else:
|
|
runner.query_start_loc = SimpleNamespace(cpu=query_start_loc)
|
|
runner.omni_prefix_cache = object()
|
|
runner.speculative_config = None
|
|
runner.routed_experts_initialized = False
|
|
runner.requests = {}
|
|
runner.supports_mm_inputs = False
|
|
runner.use_async_scheduling = False
|
|
runner._omni_num_scheduled_tokens_np = None
|
|
runner.vllm_config = SimpleNamespace(
|
|
model_config=SimpleNamespace(engine_output_type="audio"),
|
|
)
|
|
runner._async_chunk = False
|
|
|
|
monkeypatch.setattr(
|
|
GPUARModelRunner, "_sample", lambda self, logits, spec_decode_metadata: SimpleNamespace(sampled_token_ids=[])
|
|
)
|
|
monkeypatch.setattr(GPUARModelRunner, "_update_states_after_model_execute", lambda *args, **kwargs: None)
|
|
monkeypatch.setattr(
|
|
GPUARModelRunner,
|
|
"_bookkeeping_sync",
|
|
lambda *args, **kwargs: (
|
|
0,
|
|
None,
|
|
[],
|
|
None,
|
|
["r1", "r2"],
|
|
{"r1": 0, "r2": 1},
|
|
[],
|
|
),
|
|
)
|
|
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),
|
|
)
|
|
monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
|
|
monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
|
|
monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
|
|
|
|
output = GPUARModelRunner.sample_tokens(runner, grammar_output=None)
|
|
|
|
# Non-async-chunk now ships the full payload to the next stage, so
|
|
# inter_stage_outputs mirrors multimodal_outputs (PR #4792).
|
|
assert output.inter_stage_outputs is not None
|
|
assert output.multimodal_outputs is not None
|
|
assert torch.equal(output.inter_stage_outputs[0]["hidden"], output.multimodal_outputs[0]["hidden"])
|
|
assert torch.equal(output.inter_stage_outputs[1]["hidden"], output.multimodal_outputs[1]["hidden"])
|
|
assert torch.equal(output.multimodal_outputs[0]["hidden"], torch.tensor([[1.0, 10.0]]))
|
|
assert torch.equal(
|
|
output.multimodal_outputs[1]["hidden"],
|
|
torch.tensor([[2.0, 20.0], [3.0, 30.0]]),
|
|
)
|
|
|
|
|
|
def test_build_omni_output_falls_back_to_mm_cpu_without_prefix_merge(monkeypatch):
|
|
"""Tail-only prefix-cache models still need per-step mm passthrough (e.g. codes.audio)."""
|
|
runner = _make_async_output_runner(engine_output_type="latent")
|
|
runner.omni_prefix_cache = object()
|
|
|
|
monkeypatch.setattr(
|
|
GPUARModelRunner,
|
|
"_resolve_pooler_payload_req_ids",
|
|
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",
|
|
lambda *args, **kwargs: (None, None, None),
|
|
)
|
|
monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
|
|
monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
|
|
monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
|
|
|
|
codes = torch.tensor([[11.0, 12.0], [21.0, 22.0]], dtype=torch.float32)
|
|
output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
|
|
runner,
|
|
scheduler_output=SimpleNamespace(
|
|
total_num_scheduled_tokens=2,
|
|
num_scheduled_tokens={"r1": 1, "r2": 1},
|
|
),
|
|
hidden_states=torch.tensor([[1.0], [2.0]]),
|
|
staged_hidden_states_cpu=None,
|
|
multimodal_outputs={"codes.audio": codes},
|
|
req_ids_output_copy=["r1", "r2"],
|
|
req_id_to_index_output_copy={"r1": 0, "r2": 1},
|
|
valid_sampled_token_ids=[[], []],
|
|
logprobs_lists=None,
|
|
prompt_logprobs_dict={},
|
|
num_nans_in_logits=None,
|
|
kv_connector_output=None,
|
|
ec_connector_output=None,
|
|
cudagraph_stats=None,
|
|
kv_extracted_req_ids=None,
|
|
num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32),
|
|
query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
|
|
)
|
|
|
|
assert torch.equal(output.inter_stage_outputs[0]["codes.audio"], codes[0:1])
|
|
assert torch.equal(output.inter_stage_outputs[1]["codes.audio"], codes[1:2])
|
|
assert output.multimodal_outputs is None
|