[AMD] Strict bit-exact SWA HiCache for DeepSeek-V4 with unified_kv: SWA-window + c4/indexer state riding across L1/L2/L3

python/sglang/srt/arg_groups/field_order.py

@@ -362,6 +362,7 @@
     "hicache_host_memory_mode",
     "hicache_ratio",
     "hicache_size",
+    "hicache_swa_offload_page_stride",
     "hicache_write_policy",
     "hicache_io_backend",
     "hicache_mem_layout",

python/sglang/srt/arg_groups/fields/memory.py

@@ -109,6 +109,14 @@ class Memory(msgspec.Struct):
         int,
         "The size of host KV cache memory pool in gigabytes. Overrides --hicache-ratio in either host memory mode.",
     ] = 0
+    hicache_swa_offload_page_stride: A[
+        int,
+        "Strict bit-exact SWA HiCache: offload one SWA sliding window every N "
+        "pages (plus the sequence tail window, always captured). 1 = per-page "
+        "(finest reuse granularity, most host memory). Larger N trades reuse "
+        "granularity for host memory; the SWA host pool is sized as "
+        "ceil(full_host_pages / N) + tail.",
+    ] = 1
     hicache_write_policy: A[
         str,
         Arg(

python/sglang/srt/environ.py

@@ -658,6 +658,20 @@ class Envs:
     SGLANG_OPT_SWA_RADIX_CACHE_COMPACT = EnvBool(False)
     SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT = EnvBool(False)
     SGLANG_OPT_SWA_RELEASE_LEAF_LOCK_AFTER_WINDOW = EnvBool(False)
+    # Strict bit-exact SWA HiCache for unified_kv (DeepSeek-V4): offload the SWA
+    # ring to host and restore it on reuse instead of reprefilling a 128-token
+    # tail. Sizing: --hicache-swa-offload-page-stride. OFF == #29417 best-effort (tail
+    # reprefill, no SWA host pool) -- cheaper for short-prefix workloads.
+    SGLANG_UNIFIED_KV_BIT_EXACT_HICACHE = EnvBool(False)
+    # Staging headroom, in absolute pages, for the c4/indexer state pools that
+    # ride the strict SWA windows. Default -1: auto = ceil(ISL/(page*stride)) + 1
+    # at per-rank batch 1, with ISL = min(200k, context_length). A row lives from
+    # capture until write_through BACKUP_HOST promote, so a long prefix can hold
+    # every strided window at once; do not scale by max_running_requests. Raise
+    # it when [SWA-HiCache] reports staging exhausted -- N concurrent prefills
+    # want roughly N times the auto value. 0 restores the legacy 1.5x-of-durable
+    # sizing.
+    SGLANG_SWA_HICACHE_STATE_STAGING_PAGES = EnvInt(-1)
 
     # ===================================================================
     # PD disaggregation runtime

python/sglang/srt/layers/attention/dsv4/compressor_v2.py

@@ -237,6 +237,41 @@ def forward_unified(
         kv_score_input = compressor.compute_kv_score(x, forward_batch)
 
         state_pool = compressor.get_state_pool(self)
+
+        # Strict SWA-HiCache c4 overlap-state capture for the UNIFIED-KV prefill
+        # path. Legacy capture (compress_extend_paged) is dead here, so without
+        # this the state staging is empty and strict reuse is rejected. Must run
+        # BEFORE compress mutates kv_score_input; gated to the extend path (decode
+        # has its own capture_compress_state_windows_decode) and ratio==4.
+        if (
+            compressor.ratio == 4
+            and forward_batch.forward_mode.is_extend()
+            and not forward_batch.forward_mode.is_target_verify()
+            and getattr(
+                self.token_to_kv_pool,
+                (
+                    "_c4_indexer_state_host_pool"
+                    if compressor.is_in_indexer
+                    else "_c4_state_host_pool"
+                ),
+                None,
+            )
+            is not None
+        ):
+            from sglang.srt.layers.attention.dsv4.compress_hip import (
+                capture_c4_state_windows_unified,
+            )
+
+            capture_c4_state_windows_unified(
+                backend=self,
+                state_pool=state_pool,
+                kv_score_input=kv_score_input,
+                forward_batch=forward_batch,
+                is_indexer=compressor.is_in_indexer,
+                layer_id=layer_id,
+                ratio=compressor.ratio,
+            )
+
         from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
             is_unified_kv_triton,
         )

python/sglang/srt/managers/schedule_batch.py

@@ -1507,6 +1507,7 @@ def init_next_round_input(
                     ),
                     req=self,
                     cow_mamba=cow_mamba,
+                    for_reuse=True,
                 )
             )
             if envs.SGLANG_RADIX_FORCE_MISS.get():
