feat(models): add bounded encoder swa replay behind --swa-encoder-replay

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.

python/freetoken/attention/dsv41_sparse.py

@@ -39,14 +39,17 @@ 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
     n: int
     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

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()

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

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")

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

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,17 +62,32 @@ 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)
 
 
 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.
@@ -85,11 +104,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 +146,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 +181,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 +211,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 +259,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:
@@ -326,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
@@ -452,4 +513,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"]