vllm-omni/tests/platforms/npu/test_npu_prefix_cache_wiring.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""NPU prefix-cache wiring: save_outputs -> sid -> materialize on CPU.

The NPU runner consumes the prefix cache in eager mode (no CUDA streams) and
hands the step id across execute_model()/sample_tokens() as the LAST field of
its positionally-packed ``ExecuteModelState``. Neither path runs in CUDA CI,
so a hit-read change can silently break them (`vllm_ascend` imports keep the
runner modules out of reach of a plain unit test).

Two layers of protection, both pure CPU:

1. AST characterization of the ``ExecuteModelState`` field order in the NPU
   and GPU runners — the pack/unpack sites pass fields positionally, so a
   reorder or drop of ``prefix_cache_step_id`` corrupts the sid without any
   import-time error.
2. The eager-mode facade contract the NPU runner relies on: controller
   auto-selects eager when CUDA is unavailable, and save -> materialize
   returns byte-identical rows for both the miss and the hit-merge paths.
"""

import ast
from pathlib import Path
from types import SimpleNamespace

import pytest
import torch

from vllm_omni.core.prefix_cache.interface import PrefixCacheConfig
from vllm_omni.core.prefix_cache.manager import OmniPrefixCacheManager

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]

_REPO_ROOT = Path(__file__).resolve().parents[3]
_NPU_RUNNER = _REPO_ROOT / "vllm_omni/platforms/npu/worker/npu_ar_model_runner.py"
_GPU_RUNNER = _REPO_ROOT / "vllm_omni/worker/gpu_ar_model_runner.py"

NUM_BLOCKS = 16
BLOCK_SIZE = 4
HIDDEN = 8
DTYPE = torch.float32


def _named_tuple_fields(path: Path, class_name: str) -> list[str]:
    tree = ast.parse(path.read_text())
    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef) and node.name == class_name:
            return [stmt.target.id for stmt in node.body if isinstance(stmt, ast.AnnAssign)]
    raise AssertionError(f"{class_name} not found in {path}")


def test_execute_model_state_keeps_prefix_cache_sid_last():
    """Both runners pack ExecuteModelState positionally: the sid must stay
    the LAST field, and the NPU field order is pinned exactly — a reorder
    passes import and every CUDA test, then ships the wrong value as the
    sid on hardware."""
    npu_fields = _named_tuple_fields(_NPU_RUNNER, "ExecuteModelState")
    assert npu_fields == [
        "scheduler_output",
        "logits",
        "spec_decode_metadata",
        "spec_decode_common_attn_metadata",
        "hidden_states",
        "sample_hidden_states",
        "aux_hidden_states",
        "attn_metadata",
        "positions",
        "ec_connector_output",
        "cudagraph_stats",
        "batch_desc",
        "multimodal_outputs",
        "prefix_cache_step_id",
    ]
    gpu_fields = _named_tuple_fields(_GPU_RUNNER, "ExecuteModelState")
    assert gpu_fields[-1] == "prefix_cache_step_id"


class _FakeView:
    """Minimal group-view double (mirrors tests/core/test_prefix_cache)."""

    def __init__(self):
        self.block_size = BLOCK_SIZE
        self.num_blocks = NUM_BLOCKS
        self.req_blocks: dict[str, list[int]] = {}
        self.order: list[str] = []
        self.computed: dict[str, int] = {}

    def slots_for(self, req_id, token_start, token_end):
        blocks = self.req_blocks[req_id]
        return torch.tensor(
            [blocks[p // BLOCK_SIZE] * BLOCK_SIZE + p % BLOCK_SIZE for p in range(token_start, token_end)],
            dtype=torch.long,
        )

    def batch_req_ids(self) -> list[str]:
        return list(self.order)

    def step_slots_cpu(self, req_ids, num_scheduled) -> torch.Tensor:
        parts = []
        for r in req_ids:
            n = int(num_scheduled.get(r, 0))
            if n > 0:
                start = self.computed.get(r, 0)
                parts.append(self.slots_for(r, start, start + n))
        return torch.cat(parts) if parts else torch.empty((0,), dtype=torch.long)


def _run_step(mgr, view, req_id, blocks, start_pos, sched, *, hit=0, finished=()):
    """One NPU-shaped step: new_step_starts -> save_outputs -> sid."""
    view.order = [req_id]
    view.req_blocks[req_id] = blocks
    view.computed[req_id] = start_pos
    slots = view.slots_for(req_id, start_pos, start_pos + sched)
    hidden = slots.to(DTYPE).unsqueeze(1).expand(sched, HIDDEN).clone()
    sched_out = SimpleNamespace(
        scheduled_new_reqs=[
            SimpleNamespace(req_id=req_id, num_computed_tokens=hit, block_ids=[list(blocks)]),
        ],
        finished_req_ids=set(finished),
        num_scheduled_tokens={req_id: sched},
    )
    mgr.new_step_starts(sched_out)
    return mgr.save_outputs(hidden, {}, num_tokens_unpadded=sched, num_tokens_padded=sched)


def _expected_rows(slots: torch.Tensor) -> torch.Tensor:
    return slots.to(DTYPE).unsqueeze(1).expand(slots.numel(), HIDDEN)


def _make_npu_mode_manager(monkeypatch) -> tuple[OmniPrefixCacheManager, _FakeView]:
    # The NPU condition: no CUDA -> the controller must auto-select eager
    # (submit() completes copy+scatter synchronously), exactly how
    # npu_model_runner builds the manager with the default eager=None.
    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
    view = _FakeView()
    config = PrefixCacheConfig(num_blocks=NUM_BLOCKS, block_size=BLOCK_SIZE)
    return OmniPrefixCacheManager(config, view), view


def test_npu_mode_auto_selects_eager_and_roundtrips(monkeypatch):
    mgr, view = _make_npu_mode_manager(monkeypatch)
    assert mgr._controller._eager is True
    sid = _run_step(mgr, view, "a", [0, 1], 0, 8)
    outs = mgr.materialize(sid, ["a"])
    assert torch.equal(outs.hidden_states["a"], _expected_rows(view.slots_for("a", 0, 8)))


def test_npu_mode_hit_merges_cached_prefix(monkeypatch):
    """The hit-read path in eager mode: consumer's merged output must be
    [cached prefix rows + this step's rows], byte-identical."""
    mgr, view = _make_npu_mode_manager(monkeypatch)
    s1 = _run_step(mgr, view, "a", [0, 1], 0, 8)
    mgr.materialize(s1, ["a"])
    s2 = _run_step(mgr, view, "b", [0, 1, 2], 8, 4, hit=8, finished=["a"])
    merged = mgr.materialize(s2, ["b"]).hidden_states["b"]
    assert merged.shape == (12, HIDDEN)
    assert torch.equal(merged[:8], _expected_rows(view.slots_for("b", 0, 8)))
    assert torch.equal(merged[8:], _expected_rows(view.slots_for("b", 8, 12)))