From f24a6165ff1612b912962e7e169ebfdd1d2a080a Mon Sep 17 00:00:00 2001 From: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com> Date: Tue, 4 Aug 2026 20:03:57 -0700 Subject: [PATCH] [Test] Prefix-cache passthrough test coverage prerequisite for W2 (#4855) (#5310) Signed-off-by: Shenglei Fu Co-authored-by: Claude --- tests/core/test_prefix_cache.py | 316 ++++++++++++++++++++ tests/core/test_prefix_cache_async_write.py | 119 ++++++++ tests/utils/test_mm_outputs.py | 147 +++++++++ tests/worker/test_gpu_ar_model_runner.py | 136 +++++++++ 4 files changed, 718 insertions(+) create mode 100644 tests/core/test_prefix_cache_async_write.py create mode 100644 tests/utils/test_mm_outputs.py diff --git a/tests/core/test_prefix_cache.py b/tests/core/test_prefix_cache.py index db7610496..4f91a95b7 100644 --- a/tests/core/test_prefix_cache.py +++ b/tests/core/test_prefix_cache.py @@ -1,3 +1,6 @@ +import math +from typing import NamedTuple + import pytest import torch @@ -317,6 +320,103 @@ def test_deferred_multimodal_cache_write_commits_on_completion(): assert torch.equal(rows[8:16], expected) +def test_discard_deferred_mm_outputs_drops_pending_chunks(): + """discard_deferred_mm_outputs is the abort path: a request that + leaves without a cache commit must not leak its staged chunks into + the cache if commit_deferred_mm_outputs is later called for it.""" + cache = get_omni_pcache() + mm_key = "codes.audio" + block_table = torch.tensor( + [ + [2, 3, 4, 5], + [6, 7, 8, 9], + ], + dtype=torch.long, + ) + input_batch = MockInputBatch( + num_computed_tokens_cpu=torch.tensor([4, 0], dtype=torch.long), + block_table=block_table, + ) + + cache.update_omni_tensor_prefix_cache( + hidden_states=None, + multimodal_outputs={mm_key: torch.arange(8, dtype=torch.long).reshape(4, 2)}, + num_tokens_unpadded=4, + slot_mapping=torch.arange(8, 12), + skip_mm_cache_keys={mm_key}, + ) + + chunk = torch.arange(8, dtype=torch.long).reshape(4, 2) + cache.stage_deferred_mm_outputs( + query_start_loc=torch.tensor([0, 4], dtype=torch.long), + input_batch=input_batch, + multimodal_outputs={mm_key: chunk}, + num_scheduled_tokens={"req1": 4, "req2": 0}, + deferred_mm_cache_keys={mm_key}, + ) + assert "req1" in cache._deferred_mm_outputs + + cache.discard_deferred_mm_outputs("req1") + assert "req1" not in cache._deferred_mm_outputs + + # A later commit call for the now-discarded req_id must be a no-op, + # not a resurrection of the dropped chunk. + cache.commit_deferred_mm_outputs({"req1"}, input_batch) + rows = cache.mm_outputs_cache[mm_key].view(-1, 2) + assert torch.all(rows[8:12] == 0) + + +def test_deferred_mm_outputs_request_id_reused_after_discard(): + """A request ID reused after discard_deferred_mm_outputs must start + clean: the new occupant's committed data must not be mixed with the + previous (discarded) occupant's staged chunks.""" + cache = get_omni_pcache() + mm_key = "codes.audio" + block_table = torch.tensor( + [ + [2, 3, 4, 5], + [6, 7, 8, 9], + ], + dtype=torch.long, + ) + input_batch = MockInputBatch( + num_computed_tokens_cpu=torch.tensor([4, 0], dtype=torch.long), + block_table=block_table, + ) + + cache.update_omni_tensor_prefix_cache( + hidden_states=None, + multimodal_outputs={mm_key: torch.arange(8, dtype=torch.long).reshape(4, 2)}, + num_tokens_unpadded=4, + slot_mapping=torch.arange(8, 12), + skip_mm_cache_keys={mm_key}, + ) + + stale_chunk = torch.full((4, 2), -1, dtype=torch.long) + cache.stage_deferred_mm_outputs( + query_start_loc=torch.tensor([0, 4], dtype=torch.long), + input_batch=input_batch, + multimodal_outputs={mm_key: stale_chunk}, + num_scheduled_tokens={"req1": 4, "req2": 0}, + deferred_mm_cache_keys={mm_key}, + ) + cache.discard_deferred_mm_outputs("req1") + + # "req1" is reused by a brand-new request with different data. + fresh_chunk = torch.arange(8, dtype=torch.long).reshape(4, 2) + cache.stage_deferred_mm_outputs( + query_start_loc=torch.tensor([0, 4], dtype=torch.long), + input_batch=input_batch, + multimodal_outputs={mm_key: fresh_chunk}, + num_scheduled_tokens={"req1": 4, "req2": 0}, + deferred_mm_cache_keys={mm_key}, + ) + cache.commit_deferred_mm_outputs({"req1"}, input_batch) + + rows = cache.mm_outputs_cache[mm_key].view(-1, 2) + assert torch.equal(rows[8:12], fresh_chunk) + + def test_deferred_multimodal_cache_can_be_merged_on_full_block_hit(): cache = get_omni_pcache() mm_key = "codes.audio" @@ -599,3 +699,219 @@ def test_get_merged_multimodal_outputs(feat_dims, num_tokens_padded, mocker): assert req2_merged_mm_outputs.shape == torch.Size([num_new_toks_req2, curr_feat_dim]) assert torch.all(req2_merged_mm_outputs == req2_new_mm_outputs) + + +@pytest.mark.parametrize( + "passthrough_shape", + [(5,), (5, 3)], + ids=["1d_token_metadata", "2d_feature_tensor"], +) +def test_get_merged_multimodal_outputs_slices_passthrough_tensor_per_request(passthrough_shape, mocker): + """A passthrough key that is itself a real tensor (not the string/list + passthrough already covered above) must be sliced per request through + to_payload_element's tensor branch, under the same mixed cache + hit/miss + uneven scheduled-length batch as test_get_merged_multimodal_outputs. + + Covers both a 1D token-aligned metadata tensor (e.g. ref_code_len, + codec_streaming -- shape (total_tokens,)) and a 2D per-token feature + tensor (shape (total_tokens, feat_dim)), since W2's batching could + easily special-case one and silently mishandle the other. + """ + feat_dims = {"foo": 100, "bar": 100} + cache = get_omni_pcache_with_mm_tensors(feat_dims, seq_len=DEFAULT_SEQ_LEN) + + orig_num_tokens_unpadded = 8 + slot_offset = 8 + orig_slot_mapping = torch.arange(slot_offset, slot_offset + orig_num_tokens_unpadded) + feature_dims = {key: val.shape[-1] for key, val in cache.mm_outputs_cache.items()} + orig_mm_outputs = { + key: torch.rand((orig_num_tokens_unpadded, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys + } + + cache.update_omni_tensor_prefix_cache( + hidden_states=None, + multimodal_outputs=orig_mm_outputs, + num_tokens_unpadded=orig_num_tokens_unpadded, + slot_mapping=orig_slot_mapping, + num_tokens_padded=None, + ) + + num_new_toks_req1 = 3 + num_new_toks_req2 = 2 + cache.add_prefix_cached_new_req_id("req1") + + num_scheduled_tokens = { + "req1": num_new_toks_req1, + "req2": num_new_toks_req2, + } + + new_mm_outputs = {} + for mm_key in cache.mm_cache_keys: + new_mm_outputs[mm_key] = torch.rand( + (num_new_toks_req1 + num_new_toks_req2, feature_dims[mm_key]), + dtype=DTYPE, + ) + total_scheduled_tokens = num_new_toks_req1 + num_new_toks_req2 + assert passthrough_shape[0] == total_scheduled_tokens + passthrough_tensor = torch.arange(int(torch.prod(torch.tensor(passthrough_shape))), dtype=DTYPE).reshape( + passthrough_shape + ) + new_mm_outputs["passthrough_tensor"] = passthrough_tensor + + input_batch = MockInputBatch(num_computed_tokens_cpu=torch.Tensor([orig_num_tokens_unpadded, 0])) + + mocker.patch( + "vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids", + new=fake_get_cached_block_ids, + ) + merged_mm_outputs = cache.get_merged_multimodal_states( + query_start_loc=[0, num_new_toks_req1], + input_batch=input_batch, + multimodal_outputs=new_mm_outputs, + num_scheduled_tokens=num_scheduled_tokens, + ) + + assert "passthrough_tensor" not in cache.mm_cache_keys + + passthrough_out = merged_mm_outputs["passthrough_tensor"] + # req1 (cache hit) gets tokens [0:3), req2 (cache miss) gets tokens [3:5) -- + # a passthrough tensor is sliced by scheduled position, not merged with + # any cached history (unlike the mm_cache_keys tensors above). + assert torch.equal(passthrough_out["req1"], passthrough_tensor[0:num_new_toks_req1]) + assert torch.equal(passthrough_out["req2"], passthrough_tensor[num_new_toks_req1:total_scheduled_tokens]) + + +class MergeHarness(NamedTuple): + """A primed cache plus the arguments get_merged_multimodal_states needs.""" + + cache: OmniTensorPrefixCache + input_batch: MockInputBatch + num_scheduled_tokens: dict[str, int] + query_start_loc: list[int] + total_scheduled_tokens: int + new_mm_outputs: dict[str, object] + + +def build_mixed_hit_miss_harness(mocker, num_new_toks_req1: int = 3, num_new_toks_req2: int = 2) -> MergeHarness: + """Two requests over one cache key: req1 hits 8 cached tokens, req2 misses, + and their scheduled lengths differ -- the same batch shape as + test_get_merged_multimodal_outputs, which is what makes a per-request + boundary error observable. Callers add their passthrough keys to + ``new_mm_outputs`` before calling get_merged_multimodal_states. + """ + cache = get_omni_pcache_with_mm_tensors({"foo": 100}, seq_len=DEFAULT_SEQ_LEN) + + orig_num_tokens_unpadded = 8 + feature_dims = {key: val.shape[-1] for key, val in cache.mm_outputs_cache.items()} + cache.update_omni_tensor_prefix_cache( + hidden_states=None, + multimodal_outputs={ + key: torch.rand((orig_num_tokens_unpadded, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys + }, + num_tokens_unpadded=orig_num_tokens_unpadded, + slot_mapping=torch.arange(8, 8 + orig_num_tokens_unpadded), + num_tokens_padded=None, + ) + cache.add_prefix_cached_new_req_id("req1") + + total_scheduled_tokens = num_new_toks_req1 + num_new_toks_req2 + mocker.patch( + "vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids", + new=fake_get_cached_block_ids, + ) + return MergeHarness( + cache=cache, + input_batch=MockInputBatch(num_computed_tokens_cpu=torch.Tensor([orig_num_tokens_unpadded, 0])), + num_scheduled_tokens={"req1": num_new_toks_req1, "req2": num_new_toks_req2}, + query_start_loc=[0, num_new_toks_req1], + total_scheduled_tokens=total_scheduled_tokens, + new_mm_outputs={ + key: torch.rand((total_scheduled_tokens, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys + }, + ) + + +def test_get_merged_multimodal_outputs_passes_per_request_tensor_list_through(mocker): + """The passthrough shape qwen3-tts actually emits: a nested ``codes``/``meta`` + dict whose ``codes.ref`` holds one tensor per request + (qwen3_tts_talker.make_omni_output), carried through to_payload_element's + pass_lists_through branch beside 1D int metadata that IS sliced per request. + + Each request must receive the list *whole*. Per-request indexing happens + later, in the runner's _unwrap_lists; a W2 that split the list here would + make that second index select a different request's reference codes. The + existing list coverage passes a list of strings, which cannot catch a + tensor-specific regression, and no other test covers an int dtype. + """ + harness = build_mixed_hit_miss_harness(mocker) + + ref_codes = [torch.tensor([11, 12], dtype=torch.long), torch.tensor([21], dtype=torch.long)] + ref_code_len = torch.arange(harness.total_scheduled_tokens, dtype=torch.int32) + harness.new_mm_outputs["codes"] = {"ref": ref_codes} + harness.new_mm_outputs["meta"] = {"ref_code_len": ref_code_len} + + merged = harness.cache.get_merged_multimodal_states( + query_start_loc=harness.query_start_loc, + input_batch=harness.input_batch, + multimodal_outputs=harness.new_mm_outputs, + num_scheduled_tokens=harness.num_scheduled_tokens, + ) + + assert "codes" not in harness.cache.mm_cache_keys + for req_id in ("req1", "req2"): + req_refs = merged["codes"][req_id]["ref"] + assert len(req_refs) == len(ref_codes) + for got, want in zip(req_refs, ref_codes): + assert torch.equal(got, want) + assert got.dtype == torch.long + + # Per-request clones: the branch exists so one request's downstream mutation + # cannot reach another request's payload or the source list. + merged["codes"]["req1"]["ref"][0] += 100 + assert torch.equal(merged["codes"]["req2"]["ref"][0], torch.tensor([11, 12], dtype=torch.long)) + assert torch.equal(ref_codes[0], torch.tensor([11, 12], dtype=torch.long)) + + # Token-aligned 1D metadata is still split per request, unlike the list. + req1_len = harness.num_scheduled_tokens["req1"] + assert torch.equal(merged["meta"]["req1"]["ref_code_len"], ref_code_len[:req1_len]) + assert torch.equal(merged["meta"]["req2"]["ref_code_len"], ref_code_len[req1_len:]) + assert merged["meta"]["req1"]["ref_code_len"].dtype == torch.int32 + + +@pytest.mark.parametrize( + "passthrough_shape", + [(3, 4), (1, 4), (7,)], + ids=["shorter_2d", "single_row_2d", "longer_1d"], +) +def test_get_merged_multimodal_outputs_broadcasts_non_token_aligned_passthrough(passthrough_shape, mocker): + """A passthrough tensor whose first dimension is not the scheduled-token + count is request-invariant: to_payload_element hands every request the whole + tensor, cloned, rather than slicing it. + + W2 batching most plausibly breaks by assuming every passthrough tensor is + token-aligned; that assumption would slice these and give each request a + truncated view, or an empty one once start exceeds the first dimension. + """ + harness = build_mixed_hit_miss_harness(mocker) + assert passthrough_shape[0] != harness.total_scheduled_tokens + + passthrough = torch.arange(math.prod(passthrough_shape), dtype=DTYPE).reshape(passthrough_shape) + expected = passthrough.clone() + harness.new_mm_outputs["speaker_embedding"] = passthrough + + merged = harness.cache.get_merged_multimodal_states( + query_start_loc=harness.query_start_loc, + input_batch=harness.input_batch, + multimodal_outputs=harness.new_mm_outputs, + num_scheduled_tokens=harness.num_scheduled_tokens, + ) + + assert "speaker_embedding" not in harness.cache.mm_cache_keys + out = merged["speaker_embedding"] + for req_id in ("req1", "req2"): + assert out[req_id].shape == passthrough_shape + assert torch.equal(out[req_id], expected) + + out["req1"] += 100 + assert torch.equal(out["req2"], expected) + assert torch.equal(passthrough, expected) diff --git a/tests/core/test_prefix_cache_async_write.py b/tests/core/test_prefix_cache_async_write.py new file mode 100644 index 000000000..fa1980227 --- /dev/null +++ b/tests/core/test_prefix_cache_async_write.py @@ -0,0 +1,119 @@ +import pytest +import torch + +from vllm_omni.core.prefix_cache import OmniTensorPrefixCache + +# `cpu` is the lane selector, not a hardware claim: the only Buildkite steps that +# reach tests/ select `core_model and cpu`, and they run on GPU hardware (l4_1). +# Marking this `cuda` collects it nowhere. Real hardware need is the skipif below. +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cpu, + pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device"), +] + +NUM_BLOCKS = 10 +BLOCK_SIZE = 4 +HIDDEN_SIZE = 2 +DTYPE = torch.float32 + + +def get_omni_pcache() -> OmniTensorPrefixCache: + """Build an OmniTensorPrefixCache, but don't init mm tensors.""" + return OmniTensorPrefixCache( + num_blocks=NUM_BLOCKS, + block_size=BLOCK_SIZE, + hidden_size=HIDDEN_SIZE, + hs_dtype=DTYPE, + ) + + +def test_drain_ready_async_writes_is_a_noop_before_the_event_is_ready(monkeypatch): + """drain_ready_async_writes() must not touch the cache while the + scheduled D2H's event has not completed, and must consume it exactly + once after it has. + + Real GPU timing (e.g. via torch.cuda._sleep to stall the copy stream) + was tried and found unreliable on this box: cross-stream wait_stream/ + wait_event did not reliably delay event.query() from reporting ready + immediately, likely a virtualization/driver quirk of this cloud GPU. + Forcing the event's query() result directly tests the branching logic + in drain_ready_async_writes/_consume_pending_write deterministically, + independent of real CUDA scheduling behavior. + """ + cache = get_omni_pcache() + mm_key = "codes.audio" + cache.maybe_init_missing_mm_cache_keys({mm_key: torch.zeros(4, HIDDEN_SIZE)}, seq_len=4) + + num_tokens = 4 + slot_mapping_gpu = torch.arange(8, 8 + num_tokens, dtype=torch.int64, device="cuda") + mm_gpu = { + mm_key: torch.arange(num_tokens * HIDDEN_SIZE, dtype=torch.float32, device="cuda").reshape( + num_tokens, HIDDEN_SIZE + ) + } + + cache.schedule_async_write( + hidden_states_gpu=None, + multimodal_outputs_gpu=mm_gpu, + slot_mapping_gpu=slot_mapping_gpu, + num_tokens_unpadded=num_tokens, + num_tokens_padded=num_tokens, + ) + + pending_event = cache._pending_write.event + rows = cache.mm_outputs_cache[mm_key].view(-1, HIDDEN_SIZE) + + monkeypatch.setattr(pending_event, "query", lambda: False) + assert cache.drain_ready_async_writes() == 0 + assert torch.all(rows[8 : 8 + num_tokens] == 0) + assert cache._pending_write is not None + + monkeypatch.setattr(pending_event, "query", lambda: True) + assert cache.drain_ready_async_writes() == 1 + assert torch.equal(rows[8 : 8 + num_tokens], mm_gpu[mm_key].cpu()) + assert cache._pending_write is None + + # Nothing pending anymore: a second drain is a no-op, not a re-scatter. + assert cache.drain_ready_async_writes() == 0 + + +def test_schedule_async_write_consumes_previous_pending_write_first(): + """Scheduling step N+1's write must first consume (scatter) step N's + pending write, even if the caller never explicitly called + drain_ready_async_writes() in between -- otherwise step N's data would + be silently dropped when overwritten by step N+1's pending write.""" + cache = get_omni_pcache() + mm_key = "codes.audio" + cache.maybe_init_missing_mm_cache_keys({mm_key: torch.zeros(4, HIDDEN_SIZE)}, seq_len=4) + + num_tokens = 4 + first_slots = torch.arange(8, 8 + num_tokens, dtype=torch.int64, device="cuda") + first_gpu = {mm_key: torch.full((num_tokens, HIDDEN_SIZE), 1.0, device="cuda")} + second_slots = torch.arange(12, 12 + num_tokens, dtype=torch.int64, device="cuda") + second_gpu = {mm_key: torch.full((num_tokens, HIDDEN_SIZE), 2.0, device="cuda")} + + cache.schedule_async_write( + hidden_states_gpu=None, + multimodal_outputs_gpu=first_gpu, + slot_mapping_gpu=first_slots, + num_tokens_unpadded=num_tokens, + num_tokens_padded=num_tokens, + ) + torch.accelerator.synchronize() + + # This call's internal _consume_pending_write() must flush the FIRST + # write before the second one is scheduled; the caller never drained. + cache.schedule_async_write( + hidden_states_gpu=None, + multimodal_outputs_gpu=second_gpu, + slot_mapping_gpu=second_slots, + num_tokens_unpadded=num_tokens, + num_tokens_padded=num_tokens, + ) + torch.accelerator.synchronize() + cache.drain_ready_async_writes() + + rows = cache.mm_outputs_cache[mm_key].view(-1, HIDDEN_SIZE) + assert torch.all(rows[8:12] == 1.0) + assert torch.all(rows[12:16] == 2.0) diff --git a/tests/utils/test_mm_outputs.py b/tests/utils/test_mm_outputs.py new file mode 100644 index 000000000..a07effb98 --- /dev/null +++ b/tests/utils/test_mm_outputs.py @@ -0,0 +1,147 @@ +import pytest +import torch + +from vllm_omni.utils import mm_outputs as mm_outputs_mod +from vllm_omni.utils.mm_outputs import build_mm_cpu + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +TOTAL_TOKENS = 5 +NUM_TOKEN_ALIGNED_TENSORS = 3 # codes.audio, meta.ref_code_len, meta.codec_streaming + + +def make_talker_passthrough_payload(num_requests: int, device: str = "cpu") -> dict: + """The passthrough set qwen3-tts actually hands to build_mm_cpu each decode + step (qwen3_tts_talker.make_omni_output): a nested codes/meta dict whose + ``codes.ref`` holds one tensor per request, alongside 1D int token-aligned + metadata. Under prefix caching only ``codes.audio`` is 2D and token-aligned, + so it is the sole cache key and this whole structure is passthrough. + """ + return { + "codes": { + "audio": torch.rand(TOTAL_TOKENS, 4, device=device), + "ref": [torch.arange(idx + 1, dtype=torch.long, device=device) for idx in range(num_requests)], + }, + "meta": { + "ref_code_len": torch.full((TOTAL_TOKENS,), 7, dtype=torch.int32, device=device), + "codec_streaming": torch.ones(TOTAL_TOKENS, dtype=torch.int8, device=device), + }, + } + + +def iter_tensors(value): + """Yield every tensor leaf, recursing the same dict/list nesting _to_cpu does.""" + if isinstance(value, torch.Tensor): + yield value + elif isinstance(value, dict): + for item in value.values(): + yield from iter_tensors(item) + elif isinstance(value, list): + for item in value: + yield from iter_tensors(item) + + +@pytest.mark.parametrize("num_requests", [1, 3]) +def test_build_mm_cpu_d2h_count_grows_per_request(num_requests, mocker): + """Pins the pre-W2 cost: build_mm_cpu issues one per-tensor copy for each + token-aligned payload tensor plus one more for every request's ``codes.ref`` + entry, so the transfer count grows linearly with batch size. That growth is + what RFC #4855 W2 sets out to batch away. + + Counting tensor leaves rather than dict keys is deliberate: a flat dict of + K tensors makes ``call_count == K`` true by construction, since build_mm_cpu + is a loop over items. Only the per-request list can move this number. + """ + spy = mocker.spy(mm_outputs_mod, "_to_cpu") + + result = build_mm_cpu(multimodal_outputs=make_talker_passthrough_payload(num_requests)) + + tensor_copies = [call for call in spy.call_args_list if isinstance(call.args[0], torch.Tensor)] + assert len(tensor_copies) == NUM_TOKEN_ALIGNED_TENSORS + num_requests + assert len(result["codes"]["ref"]) == num_requests + + +def test_build_mm_cpu_preserves_int_dtypes_of_talker_metadata(): + """ref_code_len is int32 and codec_streaming int8 at the source; a W2 that + stages passthrough through a float buffer would corrupt both silently.""" + result = build_mm_cpu(multimodal_outputs=make_talker_passthrough_payload(3)) + + assert result["meta"]["ref_code_len"].dtype == torch.int32 + assert result["meta"]["codec_streaming"].dtype == torch.int8 + assert all(ref.dtype == torch.long for ref in result["codes"]["ref"]) + + +def test_build_mm_cpu_returns_cpu_tensors_for_every_nested_leaf(): + """The output contract: every tensor leaf, including those inside the + ``codes.ref`` list, comes back on CPU. + + On this lane the inputs are already CPU, so this documents the contract + rather than proving it -- ``.to("cpu")`` is a no-op here. The GPU-gated + sibling below is the one that can actually fail. + """ + result = build_mm_cpu(multimodal_outputs=make_talker_passthrough_payload(3)) + + leaves = list(iter_tensors(result)) + assert len(leaves) == NUM_TOKEN_ALIGNED_TENSORS + 3 + assert all(leaf.device.type == "cpu" for leaf in leaves) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device") +def test_build_mm_cpu_moves_a_gpu_payload_off_the_device(): + """D2H is the entire purpose of build_mm_cpu. RFC #4855 W2 lists "skip the + CPU conversion when the consumer is the same-GPU next stage" as one of its + three options; that option leaves every tensor on the device and no CPU-lane + test can observe it, because there ``.to("cpu")`` is a silent no-op. + """ + payload = make_talker_passthrough_payload(3, device="cuda") + source_leaves = list(iter_tensors(payload)) + assert all(leaf.device.type == "cuda" for leaf in source_leaves) + + result = build_mm_cpu(multimodal_outputs=payload) + + leaves = list(iter_tensors(result)) + assert len(leaves) == len(source_leaves) + assert all(leaf.device.type == "cpu" for leaf in leaves) + + +def test_build_mm_cpu_empty_input_is_a_noop(mocker): + spy = mocker.spy(mm_outputs_mod, "_to_cpu") + + assert build_mm_cpu(multimodal_outputs={}) == {} + spy.assert_not_called() + + +def test_build_mm_cpu_preserves_nested_and_non_tensor_payload_semantics(): + """_to_cpu must recurse into nested dict/list structures, force + non-contiguous tensors to become contiguous, detach from autograd, + and leave dtype/shape/values unchanged; non-tensor leaves (scalar, + string) pass through untouched. Any W2 implementation (cat, deferral, + or a new materializer) must preserve this exact semantics.""" + non_contig = torch.arange(9, dtype=torch.float32).reshape(3, 3).t() + assert not non_contig.is_contiguous() + grad_tensor = torch.tensor([5.0, 6.0], requires_grad=True) + + payload = { + "nested": {"inner_tensor": non_contig, "inner_scalar": "keep-me"}, + "list_of_tensors": [torch.tensor([1.0, 2.0]), "marker", torch.tensor([3.0, 4.0])], + "grad_key": grad_tensor, + "scalar": 42, + } + result = build_mm_cpu(multimodal_outputs=payload) + + inner = result["nested"]["inner_tensor"] + assert inner.is_contiguous() + assert inner.dtype == non_contig.dtype + assert inner.shape == non_contig.shape + assert torch.equal(inner, non_contig) + assert result["nested"]["inner_scalar"] == "keep-me" + + lst = result["list_of_tensors"] + assert torch.equal(lst[0], torch.tensor([1.0, 2.0])) + assert lst[1] == "marker" + assert torch.equal(lst[2], torch.tensor([3.0, 4.0])) + + assert result["grad_key"].requires_grad is False + assert torch.equal(result["grad_key"], grad_tensor.detach()) + + assert result["scalar"] == 42 diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index 050063978..e8c40c690 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/tests/worker/test_gpu_ar_model_runner.py @@ -897,3 +897,139 @@ def test_build_omni_output_splits_client_mm_from_inter_stage_keys(monkeypatch): assert output.inter_stage_outputs is not None assert "hidden" in output.inter_stage_outputs[0] assert "audio" not in output.inter_stage_outputs[0] + + +def test_build_omni_output_filters_multimodal_by_partial_downstream_batch(monkeypatch): + """When only a subset of req_ids_output_copy need the downstream + pooler payload, multimodal_outputs (not just hidden_states, which + test_build_omni_output_copies_hidden_for_partial_downstream_batch + already covers) must be attached only to the downstream requests, in + the right per-request slot -- a skipped middle request must not shift + a later request's slice onto the wrong index.""" + runner = _make_async_output_runner(engine_output_type="latent") + runner.requests = {"r1": object(), "r2": object(), "r3": object()} + runner.omni_prefix_cache = object() + + monkeypatch.setattr( + GPUARModelRunner, + "_resolve_pooler_payload_req_ids", + 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", + 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], [31.0, 32.0]], dtype=torch.float32) + 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], [2.0], [3.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=[[], [], []], + 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), + ) + + # Exact indices pin the mapping: r3 keeps slot 2 rather than shifting into + # the skipped r2's slot. + assert torch.equal(output.inter_stage_outputs[0]["codes.audio"], codes[0:1]) + assert output.inter_stage_outputs[1] is None + assert torch.equal(output.inter_stage_outputs[2]["codes.audio"], codes[2:3]) + + +def test_build_omni_output_uses_combined_prefix_cache_mm_payload_for_partial_downstream_batch(monkeypatch): + """When the prefix-cache merge path is engaged (i.e. + _prepare_prefix_cache_pooler_payload_sources returns a non-None + combined_multimodal_outputs), _build_omni_mm_payload must take the + _build_combined_prefix_cache_mm_payload branch -- keyed by req_id, + not by row/slice -- rather than falling back to build_mm_cpu, and + must still map correctly onto a partial downstream batch. This is a + distinct code path from test_build_omni_output_filters_multimodal_by_partial_downstream_batch, + which only exercises the no-prefix-cache-merge (mm_cpu) fallback. + + The merge itself is stubbed out here; what this pins is the branch selection + plus _unwrap_lists' batch-index selection for a list-valued key. + """ + runner = _make_async_output_runner(engine_output_type="latent") + runner.requests = {"r1": object(), "r2": object(), "r3": object()} + runner.omni_prefix_cache = object() + + # codes.ref arrives whole for every request: the prefix-cache merge uses + # pass_lists_through, so selecting this request's entry is _unwrap_lists' + # job, and it selects by batch index (r3 -> 2), not by downstream position. + ref_codes = [torch.tensor([1]), torch.tensor([2]), torch.tensor([3])] + combined_mm = { + "codes.audio": { + "r1": torch.tensor([[11.0, 12.0]]), + "r3": torch.tensor([[31.0, 32.0]]), + }, + "codes.ref": {"r1": list(ref_codes), "r3": list(ref_codes)}, + } + + monkeypatch.setattr( + GPUARModelRunner, + "_resolve_pooler_payload_req_ids", + 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", + lambda *args, **kwargs: (None, None, combined_mm), + ) + 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._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], [2.0], [3.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, 1, 1], dtype=np.int32), + query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.long), + ) + + assert torch.equal(output.inter_stage_outputs[0]["codes.audio"], combined_mm["codes.audio"]["r1"]) + assert output.inter_stage_outputs[1] is None + assert torch.equal(output.inter_stage_outputs[2]["codes.audio"], combined_mm["codes.audio"]["r3"]) + + assert torch.equal(output.inter_stage_outputs[0]["codes.ref"], ref_codes[0]) + assert torch.equal(output.inter_stage_outputs[2]["codes.ref"], ref_codes[2])