@@ -2605,6 +2606,15 @@ def prepare_for_extend(self):
             self
         )
 
+        # unified_kv SWA HiCache: restore reused sliding-window KV into each
+        # request positional ring now that req_pool_idx is known, before the
+        # first forward reads it. No-op unless a req stashed a window at
+        # load_back (see SWAComponent.restore_pending_swa_windows).
+        if self.tree_cache is not None and hasattr(
+            self.tree_cache, "restore_swa_windows"
+        ):
+            self.tree_cache.restore_swa_windows(reqs, req_pool_indices_cpu)
+
         # Set fields
         input_embeds = []
         all_replace_embeds: List[torch.Tensor] = []

python/sglang/srt/managers/schedule_policy.py

@@ -167,6 +167,7 @@ def match_prefix_for_req(
             ),
             cow_mamba=cow_mamba,
             req=req if include_req else None,
+            for_reuse=True,
         )
     )
     if envs.SGLANG_RADIX_FORCE_MISS.get():
@@ -370,7 +371,8 @@ def _compute_prefix_matches(
                             token_ids=prefix_ids,
                             extra_key=extra_key,
                             cache_salt=cache_salt,
-                        )
+                        ),
+                        for_reuse=True,
                     )
                 )
                 if envs.SGLANG_RADIX_FORCE_MISS.get():

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -56,6 +56,12 @@ class MatchPrefixParams:
     cow_mamba: bool = False
     req: Optional[Req] = None
 
+    # True for cross-request reuse matches (scheduler-side lookups against the
+    # shared radix tree). False (default) for self-match lookups such as the
+    # one in `cache_unfinished_req`, which must keep trusting the device-only
+    # validators for the request's own freshly-computed nodes.
+    for_reuse: bool = False
+
 
 @dataclasses.dataclass
 class InsertParams:
@@ -392,8 +398,12 @@ def free_kv_row(self, kv: Any, ranges: list[tuple[int, int]]) -> None:
         """Give back ascending, disjoint, half-open row-position ranges
         of the ``kv`` record's row; one call keeps a shared page freed once.
         """
-        from sglang.srt.mem_cache.common import free_kv_row_segments
+        from sglang.srt.mem_cache.common import (
+            free_kv_row_segments,
+            merge_adjacent_kv_row_ranges,
+        )
 
