[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>
2026-08-04 20:03:57 -07:00
1 parent bfaaa3c2c3
commit f24a6165ff
@@ -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])