[Test] Prefix-cache passthrough test coverage prerequisite for W2 (#4855) (#5310)

Signed-off-by: Shenglei Fu <sfu@confluent.io>
Co-authored-by: Claude <noreply@anthropic.com>

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)

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)

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

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])