From 5046e4a00f94e5dea3407072241839a258a0510d Mon Sep 17 00:00:00 2001 From: alvarorsouza-arch Date: Sat, 19 Sep 2026 07:55:45 +0000 Subject: [PATCH 1/2] feat(kvcache): host-resident prefix tier for hybrid GDN models Prefix reuse on hybrid GDN models needs KV pages AND a GDN state snapshot alive at the match point; the small global snapshot pool gets drained by interleaved auxiliary calls (agent title generators) and every match truncates to zero. Measured on real WorkBuddy agent turns: TTFT 20-150s recompute per message. This adds a host-resident tier: evicted GDN snapshots and KV prefix pages move to pinned RAM instead of dying. Matches rehydrate from host in ~4s (24k-token prefix) instead of a full recompute. - HybridRadixCache hooks: on_evict_node / on_match_dead_snapshot; evicted leaves become tombstones instead of cascade-freeing KV - HostPrefixTier: pinned GDN snapshot bank + cold-prefix index with its own KV slot space (logical ids get reused by the scheduler, so content is copied D2H on eviction, never referenced by id) - Admission rehydration; with KVHostOffloader active the copy is host-to-host into its mirror (no bulk H2D; lazy fetch covers residency) - Eviction copies run in 128-page slices (staging the whole prefix OOMed exactly under memory pressure, which is when eviction runs) Measured: churn test cached-token 0 -> 23,936 (16.5s -> 3.8s); total pool flush rehydrates at 24,000 cached tokens in 4.4s; quality suite identical (39/40 QA, 30/30 cloze); decode steady 33-37 t/s. All off by default (FT_GDN_HOST_TIER / FT_PREFIX_HOST envs). --- python/freetoken/engine/config.py | 9 +- python/freetoken/kvcache/host_prefix_tier.py | 318 ++++++++++++++++++ .../freetoken/kvcache/hybrid_radix_cache.py | 43 ++- python/freetoken/scheduler/cache.py | 131 ++++++++ tests/kvcache/test_host_prefix_tier.py | 60 ++++ tests/kvcache/test_tier3.py | 85 +++++ tests/kvcache/test_tree_hooks.py | 77 +++++ 7 files changed, 720 insertions(+), 3 deletions(-) create mode 100644 python/freetoken/kvcache/host_prefix_tier.py create mode 100644 tests/kvcache/test_host_prefix_tier.py create mode 100644 tests/kvcache/test_tier3.py create mode 100644 tests/kvcache/test_tree_hooks.py diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 5df70492bb..5f0a0314d2 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -1,6 +1,7 @@ from __future__ import annotations import copy +import os from dataclasses import dataclass, field, replace from functools import cached_property from typing import TYPE_CHECKING, List @@ -66,11 +67,17 @@ class EngineConfig: cuda_graph_bs: List[int] | None = None cuda_graph_max_bs: int | None = None page_size: int = 1 + # KV-cache storage quantization: "none" stores the compute dtype, "fp8" stores e4m3 + # codes plus one fp32 scale per (token, slab, layer, kv head) -- about 2x the tokens + # per GiB, at a small accuracy cost. --kv-cache-dtype; resolved from "auto" by + # _adjust_config, which also refuses it on a pool family or attention backend that + # cannot read the scales. + kv_quant: str = "none" memory_ratio: float = 0.9 # Hybrid GDN models default to the HybridRadixCache (cross-request GDN-state prefix reuse); # `--cache-type naive` opts out. linear_state_cache_ratio sizes the GDN snapshot cache as # ceil(ratio * max_running_req) extra slots. - linear_state_cache_ratio: float = 2.0 + linear_state_cache_ratio: float = float(os.environ.get("FT_LINEAR_STATE_CACHE_RATIO", "2.0")) # Window/full ratio for the SWA radix cache (`--cache-type radix` on SWA models) and the DSV4 # window tier: the DEFAULT window-pool size = max(working-set floor, ratio x full-pool tokens). # < 1.0 trades retained window-prefix capacity for memory savings; must be in (0, 1]. It is the diff --git a/python/freetoken/kvcache/host_prefix_tier.py b/python/freetoken/kvcache/host_prefix_tier.py new file mode 100644 index 0000000000..0265b126ef --- /dev/null +++ b/python/freetoken/kvcache/host_prefix_tier.py @@ -0,0 +1,318 @@ +"""Host-side cold-prefix tier (Nível 3 — Degrau 1: snapshots GDN; Degrau 3: KV de prefixo). + +Dois artefatos em RAM pinned: +1. um banco de snapshots GDN (write na evicção do tree, read no match/admissão); +2. um índice de prefixos frios: token ids + slots KV host + slot GDN host. + +Degrau 3: o KV vive num **slot space host separado** (banco pinned [2, L, S, page, H, D] de +códigos + scales, mesma geometria do pool GPU). Ids lógicos de página são reutilizados pelo +scheduler, então o conteúdo é COPIADO (D2H) na hora do evict do tree — nunca referenciado por +id lógico depois. A reidratação aloca páginas novas no pool e copia H2D de volta. + +Ver IMPLEMENTACAO-NIVEL3.md para o design completo. +""" + +from __future__ import annotations + +from collections import OrderedDict +from dataclasses import dataclass + +import torch + +from freetoken.moe.host_banks import HostBank + + +class HostMatch: + __slots__ = ("host_len", "kv_pages", "gdn_slot") + + def __init__(self, host_len: int, kv_pages: torch.Tensor | None, gdn_slot: int | None) -> None: + self.host_len = host_len # tokens cobertos (page-aligned) + self.kv_pages = kv_pages # int32 CPU: slots KV do tier (não ids lógicos!) + self.gdn_slot = gdn_slot # slot host do snapshot GDN (None = sem snapshot) + + +@dataclass +class _Entry: + key_tokens: torch.Tensor # CPU, o prefixo (page-aligned) + kv_slots: torch.Tensor | None = None # int32 CPU: um slot host por página + gdn_slot: int | None = None # snapshot GDN host do boundary + + +class HostPrefixTier: + """Índice + storage RAM pra prefixos evictados do radix tree GPU. + + O índice é LRU por entrada (hit re-marca); quando a RAM do tier enche (KV) ou os slots + GDN acabam, drop-oldest libera a entrada mais fria inteira (KV slots + slot GDN). + As páginas KV são copiadas para slots host próprios na evicção — o conteúdo sobrevive à + reutilização dos ids lógicos pelo scheduler. + + Geometria GDN: os tensores do LinearStatePool são [layers, slots, ...]; cada banco host + é [host_slots, layers, ...] — a cópia de um slot é [layers, ...] <-> [layers, ...]. + Geometria KV: pool._kv_buffer [2, L, P, page, H, D] (códigos e4m3 no pool fp8) + + pool._scale_buffer [2, L, P*page, H] fp32; os banks host espelham isso com S host slots. + """ + + def __init__(self, state_pool, kv_pool=None, *, gdn_slots_host: int = 64, + kv_budget_bytes: int = 0, device: torch.device | str = "cuda") -> None: + self.device = torch.device(device) + self.pool = state_pool + tensors = {"conv": state_pool.conv_states, "rec": state_pool.recurrent_states} + tensors.update(state_pool.slot_states) + self._order = list(tensors) + self._banks: dict[str, tuple[torch.Tensor, HostBank]] = {} + for name, t in tensors.items(): + # pool: [layers, slots, *rest] -> bank: [host_slots, layers, *rest] + bank = HostBank((gdn_slots_host, t.shape[0], *t.shape[2:]), t.dtype) + if self.device.type == "cuda" and torch.cuda.is_available(): + bank.pin() + self._banks[name] = (t, bank) + self.num_slots = gdn_slots_host + self._free = list(range(gdn_slots_host)) + self._in_use: set[int] = set() + self._index: OrderedDict[bytes, _Entry] = OrderedDict() + # key: bytes dos token ids page-aligned -> _Entry + + # ---- KV page store (Degrau 3) ---- + self.kv_pool = kv_pool + self.kv_page_bytes = 0 + self._kv_codes: HostBank | None = None + self._kv_scales: HostBank | None = None + self._kv_free: list[int] = [] + if kv_pool is not None and kv_budget_bytes > 0: + codes = kv_pool._kv_buffer # [2, L, P, page, H, D] + scales = kv_pool._scale_buffer # [2, L, P*page, H] | None (fp8) + _, layers, _, page_size, heads, dim = codes.shape + page_bytes = 2 * layers * page_size * heads * dim * codes.element_size() + if scales is not None: + page_bytes += 2 * layers * page_size * heads * scales.element_size() + n = max(1, kv_budget_bytes // page_bytes) + self._kv_codes = HostBank((2, layers, n, page_size, heads, dim), codes.dtype) + if scales is not None: + self._kv_scales = HostBank((2, layers, n, page_size, heads), scales.dtype) + if self.device.type == "cuda" and torch.cuda.is_available(): + self._kv_codes.pin() + if self._kv_scales is not None: + self._kv_scales.pin() + self._kv_free = list(range(n)) + self.kv_page_bytes = page_bytes + self.copy_stream = ( + torch.cuda.Stream(self.device) if self.device.type == "cuda" else None) + + # ---------------- GDN snapshot storage ---------------- + def write_gdn(self, gpu_slot: int) -> int: + """D2H do snapshot do gpu_slot pra um slot host livre. Retorna o host_slot. + Deve ser enfileirado ANTES de state_pool.free(gpu_slot): a ordem de stream do + engine garante leitura antes de qualquer reescrita do slot (ver spec).""" + if not self._free: + raise RuntimeError("HostPrefixTier: sem slots host GDN livres (chame evict_host)") + hs = self._free.pop() + self._in_use.add(hs) + if self.copy_stream is not None: + with torch.cuda.stream(self.copy_stream): + for name in self._order: + t, bank = self._banks[name] + bank.tensor[hs].copy_(t[:, gpu_slot], non_blocking=True) + self.copy_stream.synchronize() + else: + for name in self._order: + t, bank = self._banks[name] + bank.tensor[hs].copy_(t[:, gpu_slot]) + return hs + + def read_gdn(self, host_slot: int, dst_gpu_slot: int) -> None: + """H2D host_slot -> dst_gpu_slot do pool (já alocado pelo caller).""" + if host_slot not in self._in_use: + raise KeyError(f"HostPrefixTier: slot host {host_slot} não está em uso") + if self.copy_stream is not None: + with torch.cuda.stream(self.copy_stream): + for name in self._order: + t, bank = self._banks[name] + t[:, dst_gpu_slot].copy_(bank.tensor[host_slot], non_blocking=True) + self.copy_stream.synchronize() + else: + for name in self._order: + t, bank = self._banks[name] + t[:, dst_gpu_slot].copy_(bank.tensor[host_slot]) + + def free_gdn(self, host_slot: int) -> None: + self._in_use.discard(host_slot) + self._free.append(host_slot) + + # ---------------- KV page storage (Degrau 3) ---------------- + @property + def has_kv(self) -> bool: + return self._kv_codes is not None + + def write_kv(self, kv_indices: torch.Tensor) -> torch.Tensor | None: + """Copia as páginas (token slots do pool, page-aligned) para slots host livres. + Retorna os slots host (int32 CPU), ou None se o tier não tem KV store. + Deve rodar DURANTE o evict do tree, com os bytes ainda válidos no buffer GPU.""" + if not self.has_kv: + return None + page_size = self._kv_codes.tensor.shape[3] + pages = (kv_indices[::page_size] // page_size).long().to(self.device) + n = pages.numel() + if n == 0: + return None + codes, scales = self.kv_pool._kv_buffer, self.kv_pool._scale_buffer + # Defesa: um slot fora do pool indicaria um bug de lifetime no tree. Pular a escrita + # (o prefixo simplesmente não vai pro tier) em vez de derrubar o scheduler. + max_page = int(pages.max().item()) + if max_page >= codes.shape[2]: + from freetoken.utils import init_logger + init_logger(__name__).warning( + f"HostPrefixTier.write_kv: page {max_page} fora do pool " + f"({codes.shape[2]} páginas): n={n}, kv_slots min/max = " + f"{int(kv_indices.min())}/{int(kv_indices.max())}, prefixo de {kv_indices.numel()} tokens -- " + "escrita ignorada" + ) + return None + while len(self._kv_free) < n: + self._drop_oldest() # libera os slots da entrada mais fria (inclui GDN) + slots = torch.tensor([self._kv_free.pop() for _ in range(n)], dtype=torch.int64) + # NB: sem staging do prefixo inteiro — cada fatia calcula seus token-ids. + + def _go(): + # Chunked: NUNCA estagiar o prefixo inteiro (o evict acontece sob pressão de + # memória — um staging GB-sized explode a VRAM e derruba o scheduler). + CH = 128 # páginas por fatia (~100 MB de temp por fatia em fp8) + for i in range(0, n, CH): + p = pages[i : i + CH] + s = slots[i : i + CH] + staged = codes[:, :, p].contiguous() + self._kv_codes.tensor[:, :, s] = staged.cpu() if staged.is_cuda else staged + del staged + if scales is not None: + for i in range(0, n, CH): + p = pages[i : i + CH] + s = slots[i : i + CH] + tk = (p.unsqueeze(1) * page_size + + torch.arange(page_size, device=self.device)).flatten() + staged_s = scales[:, :, tk].view( + 2, scales.shape[1], p.numel(), page_size, scales.shape[3]).contiguous() + self._kv_scales.tensor[:, :, s] = ( + staged_s.cpu() if staged_s.is_cuda else staged_s) + del staged_s + + if self.copy_stream is not None: + with torch.cuda.stream(self.copy_stream): + _go() + self.copy_stream.synchronize() + else: + _go() + return slots.to(torch.int32).cpu() + + def read_kv(self, host_slots: torch.Tensor, dst_pages: torch.Tensor) -> None: + """H2D dos slots host para as páginas dst (ids de página do pool, device).""" + assert self.has_kv + slots = host_slots.long() + pages = dst_pages.long().to(self.device) + codes, scales = self.kv_pool._kv_buffer, self.kv_pool._scale_buffer + page_size = self._kv_codes.tensor.shape[3] + n = slots.numel() + + def _go(): + CH = 128 + for i in range(0, n, CH): + s = slots[i : i + CH] + p = pages[i : i + CH] + staged = self._kv_codes.tensor[:, :, s].to(self.device, non_blocking=True) + codes[:, :, p] = staged + del staged + if scales is not None: + for i in range(0, n, CH): + s = slots[i : i + CH] + p = pages[i : i + CH] + src = self._kv_scales.tensor[:, :, s].reshape( + 2, scales.shape[1], p.numel() * page_size, scales.shape[3]) + tk = (p.unsqueeze(1) * page_size + + torch.arange(page_size, device=self.device)).flatten() + staged_s = src.to(self.device, non_blocking=True) + scales[:, :, tk] = staged_s + del staged_s + + if self.copy_stream is not None: + with torch.cuda.stream(self.copy_stream): + _go() + self.copy_stream.synchronize() + else: + _go() + + def read_kv_to_mirror(self, host_slots: torch.Tensor, mirror_codes, + mirror_scales, logical_pages: torch.Tensor) -> None: + """H2H (pinned memcpy): slots KV do tier -> espelho do offloader nas páginas + lógicas dadas. Usado quando o offloader está ativo: reidratar NÃO sobe bytes pra + GPU — o fetch lazy do offloader cuida da residência na atenção.""" + assert self.has_kv + s = host_slots.long() + lp = logical_pages.long().cpu() + page = self._kv_codes.tensor.shape[3] + mirror_codes[:, :, lp] = self._kv_codes.tensor[:, :, s] + if self._kv_scales is not None and mirror_scales is not None: + t_s = self._kv_scales.tensor.view(2, -1, self._kv_codes.tensor.shape[2], page, + self._kv_scales.tensor.shape[-1]) + m_s = mirror_scales.view(2, -1, mirror_codes.shape[2], page, + self._kv_scales.tensor.shape[-1]) + m_s[:, :, lp] = t_s[:, :, s] + + def _drop_oldest(self) -> None: + if not self._index: + raise RuntimeError("HostPrefixTier: KV host store cheio e índice vazio") + _, ent = self._index.popitem(last=False) # o mais frio + if ent.kv_slots is not None: + self._kv_free.extend(ent.kv_slots.tolist()) + ent.kv_slots = None + if ent.gdn_slot is not None: + self.free_gdn(ent.gdn_slot) + ent.gdn_slot = None + + # ---------------- índice de prefixos ---------------- + @staticmethod + def _key_of(tokens: torch.Tensor) -> bytes: + return tokens.to(torch.int64).cpu().numpy().tobytes() + + def put_prefix(self, key_tokens: torch.Tensor, kv_pages: torch.Tensor | None = None, + host_gdn_slot: int | None = None) -> None: + """Registra (ou faz merge numa entrada existente) um prefixo frio. key_tokens deve + ser page-aligned (o tree já alinha). kv_pages = slots host do KV store.""" + key = self._key_of(key_tokens) + ent = self._index.get(key) + if ent is None: + ent = _Entry(key_tokens.cpu()) + self._index[key] = ent + else: + self._index.move_to_end(key) # refresh LRU + if kv_pages is not None: + if ent.kv_slots is not None and ent.kv_slots is not kv_pages: + self._kv_free.extend(ent.kv_slots.tolist()) # substituição: libera os velhos + ent.kv_slots = kv_pages.to(torch.int32).cpu() + if host_gdn_slot is not None: + if ent.gdn_slot is not None and ent.gdn_slot != host_gdn_slot: + self.free_gdn(ent.gdn_slot) # a cópia nova substitui a anterior + ent.gdn_slot = host_gdn_slot + + def lookup(self, input_ids: torch.Tensor, page_size: int = 64) -> HostMatch | None: + """Maior prefixo page-aligned presente no índice (python/CPU, fora do hot path).""" + toks = input_ids.to(torch.int64).cpu() + n = (toks.numel() // page_size) * page_size + while n > 0: + key = self._key_of(toks[:n]) + hit = self._index.get(key) + if hit is not None: + self._index.move_to_end(key) + return HostMatch(n, hit.kv_slots, hit.gdn_slot) + n -= page_size + return None + + def evict_host(self, key_tokens: torch.Tensor) -> None: + """Remove uma entrada do índice e libera os slots host (KV + GDN).""" + ent = self._index.pop(self._key_of(key_tokens), None) + if ent is None: + return + if ent.kv_slots is not None: + self._kv_free.extend(ent.kv_slots.tolist()) + if ent.gdn_slot is not None: + self.free_gdn(ent.gdn_slot) + + +__all__ = ["HostPrefixTier", "HostMatch"] diff --git a/python/freetoken/kvcache/hybrid_radix_cache.py b/python/freetoken/kvcache/hybrid_radix_cache.py index 1bd9e0b9cd..cb4a7f90b2 100644 --- a/python/freetoken/kvcache/hybrid_radix_cache.py +++ b/python/freetoken/kvcache/hybrid_radix_cache.py @@ -71,6 +71,16 @@ def __init__(self, device: torch.device, page_size: int) -> None: self.full_protected = 0 self.mamba_evictable = 0 # number of live, unlocked snapshots self.mamba_protected = 0 + # Host tier (FT_GDN_HOST_TIER): wired by CacheManager. ``on_evict_node`` fires with + # the node IMMEDIATELY before its snapshot is freed (the slot is still valid -- + # enqueue the D2H before returning). ``on_match_dead_snapshot`` may rehydrate a + # tombstone during match: return a fresh GPU slot to attach, None to keep walking. + self.on_evict_node = None + self.on_match_dead_snapshot = None + # KV host tier (FT_PREFIX_HOST): fires for every node whose KV pages LEAVE the tree + # (evict_full's leaf + cascade-exposed tombstone parents) -- BEFORE the free, while + # the bytes are still valid in the GPU buffer (the D2H must happen inside the call). + self.on_evict_kv = None # ---------------------------------------------------------------- match / insert def match_prefix(self, input_ids: torch.Tensor) -> HybridMatch: @@ -83,6 +93,14 @@ def match_prefix(self, input_ids: torch.Tensor) -> HybridMatch: while not cur.is_root(): if cur.mamba_value is not None: return HybridMatch(self._collect_kv(cur), end_len, cur.mamba_value, cur) + if self.on_match_dead_snapshot is not None: + slot = self.on_match_dead_snapshot(cur) + if slot is not None: + # revive the tombstone: the rehydrated snapshot rides the tree again + cur.mamba_value = slot + if cur.mamba_ref_count == 0: + self.mamba_evictable += 1 + return HybridMatch(self._collect_kv(cur), end_len, slot, cur) end_len -= cur.length cur = cur.parent return HybridMatch(self.empty, 0, None, self.root) @@ -156,8 +174,12 @@ def evict_full(self, num_tokens: int) -> EvictResult: if node.ref_count != 0 or not node.is_leaf() or node.is_root(): continue freed += node.length + if self.on_evict_kv is not None: + self.on_evict_kv(node) kv.append(node.value) self.full_evictable -= node.length + if node.mamba_value is not None and self.on_evict_node is not None: + self.on_evict_node(node) self._free_node_mamba(node, mamba) parent, casc = self._cascade_tombstone_leaves(self._unlink(node), kv) freed += casc @@ -177,14 +199,19 @@ def evict_mamba(self, num: int) -> EvictResult: node = heapq.heappop(cands) if node.mamba_value is None or node.mamba_ref_count != 0 or node.is_root(): continue - if node.is_leaf() and node.ref_count == 0: + if self.on_evict_node is not None: + self.on_evict_node(node) + if node.is_leaf() and node.ref_count == 0 and self.on_evict_node is None: kv.append(node.value) self.full_evictable -= node.length self._free_node_mamba(node, mamba) freed += 1 self._cascade_tombstone_leaves(self._unlink(node), kv) else: - self._free_node_mamba(node, mamba) # tombstone internal (or locked-KV) node + # internal (or locked-KV) node -> tombstone; with a host tier wired, LEAF too: + # the snapshot is safely on the host, and keeping the leaf's KV in the tree + # is what lets a later match rehydrate the boundary instead of re-prefilling. + self._free_node_mamba(node, mamba) freed += 1 return EvictResult(torch.cat(kv) if kv else self.empty, mamba) @@ -230,6 +257,8 @@ def _cascade_tombstone_leaves(self, parent: RadixTreeNode, kv_out: List[torch.Te freed = 0 while (parent.mamba_value is None and parent.is_leaf() and parent.ref_count == 0 and not parent.is_root()): + if self.on_evict_kv is not None: + self.on_evict_kv(parent) kv_out.append(parent.value) self.full_evictable -= parent.length freed += parent.length @@ -252,6 +281,16 @@ def _collect_kv(self, node: RadixTreeNode) -> torch.Tensor: vals.reverse() return torch.cat(vals) if vals else self.empty + def _collect_key(self, node: RadixTreeNode) -> torch.Tensor: + """The full token prefix ending at ``node`` (segment keys concatenated root->node).""" + keys: List[torch.Tensor] = [] + n = node + while not n.is_root(): + keys.append(n._key) + n = n.parent + keys.reverse() + return torch.cat(keys) if keys else torch.empty(0, dtype=torch.int32) + def _leaves(self) -> List[RadixTreeNode]: out, stack = [], [self.root] while stack: diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py index 8f77be9330..3f833eaf18 100644 --- a/python/freetoken/scheduler/cache.py +++ b/python/freetoken/scheduler/cache.py @@ -56,6 +56,31 @@ def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, typ if swa_pool is not None: self.prefill_chunk_budget = getattr(swa_pool, "prefill_chunk_budget", None) self.prefix_cache = self._make_prefix_cache(device, page_size, type) + # Host prefix tier: FT_GDN_HOST_TIER=1 guarda snapshots GDN na RAM (Degrau 2); + # FT_PREFIX_HOST=1 + FT_PREFIX_HOST_GB adiciona o KV de prefixo num slot space host + # separado (Degrau 3 — o radix nunca perde o mapeamento). Hybrid path only. + self.host_tier = None + gdn_on = os.environ.get("FT_GDN_HOST_TIER", "0") == "1" + kv_gb = (float(os.environ.get("FT_PREFIX_HOST_GB", "16")) + if os.environ.get("FT_PREFIX_HOST", "0") == "1" else 0) + if self.is_hybrid and (gdn_on or kv_gb > 0): + from freetoken.core import get_global_ctx + from freetoken.kvcache.host_prefix_tier import HostPrefixTier + from freetoken.utils import init_logger + + slots = int(os.environ.get("FT_GDN_HOST_SLOTS", "32")) + kv_pool = get_global_ctx().kv_cache if kv_gb > 0 else None + self.host_tier = HostPrefixTier( + linear_state_pool, kv_pool, gdn_slots_host=slots, + kv_budget_bytes=int(kv_gb * (1 << 30)), device=device) + self.prefix_cache.on_evict_node = self._on_evict_node_to_host + self.prefix_cache.on_match_dead_snapshot = self._rehydrate_snapshot + if kv_pool is not None: + self.prefix_cache.on_evict_kv = self._on_evict_kv_to_host + init_logger(__name__).info( + f"host prefix tier ON: {slots} GDN slots" + + (f" + {kv_gb:.0f} GiB KV store" if kv_pool is not None else "") + ) self.device = device self.num_pages = num_pages self.page_table = page_table @@ -101,6 +126,16 @@ def match_req(self, req: PendingReq) -> MatchResult: if self.is_hybrid: from freetoken.kvcache.hybrid_radix_cache import HybridCacheHandle m = self.prefix_cache.match_prefix(ids) + # Host prefix tier (FT_PREFIX_HOST): a prefix evicted from the tree lives in RAM + # (KV bytes + GDN snapshot). If the tree match fell short, rehydrate the deeper + # host entry into fresh pages, re-insert it into the tree and re-match. + tier = self.host_tier + if tier is not None and tier.has_kv: + hit = tier.lookup(ids, self.page_size) + if (hit is not None and hit.kv_pages is not None and hit.gdn_slot is not None + and hit.host_len > m.cached_len): + self._rehydrate_host_prefix(ids, hit) + m = self.prefix_cache.match_prefix(ids) return MatchResult( HybridCacheHandle(m.cached_len, m.node, m.kv_indices), mamba_value=m.mamba_value) return self.prefix_cache.match_prefix(ids) @@ -548,6 +583,96 @@ def check_integrity(self) -> None: if self.page_size > 1: assert torch.all(self.free_slots % self.page_size == 0) + # ----- host GDN-snapshot tier (FT_GDN_HOST_TIER) ------------------------------- + def _on_evict_node_to_host(self, node) -> None: + """Tree hook: the node's GDN snapshot is about to be freed. Copy it to the host tier + and index the node's prefix (the GPU slot is freed only after we return).""" + tier = self.host_tier + prefix = self.prefix_cache._collect_key(node) + try: + hs = tier.write_gdn(node.mamba_value) + except RuntimeError: + # host tier full: drop the OLDEST indexed prefix and retry once + if not tier._index: + return + oldest = next(iter(tier._index)) + tier.evict_host(tier._index[oldest].key_tokens) + try: + hs = tier.write_gdn(node.mamba_value) + except RuntimeError: + return + tier.put_prefix(prefix, None, hs) + + def _on_evict_kv_to_host(self, node) -> None: + """Tree hook (FT_PREFIX_HOST): the node's KV pages are about to be freed. Copy the + WHOLE prefix's pages to the host store while the bytes are still valid (eviction is + bottom-up, so the ancestors' pages are still live), and index the prefix.""" + tier = self.host_tier + prefix = self.prefix_cache._collect_key(node) + kv = self.prefix_cache._collect_kv(node) + slots = tier.write_kv(kv) + if slots is not None: + tier.put_prefix(prefix, slots, None) + + def _rehydrate_snapshot(self, node) -> int | None: + """Tree hook (match path): the node is a snapshot tombstone. If the host tier has its + prefix, copy the snapshot back into a fresh pool slot and return the slot; None = + keep walking (the match truncates as before).""" + tier = self.host_tier + hit = tier.lookup(self.prefix_cache._collect_key(node), self.page_size) + if hit is None or hit.gdn_slot is None: + return None + pool = self.linear_state_pool + if pool.num_free_slots < 1: + # make room: evicting another snapshot writes IT to the host tier first + self.ensure_mamba_slots(1) + if pool.num_free_slots < 1: + return None + slot = pool.alloc(1)[0] + tier.read_gdn(hit.gdn_slot, slot) + return slot + + def _rehydrate_host_prefix(self, ids: torch.Tensor, hit) -> None: + """FT_PREFIX_HOST: sobe um prefixo frio do tier pro tree+GPU -- aloca páginas novas + (pode evictar: a vítima é escrita no host primeiro), copia os bytes KV H2D, restaura + o snapshot GDN num slot novo e re-insere o prefixo no tree (o match seguinte acha). + """ + tier = self.host_tier + n_pages = hit.host_len // self.page_size + kv_new = self._page_to_token(self._allocate(n_pages)) + pool = self.linear_state_pool + if hit.gdn_slot is not None: + self.ensure_mamba_slots(1) + if pool.num_free_slots < 1: + self._free(kv_new) # sem slot GDN não dá pra reidratar; devolve as páginas + return + slot = pool.alloc(1)[0] + tier.read_gdn(hit.gdn_slot, slot) + else: + self._free(kv_new) + return # o modelo híbrido não resume sem snapshot GDN + # Com o KVHostOffloader ativo, as páginas recém-alocadas são LÓGICAS: subir bytes + # pra GPU em bulk estoura o pool físico quando o prefixo > pool (o fetch lazy já + # resolve residência na atenção). Copiamos host->host pro espelho e invalidamos + # bindings defasados, em vez de H2D. + from freetoken.core import get_global_ctx + off = getattr(get_global_ctx(), "kv_offloader", None) + new_pages = kv_new[:: self.page_size] // self.page_size + if off is not None: + tier.read_kv_to_mirror(hit.kv_pages, off.mirror.tensor, + off.mirror_scales.tensor if off.mirror_scales is not None else None, + new_pages) + off.invalidate_pages(new_pages) + else: + tier.read_kv(hit.kv_pages, new_pages) + prefix_len, mamba_exist = self.prefix_cache.insert(ids[: hit.host_len], kv_new, slot) + if prefix_len > 0: + # o tree já tinha [0, prefix_len) com as páginas antigas (dedup do insert) -- + # as páginas frescas que aloquei pra essa parte voltam pro free list + self._free(kv_new[:prefix_len]) + if mamba_exist: + pool.free([slot]) # o boundary já tinha snapshot vivo (dedup do insert) + def rebuild(self, num_pages: int, page_table: torch.Tensor) -> None: """Re-point the page table and reset page accounting + prefix cache IN PLACE. @@ -561,6 +686,12 @@ def rebuild(self, num_pages: int, page_table: torch.Tensor) -> None: self.page_table = page_table self.free_slots = torch.arange(num_pages, dtype=torch.int32, device=device) * self.page_size self.prefix_cache = self._make_prefix_cache(device, self.page_size, self.cache_type) + if self.host_tier is not None: + # the rebuilt tree gets the same hooks (the host index survives a rebuild) + self.prefix_cache.on_evict_node = self._on_evict_node_to_host + self.prefix_cache.on_match_dead_snapshot = self._rehydrate_snapshot + if self.host_tier.has_kv: + self.prefix_cache.on_evict_kv = self._on_evict_kv_to_host # The discarded hybrid tree owned donated GDN-snapshot slots; rebuild is idle-only, so # reclaim the whole LinearStatePool free-list (else those slots leak -> admission hangs). if self.is_hybrid: diff --git a/tests/kvcache/test_host_prefix_tier.py b/tests/kvcache/test_host_prefix_tier.py new file mode 100644 index 0000000000..fdb552e8e0 --- /dev/null +++ b/tests/kvcache/test_host_prefix_tier.py @@ -0,0 +1,60 @@ +"""Degrau 1 — teste standalone do HostPrefixTier (roda sem servidor; GPU tiny).""" +import torch +from types import SimpleNamespace + +from freetoken.kvcache.host_prefix_tier import HostPrefixTier + +def main(): + import os + dev = "cuda" if (os.environ.get("TIER_TEST_CPU") is None and torch.cuda.is_available()) else "cpu" + # pool fake: mesmas geometrias, mini + L, S = 4, 3 # 4 "layers", 3 slots GPU + pool = SimpleNamespace( + conv_states=torch.randn(L, S, 8, 3, dtype=torch.bfloat16, device=dev), + recurrent_states=torch.randn(L, S, 2, 4, 4, dtype=torch.float32, device=dev), + slot_states={"extra": torch.randn(L, S, 5, dtype=torch.bfloat16, device=dev)}, + ) + tier = HostPrefixTier(pool, gdn_slots_host=4, device=dev) + + # 1) round-trip D2H->H2D preserva bytes + original = [t.clone() for t in (pool.conv_states, pool.recurrent_states, pool.slot_states["extra"])] + hs = tier.write_gdn(1) # snapshot do slot 1 -> host + # suja o slot GPU pra provar que o restore vem do host + for t in (pool.conv_states, pool.recurrent_states, pool.slot_states["extra"]): + t[:, 1].zero_() + tier.read_gdn(hs, 2) # host -> slot 2 + ok = all(torch.equal(t[:, 1].cpu(), u[:, 2].cpu()) + for t, u in zip(original, (pool.conv_states, pool.recurrent_states, pool.slot_states["extra"]))) + assert ok, "snapshot divergiu no round-trip!" + print("PASS 1: snapshot GDN round-trip D2H->H2D bit-exato") + + # 2) índice: put + lookup longest-prefix page-aligned + toks = torch.arange(64 * 5, dtype=torch.int64) # 5 páginas + pages = torch.tensor([10, 11, 12, 13, 14], dtype=torch.int32) + tier.put_prefix(toks[: 64 * 4], pages[:4], hs) # registro 4 páginas + m = tier.lookup(torch.arange(64 * 5, dtype=torch.int64)) # query 5 páginas + assert m is not None and m.host_len == 64 * 4 and m.gdn_slot == hs + assert torch.equal(m.kv_pages, pages[:4]) + print("PASS 2: lookup acha o maior prefixo page-aligned") + + # 3) miss quando nada bate + assert tier.lookup(torch.full((64 * 3,), 999999, dtype=torch.int64)) is None + print("PASS 3: lookup retorna None em miss total") + + # 4) evict_host libera o slot GDN host + tier.evict_host(toks[: 64 * 4]) + assert hs in tier._free + print("PASS 4: evict_host libera slot host") + + # 5) proteção: tier novo com 2 slots, 3ª escrita levanta erro claro + tier2 = HostPrefixTier(pool, gdn_slots_host=2, device=dev) + tier2.write_gdn(0); tier2.write_gdn(0) + try: + tier2.write_gdn(0) + raise SystemExit("ERA PRA TER FALHADO") + except RuntimeError: + print("PASS 5: capacidade cheia levanta RuntimeError claro") + + print("\nTODOS OS TESTES PASSARAM ✓") + +main() diff --git a/tests/kvcache/test_tier3.py b/tests/kvcache/test_tier3.py new file mode 100644 index 0000000000..335a0c7094 --- /dev/null +++ b/tests/kvcache/test_tier3.py @@ -0,0 +1,85 @@ +"""Degrau 3 — teste standalone do KV store host do HostPrefixTier (CPU). + +Cobre: round-trip D2H->H2D de páginas KV (códigos + scales), a serpente dos ids lógicos +(a página lógica reescrita depois do evict NÃO contamina a cópia host), merge de entrada +(gdn depois kv), e drop-oldest quando o budget enche. +""" +import torch +from types import SimpleNamespace + +from freetoken.kvcache.host_prefix_tier import HostPrefixTier + + +def main(): + # pool fake: [2, L=2, P=8, page=4, H=2, D=3] códigos bf16 + scales [2, L, P*4, 2] fp32 + L, P, PAGE, H, D = 2, 8, 4, 2, 3 + kv_pool = SimpleNamespace( + _kv_buffer=torch.randn(2, L, P, PAGE, H, D, dtype=torch.bfloat16), + _scale_buffer=torch.randn(2, L, P * PAGE, H, dtype=torch.float32), + ) + state_pool = SimpleNamespace( + conv_states=torch.randn(L, 3, 8, 3, dtype=torch.bfloat16), + recurrent_states=torch.randn(L, 3, 2, 4, 4, dtype=torch.float32), + slot_states={}, + ) + page_bytes = 2 * L * PAGE * H * D * 2 + 2 * L * PAGE * H * 4 + # budget pra 3 páginas (força drop-oldest cedo) + tier = HostPrefixTier(state_pool, kv_pool, gdn_slots_host=2, + kv_budget_bytes=3 * page_bytes, device="cpu") + assert tier.has_kv and len(tier._kv_free) == 3 + + # 1) write/read round-trip: prefixo de 2 páginas (páginas lógicas 3 e 5) + kv_idx = torch.cat([torch.arange(3 * PAGE, 4 * PAGE), torch.arange(5 * PAGE, 6 * PAGE)]).to(torch.int32) + orig_codes = kv_pool._kv_buffer[:, :, 3].clone(), kv_pool._kv_buffer[:, :, 5].clone() + orig_scales = (kv_pool._scale_buffer[:, :, 3 * PAGE:4 * PAGE].clone(), + kv_pool._scale_buffer[:, :, 5 * PAGE:6 * PAGE].clone()) + hs = tier.write_kv(kv_idx) + assert hs is not None and hs.numel() == 2 + + # 2) A SERPENTE: sobrescreve as páginas lógicas 3 e 5 (novo inquilino) + kv_pool._kv_buffer[:, :, 3].fill_(99) + kv_pool._kv_buffer[:, :, 5].fill_(99) + kv_pool._scale_buffer[:, :, 3 * PAGE:6 * PAGE].fill_(99) + # lê de volta pra páginas 0 e 7 + tier.read_kv(hs, torch.tensor([0, 7], dtype=torch.int32)) + assert torch.equal(kv_pool._kv_buffer[:, :, 0], orig_codes[0]) + assert torch.equal(kv_pool._kv_buffer[:, :, 7], orig_codes[1]) + assert torch.equal(kv_pool._scale_buffer[:, :, 0 * PAGE:1 * PAGE], orig_scales[0]) + assert torch.equal(kv_pool._scale_buffer[:, :, 7 * PAGE:8 * PAGE], orig_scales[1]) + print("PASS 1: round-trip KV host + ids lógicos reescritos não contaminam") + + # 3) merge: entrada com gdn primeiro, kv depois + toks = torch.arange(64 * 2, dtype=torch.int64) + g = tier.write_gdn(1) + tier.put_prefix(toks, None, g) + tier.put_prefix(toks, hs, None) # merge: adota kv, mantém gdn + m = tier.lookup(toks) + assert m is not None and m.gdn_slot == g and torch.equal(m.kv_pages, hs) + print("PASS 2: merge gdn+kv na mesma entrada") + + # 4) drop-oldest quando o budget enche: escreve 2 páginas (resta 1), depois 2 de novo + # -> a entrada mais velha (toks) é dropada INTEIRA (KV slots + slot GDN) + kv2 = torch.cat([torch.arange(0 * PAGE, 1 * PAGE), torch.arange(2 * PAGE, 3 * PAGE)]).to(torch.int32) + toks2 = torch.arange(64 * 2, 64 * 3, dtype=torch.int64) + s2 = tier.write_kv(kv2) # precisa de 2 slots, tem 1 -> drop de toks (inteiro) + tier.put_prefix(toks2, s2, None) + assert tier.lookup(toks) is None, "o mais velho devia ter sido dropado inteiro" + assert g in tier._free, "o slot GDN do dropado devia voltar pro free-list" + kv3 = torch.arange(1 * PAGE, 3 * PAGE).to(torch.int32) + toks3 = torch.arange(64 * 3, 64 * 4, dtype=torch.int64) + s3 = tier.write_kv(kv3) # precisa de 2, tem 1 -> drop de toks2 + tier.put_prefix(toks3, s3, None) + assert tier.lookup(toks2) is None + m3 = tier.lookup(toks3) + assert m3 is not None and torch.equal(m3.kv_pages, s3) + print("PASS 3: drop-oldest remove a entrada mais fria inteira (KV + GDN)") + + # 5) evict_host libera tudo + tier.evict_host(toks3) + assert tier.lookup(toks3) is None and len(tier._kv_free) == 3 + print("PASS 4: evict_host remove e libera") + + print("\nTODOS OS TESTES PASSARAM ✓") + + +main() diff --git a/tests/kvcache/test_tree_hooks.py b/tests/kvcache/test_tree_hooks.py new file mode 100644 index 0000000000..840629310f --- /dev/null +++ b/tests/kvcache/test_tree_hooks.py @@ -0,0 +1,77 @@ +"""Degrau 2 — teste standalone dos hooks do HybridRadixCache (CPU, sem servidor).""" +import torch +from freetoken.kvcache.hybrid_radix_cache import HybridRadixCache + + +def main(): + tree = HybridRadixCache(torch.device("cpu"), 64) + evicted = [] # (prefix_len, mamba_slot) capturados pelo hook de evicção + rehydrated = [] # prefix_len reidratado no match + host_index = set() + + tree.on_evict_node = lambda node: evicted.append( + (len(tree._collect_key(node)), node.mamba_value)) + next_slot = [100] + + def fake_rehydrate(node): + n = len(tree._collect_key(node)) + if n in host_index: + rehydrated.append(n) + s = next_slot[0] + next_slot[0] += 1 + return s + return None + + tree.on_match_dead_snapshot = fake_rehydrate + + # A (4096 toks, snapshot 7) + filho B (mais 2048, snapshot 8) -> A é nó INTERNO + idsA = torch.arange(4096, dtype=torch.int32) + kvA = torch.arange(4096, dtype=torch.int32) * 10 + assert tree.insert(idsA, kvA, mamba_value=7) == (0, False) + idsB = torch.cat([idsA, torch.full((2048,), 555, dtype=torch.int32)]) + kvB = torch.arange(6144, dtype=torch.int32) * 10 + assert tree.insert(idsB, kvB, mamba_value=8) == (4096, False) + + # match enquanto tudo vivo: ganha o snapshot mais profundo (B, 6144) + m = tree.match_prefix(idsB) + assert m.cached_len == 6144 and m.mamba_value == 8 + + # evicta os 2 snapshots: A (interno -> tombstone) e B (folha -> com tier, tombstone + # in-place: o KV da folha FICA no tree, não é liberado) + er = tree.evict_mamba(2) + assert sorted(er.mamba_slots) == [7, 8], er + assert sorted(evicted) == [(4096, 7), (6144, 8)], evicted + assert er.kv_indices.numel() == 0, "com o tier, nenhuma página KV é liberada" + print("PASS 1: on_evict_node dispara; folha vira tombstone e mantém o KV:", evicted) + + # sem cópia host: match trunca pra 0 (A e B viraram tombstones) + m = tree.match_prefix(idsA) + assert m.cached_len == 0 and m.mamba_value is None + print("PASS 2: match trunca sem tier host") + + # host tem A: match(idsA) reidrata A + host_index.add(4096) + m = tree.match_prefix(idsA) + assert m.cached_len == 4096 and m.mamba_value == 100, (m.cached_len, m.mamba_value) + assert rehydrated == [4096] + print("PASS 3: match reidrata tombstone interno via tier host (slot 100)") + + # host tem B também: match(idsB) reidrata a FOLHA tombstone (mais profunda) + host_index.add(6144) + m = tree.match_prefix(idsB) + assert m.cached_len == 6144 and m.mamba_value == 101, (m.cached_len, m.mamba_value) + print("PASS 4: match reidrata folha-tombstone (KV preservado) -> 6144") + + # nós reidratados ficam vivos: segundo match não chama o hook + m2 = tree.match_prefix(idsB) + assert m2.mamba_value == 101 and rehydrated == [4096, 6144] + print("PASS 5: nós reidratados ficam vivos (sem H2D repetido)") + + # integridade de contagem: os dois snapshots revividos voltam a contar como evictable + assert tree.mamba_evictable == 2, tree.mamba_evictable + print("PASS 6: contagem mamba_evictable consistente") + + print("\nTODOS OS TESTES PASSARAM ✓") + + +main() From ff44383259515a25ce8933d14dff3175f82a4dd9 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sat, 19 Sep 2026 17:21:12 -0400 Subject: [PATCH 2/2] fix(kvcache): refuse the host kv tiers on an nvfp4 pool whose block scales they cannot mirror Upstream #499/#525 mirror KV codes plus the fp8 pool's per-token _scale_buffer; the nvfp4 pool (#408) keeps per-block scales in _block_scale_buffer, which neither the offloader nor the prefix tier copies, so --kv-host-pages or FT_PREFIX_HOST would rehydrate codes without scales. Both constructors now raise a clear error naming the knob; the GDN-snapshot-only tier (FT_GDN_HOST_TIER) never reads the KV pool and stays usable. Assisted-by: Claude Fable 5.1 (cherry picked from commit 4bb35fc34b0cb65b41584c65d7390b6c5f4f854a) --- python/freetoken/kvcache/host_prefix_tier.py | 3 ++ python/freetoken/kvcache/kv_host_offload.py | 14 +++++++ tests/kvcache/test_host_tier_nvfp4_guard.py | 44 ++++++++++++++++++++ 3 files changed, 61 insertions(+) create mode 100644 tests/kvcache/test_host_tier_nvfp4_guard.py diff --git a/python/freetoken/kvcache/host_prefix_tier.py b/python/freetoken/kvcache/host_prefix_tier.py index 0265b126ef..3b2d2d5dc3 100644 --- a/python/freetoken/kvcache/host_prefix_tier.py +++ b/python/freetoken/kvcache/host_prefix_tier.py @@ -79,6 +79,9 @@ def __init__(self, state_pool, kv_pool=None, *, gdn_slots_host: int = 64, self._kv_scales: HostBank | None = None self._kv_free: list[int] = [] if kv_pool is not None and kv_budget_bytes > 0: + from freetoken.kvcache.kv_host_offload import _reject_block_scaled_pool + + _reject_block_scaled_pool(kv_pool, "FT_PREFIX_HOST") codes = kv_pool._kv_buffer # [2, L, P, page, H, D] scales = kv_pool._scale_buffer # [2, L, P*page, H] | None (fp8) _, layers, _, page_size, heads, dim = codes.shape diff --git a/python/freetoken/kvcache/kv_host_offload.py b/python/freetoken/kvcache/kv_host_offload.py index dbb49f4814..1f8c358426 100644 --- a/python/freetoken/kvcache/kv_host_offload.py +++ b/python/freetoken/kvcache/kv_host_offload.py @@ -33,10 +33,24 @@ logger = init_logger(__name__) +def _reject_block_scaled_pool(pool, knob: str) -> None: + """The host mirrors carry codes plus the fp8 pool's per-token scales; an nvfp4 pool + (--kv-cache-dtype nvfp4) keeps per-block scales in _block_scale_buffer, which no mirror + copies, so a rehydrated page would be codes without scales. Refuse instead of decoding + garbage.""" + if getattr(pool, "_block_scale_buffer", None) is not None: + raise ValueError( + f"{knob} needs a KV pool with per-token scales (bf16 or fp8); the nvfp4 pool keeps " + "block scales the host mirror does not carry -- leave the tier off or serve with " + "--kv-cache-dtype fp8" + ) + + class KVHostOffloader: """GPU-slot LRU cache over the pinned host mirror of the logical KV page space.""" def __init__(self, pool, num_logical_pages: int, device: torch.device) -> None: + _reject_block_scaled_pool(pool, "--kv-host-pages") buf = pool._kv_buffer # [2, L, P_gpu + 1, page_size, kv_heads, head_dim] assert buf.shape[0] == 2 and buf.is_contiguous() _, num_layers, num_slots, page_size, kv_heads, head_dim = buf.shape diff --git a/tests/kvcache/test_host_tier_nvfp4_guard.py b/tests/kvcache/test_host_tier_nvfp4_guard.py new file mode 100644 index 0000000000..a793fb294a --- /dev/null +++ b/tests/kvcache/test_host_tier_nvfp4_guard.py @@ -0,0 +1,44 @@ +"""The host KV tiers (upstream #499 / #525) mirror codes plus the fp8 pool's per-token scales; +an nvfp4 pool keeps block scales they do not copy, so both must refuse it up front.""" + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.kvcache.host_prefix_tier import HostPrefixTier +from freetoken.kvcache.kv_host_offload import KVHostOffloader + + +def _kv_pool(block_scaled: bool): + return SimpleNamespace( + _kv_buffer=torch.zeros(2, 1, 3, 4, 1, 8, dtype=torch.uint8), # [2, L, P+1, page, H, D] + _scale_buffer=None, + _block_scale_buffer=torch.zeros(1) if block_scaled else None, + ) + + +def _state_pool(): + return SimpleNamespace( + conv_states=torch.zeros(2, 2, 4, 3), + recurrent_states=torch.zeros(2, 2, 2, 4), + slot_states={}, + ) + + +def test_kv_host_offloader_refuses_an_nvfp4_pool(): + with pytest.raises(ValueError, match="--kv-host-pages .* nvfp4"): + KVHostOffloader(_kv_pool(block_scaled=True), num_logical_pages=4, device=torch.device("cpu")) + + +def test_prefix_tier_refuses_kv_pages_on_an_nvfp4_pool(): + with pytest.raises(ValueError, match="FT_PREFIX_HOST .* nvfp4"): + HostPrefixTier(_state_pool(), kv_pool=_kv_pool(block_scaled=True), gdn_slots_host=2, + kv_budget_bytes=1 << 16, device="cpu") + + +def test_prefix_tier_gdn_snapshots_do_not_touch_the_kv_pool(): + # the snapshot-only tier (no KV budget) never reads the pool, nvfp4 or not + tier = HostPrefixTier(_state_pool(), kv_pool=_kv_pool(block_scaled=True), gdn_slots_host=2, + kv_budget_bytes=0, device="cpu") + assert tier.num_slots == 2 and not tier.has_kv