"""Trace vLLM prefix-cache lookup for Gemma 4 with and without MTP (in-process engine, monkeypatched coordinator).

Written against vLLM 0.29.0 internals (tokenflow.dev/notes/gemma-4-mtp-prefix-cache); other versions may rename these hooks.

  python mtp_cache_trace.py nomtp            # no speculative decoding
  python mtp_cache_trace.py mtp              # MTP K=7, default retention
  python mtp_cache_trace.py mtp 1024         # MTP K=7, --prefix-cache-retention-interval 1024

Expects the target model at /model and the MTP drafter at /drafter.
The note's table came from the first [cached-blocks] line; the second line was added after review.
"""
import os, sys
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
from vllm import LLM, SamplingParams
from vllm.v1.core import kv_cache_coordinator as kc, single_type_kv_cache_manager as st

MTP = len(sys.argv) > 1 and sys.argv[1] == "mtp"
LOG = []

def wrap_find(cls):
    orig = cls.find_longest_cache_hit.__func__
    def f(klass, **kw):
        if klass is st.SlidingWindowManager:
            from vllm.v1.core.single_type_kv_cache_manager import resolve_block_hashes
            bp = kw["block_pool"]; spec = kw["kv_cache_spec"]
            bh = resolve_block_hashes(kw["block_hashes"], bp.hash_block_size, spec.block_size,
                                      supports_fine_grained_hash_lookup=klass.supports_fine_grained_hash_lookup, alignment_tokens=kw["alignment_tokens"])
            n_blk = kw["max_length"] // spec.block_size
            per_group = {}
            for gid in kw["kv_cache_group_ids"]:
                idx = [i for i in range(n_blk) if bp.get_cached_block(bh[i], [gid])]
                rng = []; 
                for i in idx:
                    if rng and rng[-1][1] == i - 1: rng[-1][1] = i
                    else: rng.append([i, i])
                per_group[gid] = rng
            print(f"  [cached-blocks] within lookup range 0..{n_blk - 1}: {per_group}", flush=True)
            # Blocks past the lookup range can also be cached (the published note's table shows only the range above).
            beyond = {gid: [i for i in range(n_blk, len(bh)) if bp.get_cached_block(bh[i], [gid])] for gid in kw["kv_cache_group_ids"]}
            print(f"  [cached-blocks] beyond lookup range: {beyond}", flush=True)
        blocks, n = orig(klass, **kw)
        print(f"  [find] {klass.__name__} groups={kw['kv_cache_group_ids']} max_len={kw['max_length']} drop_eagle={kw.get('drop_eagle_block')} -> hit={n}", flush=True)
        return blocks, n
    cls.find_longest_cache_hit = classmethod(f)
for c in (st.FullAttentionManager, st.SlidingWindowManager):
    wrap_find(c)

orig_init = kc.HybridKVCacheCoordinator.verify_and_split_kv_cache_groups
def vinit(self):
    orig_init(self)
    for i, g in enumerate(self.kv_cache_config.kv_cache_groups):
        s = g.kv_cache_spec
        print(f"[group {i}] {type(s).__name__} block={s.block_size} sliding={getattr(s,'sliding_window',None)} eagle={g.is_eagle_group} layers={len(g.layer_names)} e.g. {g.layer_names[:2]}", flush=True)
    print("[coord] eagle_group_ids", self.eagle_group_ids, "scheduler_block", self.scheduler_block_size, "hash_block", self.hash_block_size,
          "reprefillable", getattr(self, "num_reprefillable_tokens", None), "attention_groups", [(type(a.spec).__name__, a.group_ids, a.use_eagle) for a in self.attention_groups], flush=True)
kc.HybridKVCacheCoordinator.verify_and_split_kv_cache_groups = vinit

orig_cache = kc.HybridKVCacheCoordinator.cache_blocks
def cb(self, request, num_computed_tokens):
    orig_cache(self, request, num_computed_tokens)
    LOG.append((num_computed_tokens, [(type(m).__name__, m.use_eagle, m.num_cached_block.get(request.request_id, 0)) for m in self.single_type_managers]))
kc.HybridKVCacheCoordinator.cache_blocks = cb

kw = dict(model="/model", max_model_len=32768, gpu_memory_utilization=0.25, kv_cache_memory_bytes=4 << 30, kv_cache_dtype="fp8_e4m3",
          enable_prefix_caching=True, max_num_seqs=4, limit_mm_per_prompt={"image": 0, "audio": 0}, enforce_eager=True)
if len(sys.argv) > 2:
    kw["prefix_cache_retention_interval"] = int(sys.argv[2])
if MTP:
    kw["speculative_config"] = {"method": "mtp", "model": "/drafter", "num_speculative_tokens": 7}
llm = LLM(**kw)
prompt = " ".join(f"w{i}" for i in range(1500)) + " ok?"
sp = SamplingParams(max_tokens=4, temperature=0)
for k in range(2):
    print(f"=== request {k+1}", flush=True)
    out = llm.generate([prompt], sp)
    print(f"=== request {k+1} num_cached_tokens={out[0].num_cached_tokens} prompt_len={len(out[0].prompt_token_ids)} text={out[0].outputs[0].text!r}", flush=True)
# multi-turn style: same prefix + new suffix
out = llm.generate([prompt + " And then? " + " ".join(f"z{i}" for i in range(300))], sp)
print(f"=== request 3 (prefix+suffix) num_cached_tokens={out[0].num_cached_tokens} prompt_len={len(out[0].prompt_token_ids)}", flush=True)
print("[cache_blocks calls, last 3]", LOG[-3:])