From 5b2f3c05675a0541784fd4bd551feef7f37e32e6 Mon Sep 17 00:00:00 2001 From: Mukkesh Ganesh Date: Fri, 9 Oct 2026 18:14:36 -0700 Subject: [PATCH 1/2] feat(models): add bounded encoder swa replay behind --swa-encoder-replay --- docs/models.md | 11 +- python/freetoken/attention/dsv41_sparse.py | 17 ++- python/freetoken/core.py | 13 +- python/freetoken/engine/config.py | 4 + python/freetoken/engine/engine.py | 3 + python/freetoken/kvcache/dsv4/v41_pool.py | 13 ++ python/freetoken/kvcache/swa_radix_cache.py | 71 +++++++-- python/freetoken/models/deepseek_v41/args.py | 7 +- .../models/deepseek_v41/attention.py | 10 +- python/freetoken/scheduler/cache.py | 49 +++++- python/freetoken/scheduler/mm.py | 9 +- python/freetoken/scheduler/prefill.py | 45 ++++-- python/freetoken/scheduler/scheduler.py | 6 +- python/freetoken/server/args.py | 12 ++ tests/attention/test_dsv41_backend.py | 68 ++++++++- tests/kvcache/radix/test_swa_radix.py | 62 ++++++++ tests/models/deepseek_v41/harness.py | 10 +- tests/models/deepseek_v41/reference/model.py | 127 +++++++++++++++- .../models/deepseek_v41/test_engine_config.py | 33 +++- .../deepseek_v41/test_reference_parity.py | 141 ++++++++++++++++++ tests/scheduler/test_mm.py | 29 +++- .../test_scheduler_chunked_prefill.py | 84 +++++++++++ 22 files changed, 770 insertions(+), 54 deletions(-) diff --git a/docs/models.md b/docs/models.md index cdadc2dcb1..53e5af5188 100644 --- a/docs/models.md +++ b/docs/models.md @@ -100,10 +100,13 @@ unset on PCIe GPUs: wider grids add no bandwidth there. the segment. *Encoder* replay rebuilds the encoder SWA KV behind a prefix hit from the global KV alone (approximate; the report's fallback when the SWA KV of a hit has been evicted). *Decoder* replay runs the decoder layers on each prompt's last `n_win` tokens only; their SWA KV is never - prefix-cached and post-training simulated it. FreeToken keeps encoder SWA KV in the radix cache (no - encoder replay); `--swa-decoder-replay bounded` (default) is the report's decoder replay, with a - prefix hit stopping at least `n_win` tokens before the prompt end so those tokens are prefilled; - `exact` runs the decoder on every prompt token (reference numerics). + prefix-cached and post-training simulated it. `--swa-decoder-replay bounded` (default) is the + report's decoder replay, with a prefix hit stopping at least `n_win` tokens before the prompt end + so those tokens are prefilled; `exact` runs the decoder on every prompt token (reference + numerics). `--swa-encoder-replay bounded` (default) keeps a radix hit's full global-KV match and + replays the hit's last `n_win` tokens through the encoder when that window KV was evicted, so a + hit deeper than its live window is not truncated; `exact` truncates the hit to the live window + and re-prefills the tail exactly (reference numerics). Bounded output differs from exact by construction and does not depend on the prefill chunk size. Under expert offload the prefill time is bounded by streaming each layer's experts, so decoder replay saves decoder-layer compute, not prefill time. diff --git a/python/freetoken/attention/dsv41_sparse.py b/python/freetoken/attention/dsv41_sparse.py index baaa2cf019..401bf5a1da 100644 --- a/python/freetoken/attention/dsv41_sparse.py +++ b/python/freetoken/attention/dsv41_sparse.py @@ -39,7 +39,9 @@ class PrefillSegment: ``offset``/``n`` tile the stream; ``start_pos`` is the absolute position of the first token; ``window_floor`` bounds every query's sliding window from below (0 = the exact window; bounded - replay sets it to the start of the prompt's last window). + replay sets it to the start of the prompt's last window). ``publish_floor`` is the first + position whose GLOBAL tiers this segment may publish (0 = every token); an encoder bounded + replay re-sends the hit's last window but must leave the hit's already-published tiers alone. """ offset: int @@ -47,6 +49,7 @@ class PrefillSegment: table_idx: int start_pos: int window_floor: int = 0 + publish_floor: int = 0 @property def end(self) -> int: @@ -158,6 +161,7 @@ def __init__(self, config: ModelConfig): args = config.dsv41_args self.window_size = args.window_size self.swa_decoder_replay: str = args.swa_decoder_replay + self.swa_encoder_replay: str = args.swa_encoder_replay self.capture: DSV41CaptureData | None = None self.capture_bs: List[int] = [] self.max_graph_bs = 0 @@ -183,11 +187,18 @@ def _prefill_metadata(self, batch: Batch) -> DSV41AttnMetadata: the decoder on the same stream. Decoder SWA Bounded Replay runs it on the new tokens inside the prompt's last window ``[L - W, L)`` (``L = prompt_len``) with every query's window floored at ``L - W``: the chunks of a chunked prefill each take their part of that window (earlier rows - come from the request's private ring), and a chunk that ends before it runs no decoder row.""" + come from the request's private ring), and a chunk that ends before it runs no decoder row. + An Encoder SWA Bounded Replay (``enc_replay_lo``) extends the ENCODER pass down to the replay + start with the window floored there and the hit itself as the publish floor.""" segments: List[PrefillSegment] = [] off = 0 for r in batch.reqs: - segments.append(PrefillSegment(off, r.extend_len, r.table_idx, r.cached_len)) + replay_lo = r.enc_replay_lo + if self.swa_encoder_replay == "bounded" and 0 <= replay_lo < r.cached_len: + segments.append(PrefillSegment(off, r.extend_len, r.table_idx, replay_lo, + window_floor=replay_lo, publish_floor=r.cached_len)) + else: + segments.append(PrefillSegment(off, r.extend_len, r.table_idx, r.cached_len)) off += r.extend_len if self.swa_decoder_replay == "exact": last = torch.tensor([s.n for s in segments], dtype=torch.int32, device=self.device).cumsum_(0) - 1 diff --git a/python/freetoken/core.py b/python/freetoken/core.py index 45a763c0fe..ca9af0bde8 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -55,6 +55,11 @@ class Req: mamba_restore_src: int | None = None # on a prefix hit: tree snapshot slot to COW into the live slot (first chunk only) swa_evicted_seqlen: int = 0 # SWA radix: positions < this had their swa KV freed (slid out of window) during decode decode_batch_idx: int = 0 # SWA radix: # of decode forwards done; the proactive free_swa skips the first (overlap guard) + # Encoder SWA bounded replay (DSV41): a prefix hit deeper than its live encoder window KV + # re-sends the hit's last window through the encoder. The chunk starts there while cached_len + # (the page/commit watermark) stays at the hit; -1 = no replay. Cleared by complete_one, so a + # continuation chunk or a decode step never sees it. + enc_replay_lo: int = -1 # Set once, at the first sampled tool-call opener token (scheduler detection): the state # length just after that token (its index + 1). A client-side rewrite of the echoed tool # call diverges strictly after this point, so it is the deepest reuse boundary that @@ -84,13 +89,19 @@ def _alloc_ids_buf(self) -> None: def remain_len(self) -> int: return self.max_device_len - self.device_len + @property + def chunk_lo(self) -> int: + """Where the current chunk starts: the replay start when an encoder replay is attached.""" + return self.enc_replay_lo if self.enc_replay_lo >= 0 else self.cached_len + @property def extend_len(self) -> int: - return self.device_len - self.cached_len + return self.device_len - self.chunk_lo def complete_one(self) -> None: self.cached_len = self.device_len self.device_len += 1 + self.enc_replay_lo = -1 def append_host(self, next_token: torch.Tensor) -> None: n = self.input_ids.numel() diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index b6fa8b8949..9392bea3f4 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -37,6 +37,10 @@ class EngineConfig: # DeepSeek-V4.1 Decoder SWA Bounded Replay: "bounded" (default) runs the decoder layers on each # prompt's last window; "exact" runs them on every token. swa_decoder_replay: str = "bounded" + # DeepSeek-V4.1 Encoder SWA Bounded Replay: "bounded" (default) lets a prefix hit keep the full + # global-KV match and replay the hit's last window through the encoder when that window KV was + # evicted; "exact" truncates the hit to the windowed-safe boundary and re-prefills exactly. + swa_encoder_replay: str = "bounded" # Expert-bank host load (--expert-load): auto|serial|parallel. "auto" reads scattered # experts in parallel but falls back to serial when free RAM can't cover the banks + the # parallel reader's extra (non-reclaimable) whole-shard buffer; "serial" forces the diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 3988caf602..75354ab6b1 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -1230,11 +1230,14 @@ def _adjust_dsv41_config(config: EngineConfig, override) -> None: args.max_seq_len = config.max_seq_len args.max_batch_size = config.max_running_req + 1 # +1 dummy args.swa_decoder_replay = config.swa_decoder_replay + args.swa_encoder_replay = config.swa_encoder_replay P = args.window_size override("page_size", P) logger.info_rank0( f"DSV41 KV pages are {P}-token window pages; page_size set to {P}; SWA decoder replay: {args.swa_decoder_replay}" + ("" if args.swa_decoder_replay == "exact" else f" (prefix hits stop {P} tokens before the prompt end)") + + f"; SWA encoder replay: {args.swa_encoder_replay}" + + ("" if args.swa_encoder_replay == "exact" else f" (an evicted hit replays the {P} tokens below the hit)") ) if getattr(config, "cache_type", "radix") != "naive": override("cache_type", "swa_radix") diff --git a/python/freetoken/kvcache/dsv4/v41_pool.py b/python/freetoken/kvcache/dsv4/v41_pool.py index 36d3986ae5..6ceb304964 100644 --- a/python/freetoken/kvcache/dsv4/v41_pool.py +++ b/python/freetoken/kvcache/dsv4/v41_pool.py @@ -113,6 +113,19 @@ def prefix_replay_tokens(self) -> int: the prompt's last window from that window's encoder outputs, which no cache keeps.""" return self.P if self.private_window_layer_ids else 0 + @property + def swa_encoder_replay(self) -> str: + return self.args.swa_encoder_replay + + def encoder_replay_start(self, cached_len: int, windowed_len: int) -> int: + """The single source of truth for the bounded encoder replay's extent: a hit of + ``cached_len`` whose windowed-safe boundary ``windowed_len`` fell short of it replays + exactly the last P tokens of the hit (uniform: never a partial span, never deeper), else + there is no replay and the prefill starts at the hit as usual.""" + if self.args.swa_encoder_replay != "bounded" or windowed_len < 0 or windowed_len >= cached_len: + return cached_len + return max(0, cached_len - self.P) + # ----- private window rings ----- def is_private_window(self, layer_id: int) -> bool: return layer_id in self.private_window_layer_ids diff --git a/python/freetoken/kvcache/swa_radix_cache.py b/python/freetoken/kvcache/swa_radix_cache.py index d2cc89c26f..d248f7280a 100644 --- a/python/freetoken/kvcache/swa_radix_cache.py +++ b/python/freetoken/kvcache/swa_radix_cache.py @@ -37,11 +37,15 @@ class SWACacheHandle(BaseCacheHandle): """Lock handle for a matched SWA prefix: the windowed-safe matched node (lock target) + the reusable full-pool KV page indices. ``swa_uuid`` is the window-boundary returned by ``inc_lock``; ``CacheManager.lock`` records it on the handle (via object.__setattr__ -- the - base handle is a frozen dataclass) and ``unlock``/``dec_lock`` consume it.""" + base handle is a frozen dataclass) and ``unlock``/``dec_lock`` consume it. Under bounded + encoder replay the handle names the FULL match (``cached_len`` = its depth) and carries the + windowed-safe boundary in ``windowed_len`` (< cached_len -> the hit's window KV was evicted + and the last window is replayed); -1 otherwise.""" node: RadixTreeNode kv_indices: torch.Tensor swa_uuid: Optional[int] = None + windowed_len: int = -1 def get_matched_indices(self) -> torch.Tensor: return self.kv_indices @@ -58,6 +62,15 @@ class SWAMatch(NamedTuple): node: RadixTreeNode # the matched windowed-safe node (lock target) +class SWAFullMatch(NamedTuple): + """The full-tier match plus the windowed-safe boundary (bounded encoder replay): the global + tiers are live for [0, full_len), the window KV only up to windowed_len.""" + kv_indices: torch.Tensor # reused full-pool KV page indices for [0:full_len) + full_len: int + windowed_len: int # <= full_len; < full_len -> the hit's window KV was evicted + node: RadixTreeNode # the full-depth node (lock target) + + class SWAEvictResult(NamedTuple): kv_indices: torch.Tensor # full-pool page indices removed from the tree (-> free_slots) swa_indices: torch.Tensor # full-pool indices whose swa slots to free_swa (tombstoned + removed) @@ -85,11 +98,9 @@ def __init__(self, device: torch.device, page_size: int, sliding_window_size: in self._revives = 0 # observability: # of insert-side tombstone revives (Branches 1/2) # ---------------------------------------------------------------- match / insert - def match_prefix(self, input_ids: torch.Tensor) -> SWAMatch: - """Match the token prefix for the full layers, then truncate the reusable length to the - windowed-reuse boundary for the swa layers: the deepest node whose run of contiguous live - (non-tombstone) tokens back to the last tombstone is ``>= sliding_window_size`` (or the - path is tombstone-free to root). Mirrors sglang ``_match_prefix_helper``.""" + def _walk_prefix(self, input_ids: torch.Tensor): + """The shared match walk: the full path's values and node, plus how many leading entries + are windowed-safe (the deepest boundary with a live run >= the window).""" node = self.root value: List[torch.Tensor] = [] # path connected to root without a tombstone is always reusable -> start at +inf. @@ -129,13 +140,31 @@ def match_prefix(self, input_ids: torch.Tensor) -> SWAMatch: if match_since_tomb >= self.sliding_window_size: best_value_len = len(value) best_node = node + return value, best_value_len, best_node, node + def match_prefix(self, input_ids: torch.Tensor) -> SWAMatch: + """Match the token prefix for the full layers, then truncate the reusable length to the + windowed-reuse boundary for the swa layers: the deepest node whose run of contiguous live + (non-tombstone) tokens back to the last tombstone is ``>= sliding_window_size`` (or the + path is tombstone-free to root). Mirrors sglang ``_match_prefix_helper``.""" + value, best_value_len, best_node, _ = self._walk_prefix(input_ids) self._stamp_path(best_node) kv = torch.cat(value[:best_value_len]) if best_value_len else self.empty return SWAMatch(kv, int(kv.numel()), best_node) + def match_prefix_full(self, input_ids: torch.Tensor) -> SWAFullMatch: + """The full-tier match WITHOUT the windowed truncation, plus the windowed-safe boundary: + what bounded encoder replay needs to keep the whole global-KV hit and replay only the + evicted window. Stamps the full path (the deeper match covers the boundary's path).""" + value, best_value_len, _, node = self._walk_prefix(input_ids) + self._stamp_path(node) + kv = torch.cat(value) if value else self.empty + windowed_len = sum(t.numel() for t in value[:best_value_len]) + return SWAFullMatch(kv, int(kv.numel()), windowed_len, node) + def insert(self, input_ids: torch.Tensor, kv_indices: torch.Tensor, - swa_evicted_seqlen: int = 0, update_kv_after_len: int = 0 + swa_evicted_seqlen: int = 0, update_kv_after_len: int = 0, + replay_reused_tree_pages: bool = False ) -> Tuple[int, torch.Tensor]: """Insert the committed full KV prefix. ``kv_indices`` are the request's full-pool page indices (the swa rides along via the full->swa mapping, live where the request allocated @@ -146,11 +175,15 @@ def insert(self, input_ids: torch.Tensor, kv_indices: torch.Tensor, it are tree-shared slots and are left untouched. - ``swa_evicted_seqlen`` (= positions whose swa the request itself freed during decode): below it the request's swa is the sentinel; at/above it it is live. + - ``replay_reused_tree_pages`` (a bounded encoder replay commit): the request's row for the + replay span IS the tree's own pages, so a matched tombstone there revives in place. A matched node that is ``swa_tombstone`` (swa freed by an earlier ``evict_swa``) and extends into the fresh region is REVIVED -- not by rewriting the mapping but by adopting the request's live-swa indices into ``node.value`` and clearing the tombstone (sglang Branches - 1/2/3). Tokens of the new suffix before ``swa_evicted_seqlen`` are inserted as a tombstone; + 1/2/3). A node whose value already IS the request's span (a bounded encoder replay rewrote + the window KV on the tree's own pages) revives in place, adopting nothing and freeing + nothing. Tokens of the new suffix before ``swa_evicted_seqlen`` are inserted as a tombstone; a live (non-tombstone) leaf is always added (the free_swa ``-page_size`` margin guarantees a live tail). Returns ``(matched_prefix_len, freed_full_indices)``: the full-pool slots the caller must return to BOTH pools (displaced old tree slots + non-adopted request dups; @@ -172,9 +205,23 @@ def insert(self, input_ids: torch.Tensor, kv_indices: torch.Tensor, child = child.split_at(match_len) seg_kv = kv_indices[total:total + match_len] # request's fresh slots for this span if update_kv_after_len < total + match_len: + # An encoder bounded replay reused the tree's OWN pages for [replay_lo, hit) and + # rewrote their window KV in place, so the "request's slots" for that span ARE the + # node's value: never hand the node's own slots back as dups. Only a replay commit + # can alias, so the compare never runs on any other model's insert. + same_slots = replay_reused_tree_pages and torch.equal(child.value, seg_kv) if child.swa_tombstone: assert child.swa_ref_count == 0, "a tombstoned node cannot hold a swa lock" - if child.ref_count > 0: + if same_slots: + # In-place revive: safe even under a full lock (the inserter's own -- the + # insert runs before unlock), because unlike Branch 1 it swaps no slots + # and frees nothing: the window KV a replay rewrote on these pages is + # live, which is all the flag and the accounting record. + child.swa_tombstone = False + child.timestamp = self._tick() + self.swa_evictable += child.length + self._revives += 1 + elif child.ref_count > 0: # A full-locked reader still gathers the node's CURRENT slots through its # own row; freeing them on revive would hand live KV to the next alloc. # Keep the tombstone + the tree's value, drop the request's dup (Branch 3 @@ -206,7 +253,9 @@ def insert(self, input_ids: torch.Tensor, kv_indices: torch.Tensor, freed.append(seg_kv.clone()) else: # Matched a live node -> tree slots are canonical; drop the request's dup. - freed.append(seg_kv.clone()) + # (same_slots is only set by a replay commit re-inserting the tree's own row.) + if not same_slots: + freed.append(seg_kv.clone()) total += match_len node = child if partial: @@ -452,4 +501,4 @@ def _all_nodes(self) -> List[RadixTreeNode]: stack.extend(n.children.values()) return out -__all__ = ["SWARadixCache", "SWAMatch", "SWAEvictResult", "SWACacheHandle"] +__all__ = ["SWARadixCache", "SWAMatch", "SWAFullMatch", "SWAEvictResult", "SWACacheHandle"] diff --git a/python/freetoken/models/deepseek_v41/args.py b/python/freetoken/models/deepseek_v41/args.py index 867dcb21a9..d9af67fd80 100644 --- a/python/freetoken/models/deepseek_v41/args.py +++ b/python/freetoken/models/deepseek_v41/args.py @@ -125,6 +125,10 @@ class DeepseekV41Args: # "bounded": the decoder runs on the prompt's last window_size tokens with the sliding window # truncated there (Decoder SWA Bounded Replay, tech report sec. 3.2.2); "exact": on every token. swa_decoder_replay: Literal["bounded", "exact"] = "bounded" + # "bounded": a prefix hit deeper than its live encoder window KV keeps the full global-KV match + # and replays the hit's last window_size tokens through the encoder; "exact": the hit truncates + # to the windowed-safe boundary and re-prefills the tail exactly. + swa_encoder_replay: Literal["bounded", "exact"] = "bounded" roles: tuple[LayerRole, ...] = field(default_factory=tuple, repr=False) def __post_init__(self) -> None: @@ -205,7 +209,8 @@ def from_hf(cls, hf_config: Any) -> "DeepseekV41Args": raw = text.to_dict() else: raw = vars(text) - names = {f.name for f in cls.__dataclass_fields__.values()} - {"roles", "max_seq_len", "max_batch_size", "swa_decoder_replay"} + names = {f.name for f in cls.__dataclass_fields__.values()} - { + "roles", "max_seq_len", "max_batch_size", "swa_decoder_replay", "swa_encoder_replay"} kwargs: dict[str, Any] = {} for key, value in raw.items(): key = _HF_TO_ARGS.get(key, key) diff --git a/python/freetoken/models/deepseek_v41/attention.py b/python/freetoken/models/deepseek_v41/attention.py index 1643ddb6d1..1d3a194d5d 100644 --- a/python/freetoken/models/deepseek_v41/attention.py +++ b/python/freetoken/models/deepseek_v41/attention.py @@ -144,7 +144,11 @@ def load_state_dict(self, state_dict, *, prefix: str = "", _internal: bool = Fal def publish_prefill(self, x: torch.Tensor, segments: List[PrefillSegment]) -> None: """Full mode: compress this layer's input for every segment into the source's main KV and index keys. Runs on EVERY new token (under bounded replay the decoder source publishes for - the whole prompt even though its attention runs on the prompt's last window only).""" + the whole prompt even though its attention runs on the prompt's last window only), except a + segment's ``publish_floor``: an encoder bounded replay re-sends the hit's last window, and + those groups' global tiers are read as-is -- republishing them would rewrite the shared + rows every other reader of the hit already sees. The floor is page-aligned and a page holds + whole groups, so no group straddles it.""" assert self.compressor is not None and self.indexer is not None attn, ratio, src = self.attn, self.role.ratio, self.layer_id for seg in segments: @@ -153,6 +157,10 @@ def publish_prefill(self, x: torch.Tensor, segments: List[PrefillSegment]) -> No if self.compressor.needs_tail_carry(seg.start_pos): tail = int(attn.window_slots_of(seg.table_idx, seg.start_pos - 1, seg.start_pos).item()) latent, starts = self.compressor.forward_prefill(x[seg.offset : seg.offset + seg.n], seg.start_pos, slots, tail) + if seg.publish_floor: + keep = starts >= seg.publish_floor + if not bool(keep.all()): + latent, starts = latent[keep], starts[keep] if not latent.shape[0]: continue rows = attn.compressed_rows_of(seg.table_idx, starts, ratio) diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py index 772297e069..cdf506d96c 100644 --- a/python/freetoken/scheduler/cache.py +++ b/python/freetoken/scheduler/cache.py @@ -101,6 +101,11 @@ def match_req(self, req: PendingReq, max_len: int | None = None) -> MatchResult: ids = req.input_ids[: input_len - 1 if max_len is None else max(0, min(input_len - 1, max_len))] if self.is_swa: from freetoken.kvcache.swa_radix_cache import SWACacheHandle + if getattr(self.swa_pool, "swa_encoder_replay", "exact") == "bounded": + # bounded encoder replay: keep the FULL global-tier match and carry the windowed-safe + # boundary alongside (the admission decides whether the last window is replayed) + m = self.prefix_cache.match_prefix_full(ids) + return MatchResult(SWACacheHandle(m.full_len, m.node, m.kv_indices, windowed_len=m.windowed_len)) m = self.prefix_cache.match_prefix(ids) return MatchResult(SWACacheHandle(m.cached_len, m.node, m.kv_indices)) if self.is_hybrid: @@ -110,6 +115,14 @@ def match_req(self, req: PendingReq, max_len: int | None = None) -> MatchResult: HybridCacheHandle(m.cached_len, m.node, m.kv_indices), mamba_value=m.mamba_value) return self.prefix_cache.match_prefix(ids) + def encoder_replay_start(self, handle: BaseCacheHandle) -> int: + """The first position a hit request's chunk must (re)send: the pool's bounded-encoder-replay + rule when the hit's window KV was evicted, else the hit itself.""" + rule = getattr(self.swa_pool, "encoder_replay_start", None) if self.swa_pool is not None else None + if rule is None: + return handle.cached_len + return rule(handle.cached_len, getattr(handle, "windowed_len", -1)) + @property def available_size(self) -> int: evictable = (self.prefix_cache.full_evictable_size if (self.is_hybrid or self.is_swa) @@ -282,6 +295,29 @@ def allocate_paged(self, reqs: List[Req]) -> None: self.ensure_swa_slots(len(allocated)) self.swa_pool.alloc_swa(allocated) _write_page_table(self.page_table, allocated, allocation_info, self.page_size) + self._rebind_replay_windows(reqs) + + def _rebind_replay_windows(self, reqs: List[Req]) -> None: + """Bounded encoder replay: the hit's dead tree pages get a FRESH window slot (their global + KV stays reachable through the same full locs, so only the window binding is redone). With + page == window the replay's single top page is always dead -- a live one would make the + run window-safe and no replay would have been scheduled -- so the dead check is a guard, + not a partial-page splitter.""" + if not self.swa_paged: + return + ps = self.page_size + for req in reqs: + lo = req.enc_replay_lo + if lo < 0: + continue + hi = req.cache_handle.cached_len + pages = self.page_table[req.table_idx, lo:hi].view(-1, ps) + dead = self.swa_pool.translate_loc_from_full_to_swa(pages[:, 0]) < 0 + if bool(dead.any()): + fresh = pages[dead].reshape(-1) + if self.is_swa: + self.ensure_swa_slots(int(fresh.numel())) + self.swa_pool.alloc_swa(fresh) def cache_req(self, req: Req, *, finished: bool) -> None: if self.is_swa: @@ -429,6 +465,16 @@ def _cache_req_swa(self, req: Req, *, finished: bool) -> None: page_indices = self.page_table[req.table_idx, : req.cached_len] insert_len = align_down(req.cached_len, self.page_size) + update_after = old_handle.cached_len + replay_lo = self.encoder_replay_start(old_handle) + replay_reused_tree_pages = replay_lo < update_after + if replay_reused_tree_pages: + # A bounded encoder replay rewrote the window KV of the tree's OWN pages below the hit, + # so those nodes must reconcile (revive in place); the "fresh" region starts at the + # replay start. Force a node boundary there first so no node straddles it (a straddling + # revive would claim window KV for positions below the replay that are still dead). + update_after = replay_lo + self.prefix_cache.match_prefix(req.input_ids[:replay_lo]) freed = page_indices[:0] if insert_len > 0: # insert reconciles tombstones (revives the in-window ones by ADOPTING the request's @@ -442,7 +488,8 @@ def _cache_req_swa(self, req: Req, *, finished: bool) -> None: _, freed = self.prefix_cache.insert( req.input_ids[:insert_len], page_indices[:insert_len], swa_evicted_seqlen=req.swa_evicted_seqlen, - update_kv_after_len=old_handle.cached_len) + update_kv_after_len=update_after, + replay_reused_tree_pages=replay_reused_tree_pages) self.unlock(old_handle) self._free_swa(freed) # idempotent: revived/out-of-window slots are already sentinel -> no-op self._free(freed) diff --git a/python/freetoken/scheduler/mm.py b/python/freetoken/scheduler/mm.py index 8c656c4729..55990961c2 100644 --- a/python/freetoken/scheduler/mm.py +++ b/python/freetoken/scheduler/mm.py @@ -65,7 +65,7 @@ def plan_mm_chunk( def plan_mm_batch(reqs, encoder_cache: EncoderCache | None) -> tuple[List[MMItem], List[Tuple[int, int, int, int, int, int]], List[int], List[int]]: - """Jobs, plan, the batch rows of every gathered embedding row and, per batch token, the end of the image span holding it (0 for text), over the reqs in batch order (each spans [cached_len, device_len)).""" + """Jobs, plan, the batch rows of every gathered embedding row and, per batch token, the end of the image span holding it (0 for text), over the reqs in batch order (each spans [chunk_lo, device_len); an encoder-bounded-replay chunk re-gathers its image-span rows from the cached item).""" jobs: List[MMItem] = [] plan: List[Tuple[int, int, int, int, int, int]] = [] rows: List[int] = [] @@ -73,7 +73,8 @@ def plan_mm_batch(reqs, encoder_cache: EncoderCache | None) -> tuple[List[MMItem offset = 0 for req in reqs: if req.mm_items: - req_jobs, req_plan = plan_mm_chunk(req.uid, req.mm_items, req.cached_len, req.device_len, encoder_cache) + lo = req.chunk_lo + req_jobs, req_plan = plan_mm_chunk(req.uid, req.mm_items, lo, req.device_len, encoder_cache) jobs.extend(req_jobs) plan.extend(req_plan) for _, _, row_lo, row_hi, _, pos in req_plan: @@ -82,8 +83,8 @@ def plan_mm_batch(reqs, encoder_cache: EncoderCache | None) -> tuple[List[MMItem block_ends = [0] * sum(r.extend_len for r in reqs) for item in req.mm_items: for span_lo, span_hi in item.offsets: - for i in range(max(span_lo, req.cached_len), min(span_hi, req.device_len)): - block_ends[offset + i - req.cached_len] = span_hi + for i in range(max(span_lo, lo), min(span_hi, req.device_len)): + block_ends[offset + i - lo] = span_hi offset += req.extend_len return jobs, plan, rows, block_ends diff --git a/python/freetoken/scheduler/prefill.py b/python/freetoken/scheduler/prefill.py index de4bddfb4a..918bcd2228 100644 --- a/python/freetoken/scheduler/prefill.py +++ b/python/freetoken/scheduler/prefill.py @@ -76,6 +76,9 @@ def _try_allocate_one(self, req: PendingReq): mr = self.cache_manager.match_req(req, max_len=max(0, req.input_len - replay) if replay else None) handle = mr.cuda_handle cached_len = handle.cached_len + # bounded encoder replay: the evicted hit keeps its full global-KV match and re-sends the + # last window of it through the encoder (cached_len stays the alloc/commit watermark) + enc_replay_lo = self.cache_manager.encoder_replay_start(handle) # TODO: better estimate policy extend_len = req.input_len - cached_len estimated_size = self._kv_reservation_size( @@ -108,6 +111,8 @@ def _try_allocate_one(self, req: PendingReq): need_swa = div_ceil( min(max(extend_len, 1), self.cache_manager.sliding_window_size) + 1, ps ) * ps + if enc_replay_lo < cached_len: + need_swa += ps # the replay rebinds a fresh window page for the hit's dead top page if self.cache_manager.swa_available_size - self.reserved_swa < need_swa: return self.cache_manager.unlock(handle) @@ -130,7 +135,7 @@ def _try_allocate_one(self, req: PendingReq): linear_slot_idx = pool.alloc(1)[0] ping_pong = tuple(pool.alloc(2)) - return handle, table_idx, linear_slot_idx, ping_pong, mr.mamba_value + return handle, table_idx, linear_slot_idx, ping_pong, mr.mamba_value, enc_replay_lo def _add_one_req( self, @@ -143,8 +148,12 @@ def _add_one_req( next_track_idx: int = 0, restore_src: int | None = None, swa_evicted_seqlen: int = 0, + enc_replay_lo: int = -1, ) -> Req | None: - remain_len = pending_req.input_len - cached_len + # a replay chunk starts BELOW cached_len (the hit stays the page/commit watermark); its + # replayed tokens ride the first chunk through the encoder like new ones + lo = enc_replay_lo if 0 <= enc_replay_lo < cached_len else cached_len + remain_len = pending_req.input_len - lo chunk_size = min(self.token_budget, remain_len) if self.cache_manager.swa_paged: # Cap this chunk by the swa the pool can back this pass. swa is allocated per token in @@ -163,18 +172,23 @@ def _add_one_req( # (the extend [cached_len, cached_len+chunk) pulls div_ceil(end,ps)-div_ceil(start,ps) # fresh pages -- the partial head page was charged by the previous chunk), and reserve # that cost, not the raw token count. Degenerates to the token math at page_size==1. - max_end = (div_ceil(cached_len, ps) + max(swa_budget, 0) // ps) * ps - chunk_size = min(chunk_size, max(max_end - cached_len, 0)) + max_end = (div_ceil(lo, ps) + max(swa_budget, 0) // ps) * ps + chunk_size = min(chunk_size, max(max_end - lo, 0)) # A continuation resumes the compressor carry at its boundary, which must be # page-aligned; the token_budget leftover (unlike max_end) is not. Align the end # down when the chunk mints a continuation; no whole page -> retry next pass. # 0 <: a chunk the swa cap collapsed to 0 must NOT bail (undersized pool -- # bailing would livelock; the floor tests pin the loud failure). if 0 < chunk_size < remain_len: - aligned = align_down(cached_len + chunk_size, ps) - cached_len + aligned = align_down(lo + chunk_size, ps) - lo if aligned <= 0: return None chunk_size = aligned + if lo < cached_len and lo + chunk_size <= cached_len: + # a replay chunk must reach PAST the hit: a later chunk starting inside the + # replay region would allocate fresh pages over the tree's -- retry next pass + # (a non-replay chunk collapsed to 0 still falls through and fails loudly) + return None align = self.cache_manager.prefill_chunk_align if align > 1 and 0 < chunk_size < remain_len: # An unaligned chunk end is correct, it just loses this prompt's snapshot boundaries -- @@ -185,15 +199,15 @@ def _add_one_req( if self.keep_images_whole and pending_req.mm_items and chunk_size < remain_len: # a cut image would attend within only the part already in the cache: end the chunk before it, decided last because the caps above only move the end earlier and would undo it unit = math.lcm(self.cache_manager.page_size if self.cache_manager.swa_paged else 1, align if align > 1 else 1) - end = mm_chunk_end(pending_req.mm_items, cached_len, cached_len + chunk_size, unit) - cut = next((hi for item in pending_req.mm_items for lo, hi in item.offsets if lo < end < hi), None) - if cut is not None and self.token_budget < self.pass_budget and cut - cached_len <= self.pass_budget: + end = mm_chunk_end(pending_req.mm_items, lo, lo + chunk_size, unit) + cut = next((hi for item in pending_req.mm_items for lo_, hi in item.offsets if lo_ < end < hi), None) + if cut is not None and self.token_budget < self.pass_budget and cut - lo <= self.pass_budget: # other requests took part of this pass; a pass of its own holds the image whole return None - chunk_size = end - cached_len + chunk_size = end - lo if self.cache_manager.swa_paged: ps = self.cache_manager.page_size - self.reserved_swa += (div_ceil(cached_len + chunk_size, ps) - div_ceil(cached_len, ps)) * ps + self.reserved_swa += (div_ceil(lo + chunk_size, ps) - div_ceil(lo, ps)) * ps is_chunked = chunk_size < remain_len CLS = ChunkedReq if is_chunked else Req self.token_budget -= chunk_size @@ -204,11 +218,11 @@ def _add_one_req( pending_req.input_len + pending_req.output_len, cached_len ) # NOTE: update the tokens ids only; new pages will be allocated in the scheduler - _slice = slice(cached_len, cached_len + chunk_size) + _slice = slice(lo, lo + chunk_size) device_ids = self.table_manager.token_pool[table_idx, _slice] device_ids.copy_(_maybe_pinned(pending_req.input_ids[_slice]), non_blocking=True) req = CLS( - input_ids=pending_req.input_ids[: cached_len + chunk_size], + input_ids=pending_req.input_ids[: lo + chunk_size], table_idx=table_idx, cached_len=cached_len, output_len=pending_req.output_len, @@ -217,6 +231,7 @@ def _add_one_req( sampling_params=pending_req.sampling_params, prompt_len=pending_req.input_len, ) + req.enc_replay_lo = lo if lo < cached_len else -1 req.mm_items = pending_req.mm_items req.mrope_positions_full = pending_req.mrope_positions_full req.mrope_delta = pending_req.mrope_delta @@ -247,7 +262,7 @@ def try_add_one(self, pending_req: PendingReq) -> Req | None: ) if resource := self._try_allocate_one(pending_req): - cache_handle, table_idx, linear_slot_idx, ping_pong, restore_src = resource + cache_handle, table_idx, linear_slot_idx, ping_pong, restore_src, enc_replay_lo = resource req = self._add_one_req( pending_req=pending_req, cache_handle=cache_handle, @@ -257,6 +272,7 @@ def try_add_one(self, pending_req: PendingReq) -> Req | None: ping_pong=ping_pong, next_track_idx=0, restore_src=restore_src, + enc_replay_lo=enc_replay_lo, ) if req is None: # no aligned chunk this pass: undo the admission (a continuation keeps its @@ -329,9 +345,10 @@ def schedule_next_batch(self, prefill_budget: int) -> Batch | None: ) if pending_req.mm_items and self.encoder_cache is not None: # claim the rows every chunk of this request will gather; the entry outlives the chunks + # (a replay re-gathers the hit's dead image-span rows from this same item) for item in pending_req.mm_items: self.encoder_cache.register( - item.hash, req.uid, mm_rows_after(item, req.cache_handle.cached_len) + item.hash, req.uid, mm_rows_after(item, req.chunk_lo) ) log_new_tokens += req.extend_len if not is_continuation: diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index d057d3b908..5e08f7fb23 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -916,10 +916,10 @@ def _make_mrope_positions(batch: Batch, device: torch.device) -> torch.Tensor: out = host[:, offset : offset + length] full = req.mrope_positions_full if full is not None and req.device_len <= full.shape[1]: - out.copy_(full[:, req.cached_len : req.device_len]) + out.copy_(full[:, req.chunk_lo : req.device_len]) else: row = torch.arange( - req.cached_len + req.mrope_delta, + req.chunk_lo + req.mrope_delta, req.device_len + req.mrope_delta, dtype=torch.int32, ) @@ -935,7 +935,7 @@ def _make_positions(batch: Batch, device: torch.device) -> torch.Tensor: for req in batch.padded_reqs: length = req.extend_len torch.arange( - req.cached_len, + req.chunk_lo, req.device_len, dtype=torch.int32, out=indices_host[offset : offset + length], diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index b8d06018d8..9fa1819b2d 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -659,6 +659,18 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--swa-encoder-replay", + default=ServerArgs.swa_encoder_replay, + choices=["bounded", "exact"], + help=( + "DeepSeek-V4.1 Encoder SWA Bounded Replay. 'bounded' (default) lets a radix prefix hit keep " + "the full global-KV match and rerun the encoder over the hit's last 128 tokens when that " + "window KV was evicted; 'exact' truncates the hit to the live window and re-prefills the " + "tail exactly (the reference numerics)." + ), + ) + parser.add_argument( "--nvfp4-backend", action=_DeprecatedAlias, diff --git a/tests/attention/test_dsv41_backend.py b/tests/attention/test_dsv41_backend.py index 387742eb63..40e28a6611 100644 --- a/tests/attention/test_dsv41_backend.py +++ b/tests/attention/test_dsv41_backend.py @@ -33,11 +33,11 @@ def _ctx(pool): return ctx -def _stack(swa_decoder_replay="exact", num_pages=32, max_seq_len=8192): +def _stack(swa_decoder_replay="exact", num_pages=32, max_seq_len=8192, index_sources=()): from freetoken.attention.dsv41_sparse import DSV41SparseAttnBackend args = DeepseekV41Args(n_layers=len(RATIOS), head_dim=512, index_head_dim=128, window_size=P, compress_ratios=RATIOS, - kv_source_layers=SOURCES, swa_decoder_replay=swa_decoder_replay) + kv_source_layers=SOURCES, index_source_layers=index_sources, swa_decoder_replay=swa_decoder_replay) pool = DSV41PagedKVCache(dsv41_pool_sizes(num_pages + 1, args, 1.0, P), args, DEVICE, n_scratch=MRR + 1) pool._init_paged_state(MRR, True) pt = torch.zeros(MRR + 1, max_seq_len, dtype=torch.int32) @@ -110,6 +110,70 @@ def test_prefill_metadata_short_prompt_runs_the_decoder_everywhere(): assert not md.decoder_pad and md.last_indices.tolist() == [49] +def test_prefill_metadata_encoder_replay_extends_the_encoder_segment_only(): + """An encoder bounded replay (``enc_replay_lo`` below the hit): the ENCODER segment starts at the + replay start, floored there, with the hit as its publish floor; the decoder pass and + last_indices are exactly the no-replay ones (the admission cap keeps the prompt's last window + past the hit, so the decoder never touches the replayed rows).""" + backend, _, _ = _stack("bounded") + req = _req(2, 384, 216, prompt_len=600) # hit 384, chunk reaches the prompt end 600 + req.enc_replay_lo = 384 - P + batch = _prefill_batch([req]) + backend.prepare_metadata(batch) + md = batch.attn_metadata + assert [(s.offset, s.n, s.table_idx, s.start_pos, s.window_floor, s.publish_floor) for s in md.segments] == [ + (0, 344, 2, 256, 256, 384)] + assert [(d.offset, d.n, d.table_idx, d.start_pos, d.window_floor) for d in md.decoder_segments] == [ + (0, 128, 2, 472, 472)] + assert md.decoder_rows.tolist() == list(range(216, 344)) + assert not md.decoder_pad and md.last_indices.tolist() == [127] + + +def test_prefill_metadata_encoder_replay_under_exact_decoder_shares_the_segment(): + """Exact decoder + bounded encoder replay: the decoder runs every chunk token, sharing the + replay segment; the window_floor masks reads below the replay start, which is what makes the + shared segment safe (those slots are dead).""" + backend, _, _ = _stack("exact") + req = _req(2, 384, 216, prompt_len=600) + req.enc_replay_lo = 384 - P + batch = _prefill_batch([req]) + backend.prepare_metadata(batch) + md = batch.attn_metadata + assert [(s.start_pos, s.window_floor, s.publish_floor) for s in md.segments] == [(256, 256, 384)] + assert md.decoder_segments is md.segments and md.decoder_rows is None + assert md.last_indices.tolist() == [343] + + +def test_publish_prefill_masks_groups_below_the_publish_floor(): + """A replay segment re-sends the hit's last window; its groups below the hit (the publish floor) + must not be re-published: those shared global rows are read as-is, and the floor is group-aligned + (page-aligned floor, page % ratio == 0), so no group straddles it.""" + from freetoken.attention.dsv41_sparse import PrefillSegment + from freetoken.models.deepseek_v41.attention import DSV41Attention + + backend, pool, pt = _stack("bounded", index_sources=SOURCES) + args = backend.config.dsv41_args + get_global_ctx().attn_backend = backend # the model layer resolves its backend per access + pt[2, :512] = torch.arange(512, dtype=torch.int32) # the hit's pages + layer = DSV41Attention(args, args.roles[1]) # the ratio-2 Full source + layer._freqs = torch.ones(1024, args.rope_head_dim // 2, dtype=torch.complex64) # CPU stand-in table + layer.compressor.norm.forward = lambda t: t # the fused RMSNorm is CUDA-only; values don't matter here + layer.indexer.k_norm.forward = lambda t: t + stored = [] + backend.compressed_rows_of = lambda ti, starts, ratio: starts # row == group start (identity) + backend.store_main = lambda latent, src, rows: stored.append(("main", rows.tolist())) + backend.store_index = lambda k, src, rows: stored.append(("idx", rows.tolist())) + layer.publish_prefill(torch.randn(344, args.dim), [PrefillSegment(0, 344, 2, 256, window_floor=256, publish_floor=384)]) + assert stored == [ + ("idx", list(range(384, 600, 2))), # groups [384, 600): published, keys before main + ("main", list(range(384, 600, 2))), + ] + # without a floor every group of the segment publishes (the pre-replay behavior) + stored.clear() + layer.publish_prefill(torch.randn(88, args.dim), [PrefillSegment(0, 88, 2, 256)]) + assert stored == [("idx", list(range(256, 344, 2))), ("main", list(range(256, 344, 2)))] + + def test_private_window_layers_address_per_request_rings(): """Under bounded replay the decoder layers keep their window KV in per-request rings (slot = table row * P + pos % P), never in the radix-shared window pages; the encoder layers stay on diff --git a/tests/kvcache/radix/test_swa_radix.py b/tests/kvcache/radix/test_swa_radix.py index 4338314bdc..4345bf3e97 100644 --- a/tests/kvcache/radix/test_swa_radix.py +++ b/tests/kvcache/radix/test_swa_radix.py @@ -556,6 +556,68 @@ def test_swa_pressure_drains_every_slot_exactly_once(sess): assert sess.kv.in_use() == set() # no leak, no double free (ledger-checked) +# --------------------------------------------------------------------------- bounded encoder replay +def test_match_prefix_full_serves_the_whole_depth_with_the_windowed_boundary(swa_spec): + """``match_prefix_full`` (the bounded-encoder-replay match): the full-tier depth and its page + indices WITHOUT the windowed truncation, plus the windowed-safe boundary the truncating match + would have returned -- the two agree exactly when the path is fully live.""" + import torch + + from freetoken.kvcache.swa_radix_cache import SWARadixCache + + P, wp, W = swa_spec.page_size, _window_pages(swa_spec), swa_spec.window + dev = torch.device("cpu") + ids = torch.tensor(_seq(swa_spec, wp + 1), dtype=torch.int64) + slots = torch.arange((wp + 1) * P, dtype=torch.int32) + + live = SWARadixCache(dev, P, W) + live.insert(ids, slots) + m = live.match_prefix_full(ids) + assert m.full_len == len(ids) and m.windowed_len == m.full_len + assert m.kv_indices.tolist() == slots.tolist() + t = live.match_prefix(ids) + assert (t.cached_len, t.kv_indices.tolist()) == (m.windowed_len, m.kv_indices.tolist()) + + if wp == 1: + pytest.skip("page == window: a live run shorter than the window is below one whole page, " + "which the suffix clamp never leaves -- the DSV4 shape is covered end-to-end " + "by the reference-parity replay test") + evicted = SWARadixCache(dev, P, W) + evicted.insert(ids, slots, swa_evicted_seqlen=wp * P) # one live page on top of a tombstoned head + m = evicted.match_prefix_full(ids) + assert m.full_len == len(ids) and m.windowed_len == 0 # the run is short -> nothing windowed-safe + assert m.kv_indices.tolist() == slots.tolist() # ... yet the FULL indices are all served + assert evicted.match_prefix(ids).cached_len == m.windowed_len + + +def test_insert_revives_in_place_when_the_request_reused_the_trees_pages(swa_spec): + """The bounded-encoder-replay commit: the request's row for the replayed span IS the tree's own + value (the window KV was rewritten in place on those pages), so insert must revive a tombstone + there WITHOUT adopting or freeing anything -- freeing would hand the tree's live pages to the + next allocation, and adopting is a no-op.""" + import torch + + from freetoken.kvcache.swa_radix_cache import SWARadixCache + + P, W = swa_spec.page_size, swa_spec.window + dev = torch.device("cpu") + cache = SWARadixCache(dev, P, W) + ids = torch.tensor(_seq(swa_spec, 3), dtype=torch.int64) + slots = torch.arange(3 * P, dtype=torch.int32) + cache.insert(ids, slots, swa_evicted_seqlen=P) # page 0 tombstone, pages 1-2 live + assert cache.swa_evictable_size == 2 * P + + # the fresh region starts at the replay start (inside the tombstone); the row == the tree. + # The flag marks the commit as a replay's (the cache never probes for the alias otherwise). + matched, freed = cache.insert(ids, slots, swa_evicted_seqlen=0, update_kv_after_len=0, + replay_reused_tree_pages=True) + assert matched == 3 * P and freed.numel() == 0 # the tree's own slots are never "dups" + assert cache.swa_evictable_size == 3 * P # the tombstone revived in place + cache.check_integrity() + m = cache.match_prefix(ids) # the windowed boundary now reaches the full depth + assert m.cached_len == 3 * P and m.kv_indices.tolist() == slots.tolist() + + # --------------------------------------------------------------------------- retention def test_finish_time_restamp_soft_pins_the_prompt_window(swa_spec, make_session): """Decode never re-matches the prompt, so at finish the prompt window carries the stamp of the diff --git a/tests/models/deepseek_v41/harness.py b/tests/models/deepseek_v41/harness.py index 98b19e2de4..f6ff1b0db2 100644 --- a/tests/models/deepseek_v41/harness.py +++ b/tests/models/deepseek_v41/harness.py @@ -28,7 +28,7 @@ class TinyEngine: - def __init__(self, checkpoint: str, *, max_seq_len: int = 1024, max_running_req: int = 2, swa_decoder_replay: str = "exact", engram_table=None, quantized: bool = False): + def __init__(self, checkpoint: str, *, max_seq_len: int = 1024, max_running_req: int = 2, swa_decoder_replay: str = "exact", swa_encoder_replay: str = "bounded", engram_table=None, quantized: bool = False): self.device = torch.device("cuda") self.checkpoint = checkpoint if try_get_tp_info() is None: @@ -40,6 +40,7 @@ def __init__(self, checkpoint: str, *, max_seq_len: int = 1024, max_running_req: args.max_seq_len = max_seq_len args.max_batch_size = max_running_req + 1 args.swa_decoder_replay = swa_decoder_replay + args.swa_encoder_replay = swa_encoder_replay self.args = args quant = NoQuantConfig() if quantized: @@ -152,11 +153,12 @@ def new_request_on_prefix(self, table_idx: int, donor: Req, tokens: list[int], m uid=table_idx, sampling_params=SamplingParams(), cache_handle=None) def prefill(self, reqs: list[Req]) -> torch.Tensor: - """Run each request's ``[cached_len, device_len)`` tokens; returns ``[B, vocab]`` logits.""" + """Run each request's ``[chunk_lo, device_len)`` tokens (the replay start when an encoder + replay is attached); returns ``[B, vocab]`` logits.""" batch = Batch(reqs=reqs, phase="prefill") batch.padded_reqs = list(reqs) - batch.input_ids = torch.cat([r.input_ids[r.cached_len : r.device_len] for r in reqs]).to(self.device) - batch.positions = torch.cat([torch.arange(r.cached_len, r.device_len) for r in reqs]).to(self.device) + batch.input_ids = torch.cat([r.input_ids[r.chunk_lo : r.device_len] for r in reqs]).to(self.device) + batch.positions = torch.cat([torch.arange(r.chunk_lo, r.device_len) for r in reqs]).to(self.device) self._bind() self.backend.prepare_metadata(batch) with self.ctx.forward_batch(batch), self.model.forward_host_ctx(batch, False): diff --git a/tests/models/deepseek_v41/reference/model.py b/tests/models/deepseek_v41/reference/model.py index 593c7518c4..c29e6ff2f2 100644 --- a/tests/models/deepseek_v41/reference/model.py +++ b/tests/models/deepseek_v41/reference/model.py @@ -1,8 +1,9 @@ # Vendored from https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash (inference/model.py), MIT License # (see LICENSE next to this file). FreeToken keeps this copy as the numerical ORACLE its port is tested # against. Edits, all marked "FreeToken:", are limited to (1) local imports with the tilelang kernels -# replaced by the pure-torch transcriptions in kernel.py, and (2) forward_bounded(), a Decoder SWA -# Bounded Replay oracle built from the unmodified layer math. +# replaced by the pure-torch transcriptions in kernel.py, (2) forward_bounded(), a Decoder SWA +# Bounded Replay oracle built from the unmodified layer math, and (3) forward_replay(), the +# Encoder SWA Bounded Replay oracle built the same way. import math from contextlib import contextmanager from dataclasses import dataclass @@ -887,6 +888,67 @@ def forward_tail(self, x: torch.Tensor, s: int, x_all: torch.Tensor | None = Non o = torch.einsum("bsgd,grd->bsgr", o, wo_a) return self.wo_b(o.flatten(2)) + # FreeToken: Encoder SWA Bounded Replay oracle. Queries are the tail rows [s, s+n) with the + # sliding window truncated at s (forward_tail's window handling); a kv source compresses only + # x_pub = the rows [pub_lo, pub_lo+m) -- the global tiers below pub_lo are read from the + # caches the prefix forward() populated and are never re-published. + def forward_replay(self, x: torch.Tensor, s: int, x_pub: torch.Tensor | None, pub_lo: int): + bsz, n, _ = x.size() + win, rd, ratio = self.window_size, self.rope_head_dim, self.compress_ratio + freqs_cis = self.freqs_cis[s : s + n] + + qr = self.q_norm(self.wq_a(x)) + q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.head_dim)) + apply_rotary_emb(q[..., -rd:], freqs_cis) + + kv = self.kv_norm(self.wkv(x)) + apply_rotary_emb(kv[..., -rd:], freqs_cis) + act_quant(kv, fp8_block_size, scale_fmt, scale_dtype, True) + for i in range(max(0, n - win), n): + self.window_kv_cache[:bsz, (s + i) % win] = kv[:, i] + topk_idxs = get_window_topk_idxs(win, bsz, n, 0).to(x.device) + + if ratio: + end = s + n + compress_len = end // ratio + latent = None + if self.is_kv_source: + assert x_pub is not None, "the kv source compresses the publish span" + latent = self.compressor(x_pub, 0) + shared_attn.compress_kv = self.compress_kv_cache + if self.indexer is not None: + if self.indexer.freqs_cis is None: + self.indexer.freqs_cis = self.freqs_cis + L = pub_lo + x_pub.size(1) + k = self.indexer.k_norm(self.indexer.wk(latent)) + apply_rotary_emb(k[..., -rd:], self.freqs_cis[pub_lo : L - L % ratio : ratio]) + fp4_act_quant(k, fp4_block_size, True) + self.indexer.k_cache[:bsz, pub_lo // ratio : pub_lo // ratio + k.size(1)] = k + shared_attn.index_k = self.indexer.k_cache + if not self.is_index_source: + idxs = shared_attn.topk_idxs + elif compress_len == 0: + idxs = torch.empty(bsz, n, 0, dtype=torch.int32, device=x.device) + else: + if self.indexer.freqs_cis is None: + self.indexer.freqs_cis = self.freqs_cis + idxs = self.indexer.forward_tail(x, qr, s, n, kv.size(1)) + shared_attn.topk_idxs = idxs + if latent is not None: + L = pub_lo + x_pub.size(1) + apply_rotary_emb(latent[..., -rd:], self.freqs_cis[pub_lo : L - L % ratio : ratio]) + fp4_act_quant(latent, 16, True, scale_dtype=torch.float8_e4m3fn) + self.compress_kv_cache[:bsz, pub_lo // ratio : pub_lo // ratio + latent.size(1)] = latent + kv = torch.cat([kv, shared_attn.compress_kv[:bsz, :compress_len]], dim=1) + topk_idxs = torch.cat([topk_idxs, idxs], dim=-1) + + o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale) + apply_rotary_emb(o[..., -rd:], freqs_cis, True) + o = o.view(bsz, n, self.n_local_groups, -1) + wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1) + o = torch.einsum("bsgd,grd->bsgr", o, wo_a) + return self.wo_b(o.flatten(2)) + class Gate(nn.Module): """MoE gating. The correction bias steers expert selection only; the routing weights come from the @@ -1118,6 +1180,28 @@ def forward_tail(self, h, s: int, pre_mix, h_all=None, pre_mix_all=None): return x, ffn_pre + # FreeToken: Encoder SWA Bounded Replay oracle. ``h`` / ``pre_mix`` are the tail streams + # [s, end); a kv source publishes its global tiers from ``pub_lo`` only, and ``row_lo`` + # narrows the attention / FFN rows to [row_lo, end) (the CED decoder source). + def forward_replay(self, h, s: int, pre_mix, pub_lo: int | None = None, row_lo: int | None = None): + attn_pre, attn_post, attn_comb = self.hc_mixes(h, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base) + x_all = self.attn_norm(self.hc_pre(h, pre_mix)) + head = 0 if row_lo is None else row_lo - s + pub = 0 if pub_lo is None else pub_lo - s + residual = h[:, head:] + attn_pre, attn_post, attn_comb = attn_pre[:, head:], attn_post[:, head:], attn_comb[:, head:] + x = self.attn.forward_replay(x_all[:, head:], s + head, x_all[:, pub:], s + pub) + x = self.hc_post(x, residual, attn_post, attn_comb) + + residual = x + ffn_pre, ffn_post, ffn_comb = self.hc_mixes(x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base) + x = self.hc_pre(x, attn_pre) + x = self.ffn_norm(x) + x = self.ffn(x, None) + x = self.hc_post(x, residual, ffn_post, ffn_comb) + return x, ffn_pre + + class ParallelHead(nn.Module): def __init__(self, vocab_size: int, dim: int, norm_eps: float = 1e-6, hc_eps: float = 1e-6): super().__init__() @@ -1420,6 +1504,45 @@ def forward_bounded(self, input_ids: torch.Tensor, n_win: int): h = layer.hc_pre(h, pre_mix) return self.head(self.norm(h)) + # FreeToken: Encoder SWA Bounded Replay oracle (roadmap swa-encoder-bounded-replay). The + # prefix forward() plays the original prefill a radix hit keeps: its GLOBAL tiers are read + # as-is below the hit F = len(prefix) and never re-published. The encoder then runs one + # bounded replay over [max(0, F - n_win), N) with every query's window truncated there; the + # decoder follows forward_bounded's rule (the prompt's last window), or with ``exact_decoder`` + # runs every replay-chunk row (windows still truncated at the replay start). + @torch.inference_mode() + def forward_replay(self, prefix_ids: torch.Tensor, suffix_ids: torch.Tensor, n_win: int, + exact_decoder: bool = False): + F = prefix_ids.size(1) + self.forward(prefix_ids, 0) + ids = torch.cat([prefix_ids, suffix_ids], dim=1) + N = ids.size(1) + s = max(0, F - n_win) # the encoder replay start (every window truncated at s) + d = s if exact_decoder else max(0, N - n_win) # the decoder's first row + assert s <= d <= N + engram_hashes = self.engram_hash(ids, 0, None) if self.engram_hash is not None else None + h = self.embed(ids[:, s:]).unsqueeze(2).repeat(1, 1, self.hc_mult, 1) + pre_mix = make_identity_pre_mix(h, self.hc_mult) + dec_start = max(i for i, layer in enumerate(self.layers) if layer.attn.is_kv_source) + for i, layer in enumerate(self.layers): + if layer.engram is not None: + assert i < dec_start, "Engram sits in the encoder" + h = layer.engram(h, engram_hashes[:, s:][:, :, layer.engram.layer_hash_index, :], None) + if i < dec_start: + h, pre_mix = layer.forward_replay(h, s, pre_mix, pub_lo=F) + elif i == dec_start: + if exact_decoder: + # the CED runs every replay-chunk row, publishing from F only + h, pre_mix = layer.forward_replay(h, s, pre_mix, pub_lo=F, row_lo=s) + else: + # the CED publishes from F (rows [F, N)) and narrows to the decoder tail [d, N) + h, pre_mix = layer.forward_replay( + h[:, F - s:], F, pre_mix[:, F - s:], pub_lo=F, row_lo=d) + else: + h, pre_mix = layer.forward_tail(h, d, pre_mix) + h = layer.hc_pre(h, pre_mix) + return self.head(self.norm(h)) + @torch.inference_mode() def forward_spec(self, input_ids: torch.Tensor, main_hidden: torch.Tensor, start_pos: int = 0): h, main_x = self.mtp[0].forward_embed(main_hidden, input_ids) diff --git a/tests/models/deepseek_v41/test_engine_config.py b/tests/models/deepseek_v41/test_engine_config.py index 116209ea51..e7f1797a60 100644 --- a/tests/models/deepseek_v41/test_engine_config.py +++ b/tests/models/deepseek_v41/test_engine_config.py @@ -30,13 +30,15 @@ def test_resolution_picks_dsv41_and_the_replay_knob(tmp_path, monkeypatch, repla monkeypatch.setattr(engine, "is_sm100_family", lambda: False) monkeypatch.setattr(engine, "is_sm90_family", lambda: True) write_tiny_checkpoint(str(tmp_path)) - config = _engine_config(str(tmp_path), attention_backend="auto", moe_strategy="offload", swa_decoder_replay=replay, max_seq_len_override=2048) + config = _engine_config(str(tmp_path), attention_backend="auto", moe_strategy="offload", swa_decoder_replay=replay, + swa_encoder_replay=replay, max_seq_len_override=2048) _adjust_config(config) assert config.attention_backend == "dsv41_sparse" assert config.page_size == 128 and config.cache_type == "swa_radix" assert resolve_pool_class(config.model_config) is DSV41PagedKVCache args = config.model_config.dsv41_args assert args.swa_decoder_replay == replay and args.max_seq_len == 2048 and args.max_batch_size == config.max_running_req + 1 + assert args.swa_encoder_replay == replay assert config.max_extend_tokens == 8192 # the prefill chunk stays bounded (whole window pages) # the cache contract follows the replay mode: a bounded-mode prefix hit must leave the prompt's last # window to the prefill (the pool's prefix_replay_tokens); the history a resume reads stays one window @@ -51,6 +53,35 @@ def test_resolution_picks_dsv41_and_the_replay_knob(tmp_path, monkeypatch, repla assert DSV41PagedKVCache.min_kv_tokens(config) // 128 == 8 + 3 * config.max_running_req + 2 * (config.max_running_req + 1) + 1 +def test_the_encoder_replay_knob_is_independent_of_the_decoder_knob(tmp_path, monkeypatch): + """--swa-encoder-replay resolves onto the args on its own: it never changes the decoder-driven + cache contract (prefix_replay_tokens, private rings) or the window-pool sizing, and its rule is + the uniform last-window replay of an evicted hit.""" + from freetoken.engine import engine + + monkeypatch.setattr(engine, "is_sm100_family", lambda: False) + monkeypatch.setattr(engine, "is_sm90_family", lambda: True) + write_tiny_checkpoint(str(tmp_path)) + from freetoken.kvcache.dsv4.v41_cost_model import dsv41_pool_sizes + + for dec in ("bounded", "exact"): + for enc in ("bounded", "exact"): + config = _engine_config(str(tmp_path), attention_backend="auto", moe_strategy="offload", + swa_decoder_replay=dec, swa_encoder_replay=enc, max_seq_len_override=2048) + _adjust_config(config) + args = config.model_config.dsv41_args + assert args.swa_encoder_replay == enc and args.swa_decoder_replay == dec + pool = DSV41PagedKVCache(dsv41_pool_sizes(16, args, 1.0, 128), args, torch.device("cpu")) + assert pool.prefix_replay_tokens == (128 if dec == "bounded" else 0) # the DECODER knob owns it + assert pool.private_window_layer_ids == (tuple(range(args.decoder_start_layer, args.n_layers)) if dec == "bounded" else ()) + if enc == "bounded": + assert pool.encoder_replay_start(640, 640) == 640 # nothing evicted: no replay + assert pool.encoder_replay_start(640, 512) == 640 - 128 # evicted: the uniform last window + assert pool.encoder_replay_start(64, 0) == 0 # a hit below one window replays it all + else: + assert pool.encoder_replay_start(640, 512) == 640 # exact mode never replays + + def test_fp8_block32_dialect_reaches_the_quant_layer(tmp_path): """A V4.1-style quantization_config on the tiny checkpoint: dense linears get the 32-block e8m0 scheme, routed experts the MXFP4 scheme with a 32-wide activation block.""" diff --git a/tests/models/deepseek_v41/test_reference_parity.py b/tests/models/deepseek_v41/test_reference_parity.py index 839cdafb82..6176ae90b8 100644 --- a/tests/models/deepseek_v41/test_reference_parity.py +++ b/tests/models/deepseek_v41/test_reference_parity.py @@ -94,6 +94,12 @@ def prefill_bounded(self, ids: torch.Tensor) -> torch.Tensor: with torch.device("cuda"): return self.model.forward_bounded(ids.reshape(-1, ids.shape[-1]).cuda(), self.margs.window_size).float() + def prefill_replay(self, prefix: torch.Tensor, suffix: torch.Tensor, n_win: int, + exact_decoder: bool = False) -> torch.Tensor: + with torch.device("cuda"): + return self.model.forward_replay(prefix.view(1, -1).cuda(), suffix.view(1, -1).cuda(), + n_win, exact_decoder).float() + def decode(self, tokens: torch.Tensor, pos: int) -> torch.Tensor: with torch.device("cuda"): return self.model(tokens.view(-1, 1).cuda(), pos)[1].float() @@ -217,3 +223,138 @@ def test_bounded_replay_matches_the_reference_oracle(checkpoint): got = eng.decode([req], [nxt]) _compare(f"decode after bounded prefill, step {step}", got, want) pos += 1 + + +# ---- Encoder SWA Bounded Replay: a prefix hit keeps the FULL global-KV match and replays the +# hit's last window through the encoder (the radix fallback when that window KV was evicted). +# The logits-level drift of a P-token replay is below the parity tolerance, so the oracle check +# is paired with structural pins: the evicted page must be regenerated, the global tiers below +# the hit must be untouched, and the regenerated rows must show the floored-window semantics. + +P = 128 + + +def _window_rows(eng, req, lo, hi, layer): + ws = eng.pool.translate_full_to_window(eng.page_table[req.table_idx, lo:hi].long()) + return eng.pool.read_window(layer, ws) + + +def _attach_replay(eng, req): + """What the scheduler does for a bounded-encoder-replay hit: the radix evicted the hit's top + window page (a later alloc would hand it to someone else, so its rows are zeroed here), the + rebind gives it a fresh window slot, and the request carries the replay start.""" + F, P = req.cached_len, eng.args.window_size + locs = eng.page_table[req.table_idx, F - P : F] + ws = eng.pool.translate_full_to_window(locs.long()) + assert bool((ws >= 0).all()), "the hit's top window page must be live before the eviction" + eng.pool.free_swa(locs) + for l in range(eng.args.n_layers): # a reused page's content is arbitrary; zero it + if not eng.pool.is_private_window(l): + eng.pool.window_pool[l][ws] = 0 + eng.pool.alloc_swa(locs) + req.enc_replay_lo = F - P + + +def _global_tiers(eng, table_idx: int, upto: int): + """Dequantized global-tier rows of [0, upto) per kv source, for the bit-exact no-overwrite check.""" + out = [] + for src in eng.args.backbone_kv_sources: + ratio = eng.args.compress_ratios[src] + starts = torch.arange(0, upto // ratio, device=eng.device) * ratio + rows = torch.div(eng.page_table[table_idx, starts].long(), ratio) + out.append((eng.pool.read_main(src, rows).clone(), eng.pool.read_index(src, rows).clone())) + return out + + +def _replay_setup(checkpoint, decoder_replay: str = "bounded"): + from .harness import TinyEngine + + folder, tensors = checkpoint + text = tiny_text_config() + ref = Reference(tensors, text, max_seq_len=2048, max_batch_size=3) + eng = TinyEngine(folder, max_seq_len=2048, max_running_req=2, swa_decoder_replay=decoder_replay) + _engram_table_for(eng, tensors, text) + prefix, suffix = _tokens(640, 31), _tokens(148, 32) + donor = eng.new_request(0, prefix.tolist() + suffix.tolist()) + eng.prefill([donor]) + eng.finish_prefill([donor]) + full = prefix.tolist() + suffix.tolist() + # the oracle's prefix is the depth the scheduler hits (new_request_on_prefix's rule): capped a + # window below the end under bounded replay, the full page-aligned prompt under exact + limit = len(full) - 1 + if eng.pool.prefix_replay_tokens: + limit = min(limit, max(0, len(full) - eng.pool.prefix_replay_tokens)) + F = min(donor.device_len, limit) // P * P + want = ref.prefill_replay(torch.tensor(full[:F]), torch.tensor(full[F:]), P, + exact_decoder=decoder_replay == "exact") + return eng, ref, donor, full, want + + +def test_encoder_bounded_replay_matches_the_reference_oracle(checkpoint): + eng, ref, donor, toks, want = _replay_setup(checkpoint) + + hit = eng.new_request_on_prefix(1, donor, toks) + assert hit.cached_len == 640 and hit.extend_len == 148 + before_g = _global_tiers(eng, 1, 640) + before_w = {l: _window_rows(eng, hit, 640 - P, 640, l) for l in (0, 2)} # a window and a compressing layer + _attach_replay(eng, hit) + assert hit.extend_len == 148 + P + got = eng.prefill([hit]) + _compare("encoder bounded replay prefill", got, want) + for (m0, i0), (m1, i1) in zip(before_g, _global_tiers(eng, 1, 640)): + assert torch.equal(m0, m1) and torch.equal(i0, i1), "the replay must not rewrite global rows [0, F)" + now0, now2 = _window_rows(eng, hit, 640 - P, 640, 0), _window_rows(eng, hit, 640 - P, 640, 2) + assert bool((now0 != 0).any()) and bool((now2 != 0).any()), "the replayed page was not regenerated" + # layer 0's window KV is a pure projection of the embedding: the recompute is bit-exact + assert torch.equal(now0, before_w[0]) + # deeper rows show the floor: position 639's window ([512, 639]) is unchanged by the replay, + # position 512's ([512, 512]) lost the 127 keys below the floor the original pass had + d = (now2 - before_w[2]).abs().mean(-1) + assert d[0].item() > 5 * max(d[-1].item(), 1e-6), f"floored-window drift {d[0]:.4f} vs full-window {d[-1]:.4f}" + eng.finish_prefill([hit]) + pos = 788 + for step in range(6): + nxt = int(want.argmax(-1).item()) + want = ref.decode(torch.tensor([nxt]), pos) + got = eng.decode([hit], [nxt]) + _compare(f"decode after encoder replay, step {step}", got, want) + pos += 1 + + +def test_encoder_bounded_replay_under_exact_decoder_matches_the_reference_oracle(checkpoint): + """Exact decoder + bounded encoder (the corner the backend metadata test pins structurally): + the uncapped hit is the tree's full page-aligned depth, the decoder runs every replay-chunk + row with the window still floored at the replay start, and the CED publishes only [hit, N).""" + eng, ref, donor, toks, want = _replay_setup(checkpoint, decoder_replay="exact") + + hit = eng.new_request_on_prefix(1, donor, toks) + assert hit.cached_len == 6 * P and hit.extend_len == 20 # no replay cap: the full committed depth + _attach_replay(eng, hit) + assert hit.extend_len == 20 + P + got = eng.prefill([hit]) + _compare("exact-decoder encoder bounded replay prefill", got, want) + assert bool((_window_rows(eng, hit, 6 * P - P, 6 * P, 2) != 0).any()), "the replayed page was not regenerated" + eng.finish_prefill([hit]) + pos = 788 + for step in range(4): + nxt = int(want.argmax(-1).item()) + want = ref.decode(torch.tensor([nxt]), pos) + got = eng.decode([hit], [nxt]) + _compare(f"decode after exact-decoder encoder replay, step {step}", got, want) + pos += 1 + + +def test_encoder_bounded_replay_attaches_to_the_first_chunk_only(checkpoint): + """The replay rides chunk 1 (which must reach past the hit F); the continuation is an ordinary + chunk whose cached watermark already covers the replayed page.""" + eng, ref, donor, toks, want = _replay_setup(checkpoint) + + hit = eng.new_request_on_prefix(1, donor, toks) + _attach_replay(eng, hit) + hit.device_len = 768 # chunk 1 = [512, 768): the replay page plus [640, 768) of new tokens + first = eng.prefill([hit]) + assert first.shape == (1, VOCAB) and torch.isfinite(first).all() + assert bool((_window_rows(eng, hit, 640 - P, 640, 2) != 0).any()), "chunk 1 did not replay the dead page" + hit.cached_len, hit.device_len = 768, 788 + hit.enc_replay_lo = -1 # complete_one clears the attach for the continuation + _compare("encoder replay chunked prefill", eng.prefill([hit]), want) diff --git a/tests/scheduler/test_mm.py b/tests/scheduler/test_mm.py index 72db7e3db8..f4aeb3cadb 100644 --- a/tests/scheduler/test_mm.py +++ b/tests/scheduler/test_mm.py @@ -90,8 +90,8 @@ def test_batch_rows_follow_the_reqs_in_batch_order(): from types import SimpleNamespace # req 1: 6 tokens, image on [2, 5); req 2: chunk [4, 10) of a prompt whose image spans [3, 8) - a = SimpleNamespace(uid=1, mm_items=[_item(7, [[2, 5]])], cached_len=0, device_len=6, extend_len=6) - b = SimpleNamespace(uid=2, mm_items=[_item(8, [[3, 8]])], cached_len=4, device_len=10, extend_len=6) + a = SimpleNamespace(uid=1, mm_items=[_item(7, [[2, 5]])], cached_len=0, device_len=6, extend_len=6, chunk_lo=0) + b = SimpleNamespace(uid=2, mm_items=[_item(8, [[3, 8]])], cached_len=4, device_len=10, extend_len=6, chunk_lo=4) jobs, plan, rows, block_ends = plan_mm_batch([a, b], None) assert [j.hash for j in jobs] == [7, 8] assert plan == [(1, 7, 0, 3, 3, 2), (2, 8, 1, 5, 5, 0)] @@ -99,6 +99,31 @@ def test_batch_rows_follow_the_reqs_in_batch_order(): assert block_ends == [0, 0, 5, 5, 5, 0, 8, 8, 8, 8, 0, 0] # every image row carries its span's end in request positions +def test_an_encoder_replay_chunk_re_gathers_image_rows_from_the_cached_item(): + """A replay chunk starts below the hit (chunk_lo < cached_len): the rows of an image span + inside the replay region -- including one crossing the replay start -- are re-gathered from + the encoder cache item the admission registered at mm_rows_after(item, chunk_lo), with no + re-encode; rows and block_ends key off chunk_lo, not cached_len.""" + from types import SimpleNamespace + + cache = EncoderCache(storage="cpu") + item = _item(7, [[16, 26], [28, 34]]) # span A crosses the replay start 20; span B sits above it + req = SimpleNamespace(uid=3, mm_items=[item], cached_len=32, device_len=44, extend_len=24, chunk_lo=20) + cache.register(item.hash, req.uid, mm_rows_after(item, 20)) + cache.put(item.hash, torch.ones(16, 8)) + assert mm_rows_after(item, 20) == 12 # exactly the rows the re-gather reads + + jobs, plan, rows, block_ends = plan_mm_batch([req], cache) + assert jobs == [] # the feature is cached: the replay re-gathers, it never re-encodes + assert plan == [(3, 7, 4, 10, 16, 0), (3, 7, 10, 16, 16, 8)] + assert rows == [0, 1, 2, 3, 4, 5, 8, 9, 10, 11, 12, 13] + assert block_ends == [26] * 6 + [0, 0] + [34] * 6 + [0] * 10 + + jobs, _, _, _ = plan_mm_batch([req], None) + assert [j.hash for j in jobs] == [7] # without the cache the item would encode again + + + def test_chunk_end_never_lands_inside_an_image_span(): items = [_item(1, [[10, 20]]), _item(2, [[22, 40]])] assert mm_chunk_end(items, 0, 8, 1) == 8 # before the images: untouched diff --git a/tests/scheduler/test_scheduler_chunked_prefill.py b/tests/scheduler/test_scheduler_chunked_prefill.py index a4c7e0659a..5ba8fce8c7 100644 --- a/tests/scheduler/test_scheduler_chunked_prefill.py +++ b/tests/scheduler/test_scheduler_chunked_prefill.py @@ -237,3 +237,87 @@ def test_the_sliding_window_pool_cap_cannot_recut_an_image(): for r in batch.reqs: r.complete_one() assert ends[0] == 4900 and ends[-1] == 6000 + + +def test_an_evicted_prefix_hit_chunks_from_the_encoder_replay_start(): + """Bounded encoder replay admission on the DSV41 shape (page == window == 128): a hit whose top + window page was evicted keeps the FULL match; chunk 1 starts a window below the hit and must + reach past it, the dead page is rebound to a fresh window slot, the continuation drops the + attach, and the commit revives the replayed nodes so the next hit needs no replay.""" + from freetoken.core import SamplingParams + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.kvcache.dsv4.v41_cost_model import dsv41_pool_sizes + from freetoken.kvcache.dsv4.v41_pool import DSV41PagedKVCache + from freetoken.models.deepseek_v41.args import DeepseekV41Args + from freetoken.scheduler.cache import CacheManager + from freetoken.scheduler.decode import DecodeManager + from freetoken.scheduler.prefill import ChunkedReq, PrefillManager + from freetoken.scheduler.table import TableManager + from freetoken.scheduler.utils import PendingReq + + _setup_context() + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + P, num_pages, width = 128, 32, 2048 + args = DeepseekV41Args(n_layers=5, head_dim=512, index_head_dim=128, window_size=P, + compress_ratios=(0, 2, 2, 1, 1), kv_source_layers=(1, 3), index_source_layers=(1, 3)) + pool = DSV41PagedKVCache(dsv41_pool_sizes(num_pages + 1, args, 1.0, P), args, torch.device("cpu"), + n_scratch=MAX_RUNNING + 1) + pool._init_paged_state(MAX_RUNNING, True) + pt = torch.zeros((MAX_RUNNING + 1, width), dtype=torch.int32) + pool.attach_page_table(pt) + cm = CacheManager(num_pages=num_pages, page_size=P, page_table=pt, type="swa_radix", + swa_pool=pool, sliding_window_size=P) + tm = TableManager(max_running_reqs=MAX_RUNNING, page_table=pt) + pm = PrefillManager(cm, tm, DecodeManager(P)) + + donor_ids = torch.arange(7 * P, dtype=torch.int32) # a 896-token prompt + from freetoken.core import Req + + donor = PendingReq(UID, donor_ids, SamplingParams(max_tokens=1)) + h = cm.match_req(donor).cuda_handle + donor_req = Req(input_ids=donor_ids, table_idx=tm.allocate(), cached_len=0, output_len=1, uid=UID, + sampling_params=SamplingParams(max_tokens=1), cache_handle=h) + cm.lock(h) + cm.allocate_paged([donor_req]) + donor_req.complete_one() + cm.cache_req(donor_req, finished=True) + tm.free(donor_req.table_idx) + + # the radix evicted the donor's top window page under a full-lock holder (evict_swa's + # tombstone-in-place branch): window KV gone, full KV kept + top = cm.prefix_cache.match_prefix_full(donor_ids).node + top.swa_tombstone = True + cm.prefix_cache.swa_evictable -= top.length + cm.swa_pool.free_swa(top.value) + + pm.pending_list = [PendingReq(UID + 1, torch.cat([donor_ids, torch.arange(7 * P, 7 * P + 40, dtype=torch.int32)]), + SamplingParams(max_tokens=1))] + first = pm.schedule_next_batch(2 * P) + r1 = first.reqs[0] + assert isinstance(r1, ChunkedReq) + assert r1.cache_handle.cached_len == 6 * P and r1.cache_handle.windowed_len == 0 # the full hit, window evicted + assert r1.enc_replay_lo == 5 * P and r1.cached_len == 6 * P and r1.device_len == 7 * P # chunk 1 = [640, 896) + assert r1.extend_len == 2 * P and first.log_cached_tokens == 6 * P + cm.free_swa_out_of_window_extend(first.reqs) + cm.allocate_paged(first.reqs) + ws = pool.translate_full_to_window(pt[r1.table_idx, 5 * P : 6 * P].long()) + assert bool((ws >= 0).all()), "the replayed page was not rebound to a fresh window slot" + for r in first.reqs: + r.complete_one() + + second = pm.schedule_next_batch(2 * P) + r2 = second.reqs[0] + assert not isinstance(r2, ChunkedReq) + assert r2.enc_replay_lo == -1 and r2.cached_len == 7 * P and r2.extend_len == 40 # an ordinary final chunk + cm.allocate_paged(second.reqs) + r2.complete_one() + cm.cache_req(r2, finished=True) + cm.check_integrity() + + # the commit revived the replayed nodes in place: a later hit over the same prefix keeps the + # full match with NOTHING evicted -- no second replay + probe = PendingReq(UID + 2, torch.cat([donor_ids, torch.arange(7 * P, 7 * P + 40, dtype=torch.int32)]), + SamplingParams(max_tokens=1)) + h2 = cm.match_req(probe, max_len=len(probe.input_ids) - P).cuda_handle + assert h2.cached_len == 6 * P and h2.windowed_len == 6 * P From 39e9e6a542479b3a4d47b59f3cfb47f21bac81b3 Mon Sep 17 00:00:00 2001 From: Mukkesh Ganesh Date: Sat, 10 Oct 2026 11:27:51 -0700 Subject: [PATCH 2/2] fix(kvcache): keep full kv past swa eviction for bounded encoder replay --- python/freetoken/kvcache/swa_radix_cache.py | 14 ++++++++- python/freetoken/scheduler/cache.py | 4 ++- tests/kvcache/radix/test_swa_radix.py | 34 +++++++++++++++++++++ 3 files changed, 50 insertions(+), 2 deletions(-) diff --git a/python/freetoken/kvcache/swa_radix_cache.py b/python/freetoken/kvcache/swa_radix_cache.py index d248f7280a..f48879fc1f 100644 --- a/python/freetoken/kvcache/swa_radix_cache.py +++ b/python/freetoken/kvcache/swa_radix_cache.py @@ -77,11 +77,17 @@ class SWAEvictResult(NamedTuple): class SWARadixCache: - def __init__(self, device: torch.device, page_size: int, sliding_window_size: int) -> None: + def __init__(self, device: torch.device, page_size: int, sliding_window_size: int, + keep_full_on_swa_evict: bool = False) -> None: assert sliding_window_size > 0, "SWARadixCache requires a positive sliding window" self.device = device self.page_size = page_size self.sliding_window_size = sliding_window_size + # bounded encoder replay: evicting a free leaf's swa KV tombstones it in place (full KV + # kept) instead of freeing both pools and unlinking, so the full-depth match can still + # walk it and replay only the window KV (#643 lane; evict_full still reclaims it by the + # full pool's own LRU). + self.keep_full_on_swa_evict = keep_full_on_swa_evict self.key_fn = _get_key_fn(page_size) # Page-index dtype: matches the engine's int32 page_table / free_slots so an empty result # (returned by insert/evict/match) never promotes them to int64 when concatenated. @@ -375,6 +381,12 @@ def evict_swa(self, num_tokens: int) -> SWAEvictResult: if node.swa_tombstone or node.swa_ref_count != 0 or node.is_root(): continue if node.is_leaf() and node.ref_count == 0: + if self.keep_full_on_swa_evict: + swa.append(node.value) # tombstone in place: full KV kept for replay + self.swa_evictable -= node.length + freed += node.length + node.swa_tombstone = True + continue kv.append(node.value) swa.append(node.value) self.full_evictable -= node.length diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py index cdf506d96c..aa8e80911c 100644 --- a/python/freetoken/scheduler/cache.py +++ b/python/freetoken/scheduler/cache.py @@ -92,7 +92,9 @@ def _make_prefix_cache(self, device, page_size, type): return HybridRadixCache(device, page_size) if type == "swa_radix": from freetoken.kvcache.swa_radix_cache import SWARadixCache - return SWARadixCache(device, page_size, self.sliding_window_size) + keep_full = getattr(self.swa_pool, "swa_encoder_replay", "exact") == "bounded" + return SWARadixCache(device, page_size, self.sliding_window_size, + keep_full_on_swa_evict=keep_full) return create_prefix_cache(device=device, type=type, page_size=page_size) def match_req(self, req: PendingReq, max_len: int | None = None) -> MatchResult: diff --git a/tests/kvcache/radix/test_swa_radix.py b/tests/kvcache/radix/test_swa_radix.py index 4345bf3e97..aaf0f857a4 100644 --- a/tests/kvcache/radix/test_swa_radix.py +++ b/tests/kvcache/radix/test_swa_radix.py @@ -690,3 +690,37 @@ def test_trim_head_swa_skips_locked_leaf_and_tombstoned_nodes(sess): sess.do_unlock(held) sess.check() + + +def test_evict_swa_keeps_full_kv_for_replay_when_flagged(swa_spec): + """Bounded encoder replay retention: with ``keep_full_on_swa_evict`` a free leaf's swa + eviction tombstones it in place (full KV stays walkable for ``match_prefix_full``); the + default keeps main's behavior -- free both pools, unlink -- so only replay mode retains.""" + import torch + + from freetoken.kvcache.swa_radix_cache import SWARadixCache + + P, W = swa_spec.page_size, swa_spec.window + wp = _window_pages(swa_spec) + dev = torch.device("cpu") + ids = torch.tensor(_seq(swa_spec, wp + 1), dtype=torch.int64) + slots = torch.arange((wp + 1) * P, dtype=torch.int32) + + keep = SWARadixCache(dev, P, W, keep_full_on_swa_evict=True) + keep.insert(ids, slots) + r = keep.evict_swa(keep.swa_evictable) + assert r.kv_indices.numel() == 0, "replay retention must not surrender full-pool pages" + m = keep.match_prefix_full(ids) + assert m.full_len == len(ids) + assert m.kv_indices.tolist() == slots.tolist() + assert m.windowed_len == 0 + assert keep.match_prefix(ids).cached_len == 0 + r2 = keep.evict_full(keep.full_evictable) + assert sorted(r2.kv_indices.tolist()) == sorted(slots.tolist()), \ + "the full pool's own LRU must still reclaim the tombstoned leaf" + + default = SWARadixCache(dev, P, W) + default.insert(ids, slots) + r3 = default.evict_swa(default.swa_evictable) + assert sorted(r3.kv_indices.tolist()) == sorted(slots.tolist()), "main frees both pools" + assert default.match_prefix_full(ids).full_len == 0