[Perf] Deduplicate AR prefix cache hidden-state CPU staging (#3734)

Signed-off-by: TaffyOfficial <2324465096@qq.com>
Co-authored-by: TaffyOfficial <2324465096@qq.com>

tests/core/test_prefix_cache.py

@@ -157,6 +157,75 @@ def test_update_no_multimodal():
         assert torch.all(slot_states == new_states)
 
 
+def test_update_uses_precomputed_hidden_states_cpu():
+    """Precomputed CPU staging should be the only hidden-state cache source."""
+    cache = get_omni_pcache()
+
+    num_tokens_unpadded = 8
+    slot_offset = 8
+    slot_mapping = torch.arange(slot_offset, slot_offset + num_tokens_unpadded)
+    hidden_states = torch.zeros((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
+    hidden_states_cpu = torch.rand((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
+
+    cache.update_omni_tensor_prefix_cache(
+        hidden_states=hidden_states,
+        hidden_states_cpu=hidden_states_cpu,
+        multimodal_outputs=None,
+        num_tokens_unpadded=num_tokens_unpadded,
+        slot_mapping=slot_mapping,
+    )
+
+    hs_rows = cache.hidden_states_cache.view(NUM_BLOCKS * BLOCK_SIZE, HIDDEN_SIZE)
+    for slot_idx, new_states in zip(slot_mapping, hidden_states_cpu):
+        slot_states = hs_rows[slot_idx]
+        assert torch.all(slot_states == new_states)
+
+
+@pytest.mark.parametrize(
+    "hidden_states_cpu",
+    [
+        torch.rand((4, HIDDEN_SIZE), dtype=DTYPE),
+        torch.rand((8, HIDDEN_SIZE), dtype=OTHER_DTYPE),
+        torch.rand((8, HIDDEN_SIZE + 1), dtype=DTYPE),
+        torch.rand((HIDDEN_SIZE, 8), dtype=DTYPE).t(),
+    ],
+)
+def test_update_rejects_invalid_hidden_states_cpu(hidden_states_cpu):
+    cache = get_omni_pcache()
+
+    num_tokens_unpadded = 8
+    slot_mapping = torch.arange(num_tokens_unpadded)
+    hidden_states = torch.rand((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
+
+    with pytest.raises(RuntimeError):
+        cache.update_omni_tensor_prefix_cache(
+            hidden_states=hidden_states,
+            hidden_states_cpu=hidden_states_cpu,
+            multimodal_outputs=None,
+            num_tokens_unpadded=num_tokens_unpadded,
+            slot_mapping=slot_mapping,
+        )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_update_rejects_gpu_hidden_states_cpu():
+    cache = get_omni_pcache()
+
+    num_tokens_unpadded = 8
+    slot_mapping = torch.arange(num_tokens_unpadded)
+    hidden_states = torch.rand((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
+    hidden_states_cpu = torch.rand((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE, device="cuda")
+
+    with pytest.raises(RuntimeError):
+        cache.update_omni_tensor_prefix_cache(
+            hidden_states=hidden_states,
+            hidden_states_cpu=hidden_states_cpu,
+            multimodal_outputs=None,
+            num_tokens_unpadded=num_tokens_unpadded,
+            slot_mapping=slot_mapping,
+        )
+
+
 @pytest.mark.parametrize(
     "feat_dims",
     [
@@ -359,6 +428,76 @@ def test_get_merged_hidden_states(num_tokens_padded, mocker):
     assert torch.all(req2_merged_states == req2_new_states)
 
 
+def test_get_merged_hidden_states_uses_precomputed_hidden_states_cpu(mocker):
+    cache = get_omni_pcache()
+
+    orig_num_tokens_unpadded = 8
+    slot_offset = 8
+    orig_slot_mapping = torch.arange(slot_offset, slot_offset + orig_num_tokens_unpadded)
+    orig_hidden_states = torch.rand((orig_num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
+
+    cache.update_omni_tensor_prefix_cache(
+        hidden_states=orig_hidden_states,
+        multimodal_outputs=None,
+        num_tokens_unpadded=orig_num_tokens_unpadded,
+        slot_mapping=orig_slot_mapping,
+    )
+
+    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_hidden_states = torch.zeros(
+        (num_new_toks_req1 + num_new_toks_req2, HIDDEN_SIZE),
+        dtype=DTYPE,
+    )
+    hidden_states_cpu = torch.rand_like(new_hidden_states)
+    req1_new_states = hidden_states_cpu[:num_new_toks_req1]
+    req2_new_states = hidden_states_cpu[-num_new_toks_req2:]
+    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_states = cache.get_merged_hidden_states(
+        query_start_loc=[0, num_new_toks_req1],
+        input_batch=input_batch,
+        hidden_states=new_hidden_states,
+        hidden_states_cpu=hidden_states_cpu,
+        num_scheduled_tokens=num_scheduled_tokens,
+    )
+
+    assert torch.all(merged_states["req1"][:orig_num_tokens_unpadded] == orig_hidden_states)
+    assert torch.all(merged_states["req1"][-num_new_toks_req1:] == req1_new_states)
+    assert torch.all(merged_states["req2"] == req2_new_states)
+
+
+def test_get_merged_hidden_states_rejects_short_hidden_states_cpu(mocker):
+    cache = get_omni_pcache()
+    num_scheduled_tokens = {
+        "req1": 3,
+        "req2": 2,
+    }
+    input_batch = MockInputBatch(num_computed_tokens_cpu=torch.Tensor([0, 0]))
+
+    mocker.patch(
+        "vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids",
+        new=fake_get_cached_block_ids,
+    )
+    with pytest.raises(RuntimeError):
+        cache.get_merged_hidden_states(
+            query_start_loc=[0, 3],
+            input_batch=input_batch,
+            hidden_states=torch.rand((5, HIDDEN_SIZE), dtype=DTYPE),
+            hidden_states_cpu=torch.rand((4, HIDDEN_SIZE), dtype=DTYPE),
+            num_scheduled_tokens=num_scheduled_tokens,
+        )
+
+
 @pytest.mark.parametrize("num_tokens_padded", [None, 16])
 @pytest.mark.parametrize(
     "feat_dims",

tests/worker/test_gpu_ar_model_runner.py

@@ -1,8 +1,9 @@
 from types import SimpleNamespace
 
 import pytest
+import torch
 
-from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState, GPUARModelRunner
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -41,3 +42,83 @@ def test_resolve_pooler_payload_req_ids_downstream_stage_uses_filtered_requests(
 
     assert engine_output_type == "latent"
     assert payload_req_ids == ["r2"]
+
+
+@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
+
+    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)
+
+    assert torch.equal(output.pooler_output[0]["hidden"], torch.tensor([[1.0, 10.0]]))
+    assert torch.equal(
+        output.pooler_output[1]["hidden"],
+        torch.tensor([[2.0, 20.0], [3.0, 30.0]]),
+    )

tests/worker/test_gpu_generation_model_runner.py

@@ -1,7 +1,10 @@
 import pytest
 import torch
 
-from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner
+from vllm_omni.worker.gpu_generation_model_runner import (
+    ExecuteModelState,
+    GPUGenerationModelRunner,
+)
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -16,7 +19,8 @@ class _DummyInputBatch:
 
 def _make_runner(multimodal_outputs):
     runner = object.__new__(GPUGenerationModelRunner)
-    runner.execute_model_state = (
+    runner.execute_model_state = ExecuteModelState(
+        None,
         None,
         None,
         None,