+        ranges = merge_adjacent_kv_row_ranges(ranges)
         row = self.req_to_token_pool.req_to_token[kv.req_pool_idx]
         free_kv_row_segments(
             self.token_to_kv_pool_allocator,

python/sglang/srt/mem_cache/common.py

@@ -51,6 +51,20 @@ def page_align_floor(length: int, page_size: int) -> int:
     return (length // page_size) * page_size
 
 
+def merge_adjacent_kv_row_ranges(
+    ranges: list[tuple[int, int]],
+) -> list[tuple[int, int]]:
+    # token-adjacent ranges can still share a page (EAGLE +1 on a page-aligned
+    # SWA branch). _page_disjoint rejects that; one segment frees the page once.
+    merged: list[tuple[int, int]] = []
+    for start, end in sorted((s, e) for s, e in ranges if s < e):
+        if merged and start <= merged[-1][1]:
+            merged[-1] = (merged[-1][0], max(merged[-1][1], end))
+        else:
+            merged.append((start, end))
+    return merged
+
+
 def free_swa_out_of_window_slots(
     req: Req,
     pre_len: int,

python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py

@@ -869,6 +869,38 @@ def unified_region_buffers(self, ratio: int) -> Tuple[List[torch.Tensor], int]:
         )
         return views, item_bytes
 
+    def swa_region_buffers(self) -> Tuple[List[torch.Tensor], int]:
+        """
+        Page-granular views of the SWA ring region [0, swa_pages) of every
+        unified_kv layer, one indexed row per sliding-window page
+        (``unified_swa_ring_size`` consecutive ring rows). Mirrors
+        ``unified_region_buffers`` so the SWA host pool transfers a whole window
+        per page and device rows match the host ``item_bytes`` in transfer_kv.
+        """
+        assert self._unified_kv, "swa_region_buffers requires unified_kv layout"
+
+        swa_pages = self.unified_kv_pool.swa_pages
+        head_dim = self.unified_kv_pool.head_dim
+        rows_per_page = self.unified_swa_ring_size
+        assert swa_pages % rows_per_page == 0, (
+            f"swa_pages {swa_pages} not a multiple of ring size {rows_per_page}"
+        )
+        num_pages = swa_pages // rows_per_page
+
+        views: List[torch.Tensor] = []
+        for buf in self.unified_kv_pool.kv_buffer:
+            page_view = (
+                buf.narrow(0, 0, swa_pages)
+                .reshape(num_pages, rows_per_page * head_dim)
+                .view(torch.uint8)
+            )
+            views.append(page_view)
+
+        item_bytes = (
+            rows_per_page * head_dim * self.unified_kv_pool.kv_buffer[0].element_size()
+        )
+        return views, item_bytes
+
     def get_state_buf_infos(self) -> Tuple[List[int], List[int], List[int]]:
         data_ptrs: List[int] = []
         data_lens: List[int] = []
@@ -898,6 +930,37 @@ def get_state_buf_infos(self) -> Tuple[List[int], List[int], List[int]]:
 
         return data_ptrs, data_lens, item_lens
 
+    def get_swa_state_coupling_infos(self) -> List[Tuple[int, int]]:
+        """FlexKV integration contract: the SWA-slot -> state-row mapping.
+
+        FlexKV registers the c4 / indexer state device buffers itself (handles
+        come from ``get_state_buf_infos``); it only needs the mapping to
+        co-place / co-fetch each state row with the SWA window it rides. Returns
+        one ``(swa_page_size, ring_size)`` per sidecar state pool, in the SAME
+        order/filtering as the state-pool entries of ``get_state_buf_infos``. For
+        a SWA device location ``swa_loc`` the coupled state row is
+        ``(swa_loc // swa_page_size) * ring_size + (swa_loc % ring_size)`` when
+        ``swa_loc >= 0`` else ``-1`` (== translate_from_swa_loc_to_state_loc).
+
+        On unified_kv the state row is owned by the request slot, not by swa_loc:
+        the authoritative mapping there is ``req_pool_idx * ring_size + pos %
+        ring_size`` (translate_from_req_position_to_state_loc), which the form
+        above only reproduces when ``swa_page_size == ring_size``. A connector
+        driving a unified_kv pool must use the req-position form.
+
+        Co-lifetime invariant the connector MUST honor: for a given (rid, B) the
+        SWA window and its coupled state rows are allocated/freed together and
+        must be co-present on fetch -- if any coupled row is missing, drop the
+        whole SWA window (recompute), never fetch with a desynced state.
+        """
+        infos: List[Tuple[int, int]] = []
+        for pools in [self.compress_state_pools, self.indexer_compress_state_pools]:
+            for pool in pools:
+                if pool is None or pool.ratio == 128:
+                    continue
+                infos.append((pool.swa_page_size, pool.ring_size))
+        return infos
+
     def get_request_state_buf_infos(
         self,
     ) -> Tuple[List[int], List[int], List[int]]:

python/sglang/srt/mem_cache/hicache_storage.py

@@ -398,6 +398,16 @@ def __init__(
         if attn_cp_size > 1:
             self.config_suffix += f"_cp{attn_cp_rank}_{attn_cp_size}"
 
+        # Strict unified-kv SWA HiCache: the SWA window and its coupled c4/indexer
+        # overlap state persist as INDEPENDENT L3 pools coupled by key (sidecar
+        # C4_STATE->SWA + TRAILING_PAGES). The SWA blob is a pure window again (no
+        # A-gather tail packing), so its byte layout differs from the older
+        # ``_swapk1`` packed blobs -- namespace the whole cache by layout version so
+        # a stale packed blob cannot be read under the new independent layout. Bump
+        # the tag on any future SWA/state L3 layout change.
+        if envs.SGLANG_UNIFIED_KV_BIT_EXACT_HICACHE.get():
+            self.config_suffix += "_swaind1"
+
         if not os.path.exists(self.file_path) and tp_rank == 0 and attn_cp_rank == 0:
             os.makedirs(self.file_path)
             logger.info(f"Created HiCacheFile storage directory at {self.file_path}")