DeepInfer/tests/hicache/test_host_cache_l3.py

"""HostPrefixCache L3 (Mooncake Store) logic against a dict-backed fake store: write-behind of page groups and
states, a second cache (another rank / instance) extending its empty host pool from L3, no re-puts, and a fetch that
loses an object between exists() and get(). Same fake engine as test_host_cache.py. Needs one GPU."""

from __future__ import annotations

import ctypes

import numpy as np
import pytest
import torch

from tests.hicache.test_host_cache import BLOCK, PAGE, STATE, FakeSeq, _prefill, _tokens, engine  # noqa: F401

pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")

GROUP = 4


class FakeL3:
    """In-memory stand-in for MooncakeL3 with the same interface; puts complete synchronously."""

    def __init__(self) -> None:
        self.group_pages = GROUP
        self.objects: dict[str, bytes] = {}
        self.puts = 0
        self.drop_on_get: set[str] = set()  # keys that vanish between exists() and get()
        self._done: list[object] = []

    def register(self, ptr: int, nbytes: int) -> None:
        pass

    def group_key(self, h: bytes) -> str:
        return "dsv4kv/test/g4/" + h.hex()

    def state_key(self, h: bytes) -> str:
        return "state/test/" + h.hex()

    def submit_put(self, keys, ptrs, sizes, token) -> None:  # noqa: ANN001
        for k, ps, ss in zip(keys, ptrs, sizes, strict=True):
            if k not in self.objects:
                self.objects[k] = b"".join(ctypes.string_at(p, s) for p, s in zip(ps, ss, strict=True))
                self.puts += 1
        self._done.append(token)

    def drain_done(self) -> list[object]:
        out, self._done = self._done, []
        return out

    def exists(self, keys: list[str]) -> list[bool]:
        return [k in self.objects for k in keys]

    def get(self, keys, ptrs, sizes) -> list[bool]:  # noqa: ANN001
        ok = []
        for k, ps, ss in zip(keys, ptrs, sizes, strict=True):
            data = None if k in self.drop_on_get else self.objects.get(k)
            if data is None or len(data) != sum(ss):
                ok.append(False)
                continue
            off = 0
            for p, s in zip(ps, ss, strict=True):
                ctypes.memmove(p, data[off : off + s], s)
                off += s
            ok.append(True)
        return ok

    def stats_dict(self) -> dict:
        return {"l3_put_objects": self.puts}


@pytest.fixture
def l3(monkeypatch):
    from deepinfer.hicache import HostPrefixCache

    store = FakeL3()
    monkeypatch.setattr(HostPrefixCache, "_make_l3", lambda self: store)
    return store


def _settle(cache) -> None:  # noqa: ANN001
    torch.cuda.synchronize()
    cache.poll()  # retires D2H copies and submits the planned L3 puts
    cache.poll()  # drains the put completions (refs back to 0)


def test_second_cache_hits_from_l3_with_exact_bytes(engine, l3):  # noqa: F811
    from deepinfer.hicache import HostPrefixCache

    blocks = engine.driver_worker.cache_engine.gpu_blocks
    states = engine.driver_worker.cache_engine.gpu_states
    a_cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8)
    a = FakeSeq("a", _tokens(1000, 1))
    _prefill(engine, a_cache, a)
    _settle(a_cache)
    # 7 pages -> one full group (pages 0..3) + states at 896/768/640/512
    assert len([k for k in l3.objects if k.startswith("dsv4kv/")]) == 1
    assert len([k for k in l3.objects if k.startswith("state/")]) == 4
    assert all(e.refs == 0 for e in list(a_cache.pages.values()) + list(a_cache.states.values()))
    a_pages = [blocks[b].clone() for b in a.block_table[:4]]
    a_state = states[a.state_id, :STATE].clone()

    # Another rank / instance: empty host pool, same prefix + a new turn. Host alone would miss; L3 gives the group
    # and the state at 512 (the only boundary inside the stored group).
    b_cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8)
    b = FakeSeq("b", np.concatenate([a.tokens, _tokens(500, 2)]))
    assert b_cache.try_resume(b, "b", b.tokens) == 4 * BLOCK
    torch.cuda.synchronize()
    b_cache.poll()
    assert all(torch.equal(blocks[b.block_table[i]], a_pages[i]) for i in range(4))
    assert torch.equal(states[b.state_id, :STATE], a_state)
    s = b_cache.stats
    assert s.l3_lookups == 1 and s.l3_hits == 1 and s.l3_pages_loaded == 4 and s.l3_states_loaded == 1
    assert all(e.refs == 0 for e in list(b_cache.pages.values()) + list(b_cache.states.values()))
    assert len(b_cache.free_pages) + len(b_cache.pages) == b_cache.num_pages


def test_no_reput_of_known_objects(engine, l3):  # noqa: F811
    from deepinfer.hicache import HostPrefixCache

    cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8)
    a = FakeSeq("a", _tokens(1000, 1))
    _prefill(engine, cache, a)
    _settle(cache)
    puts = l3.puts
    again = FakeSeq("a-again", a.tokens.copy())
    _prefill(engine, cache, again)  # identical prompt: nothing new on the host, nothing new for L3
    _settle(cache)
    assert l3.puts == puts


def test_object_lost_between_exists_and_get_falls_back(engine, l3):  # noqa: F811
    from deepinfer.hicache import HostPrefixCache

    a_cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8)
    a = FakeSeq("a", _tokens(1000, 1))
    _prefill(engine, a_cache, a)
    _settle(a_cache)
    l3.drop_on_get = {k for k in l3.objects if k.startswith("dsv4kv/")}

    b_cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8)
    b = FakeSeq("b", np.concatenate([a.tokens, _tokens(500, 2)]))
    assert b_cache.try_resume(b, "b", b.tokens) == 0  # miss, sequence untouched
    assert b.block_table == [] and b_cache.stats.l3_get_failed == 1
    assert len(b_cache.pages) == 0 and len(b_cache.free_pages) == b_cache.num_pages