fix(unified-memory): four boot/correctness fixes on the hybrid model paths (#35154)

Co-authored-by: Caihua Li <caihua.li@bytedance.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>

python/sglang/kernels/jit/csrc/inkling/causal_conv1d.cuh

@@ -175,7 +175,7 @@ struct CausalConv1dKernel {
     // x may be a non-contiguous row view (stride_t arbitrary) but must be
     // channel-contiguous. cache_mask is torch-bool (verify shape/device only).
     TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(x);
-    TensorMatcher({-1, Km1, D}).with_dtype<DType>().with_device(dev).verify(cache);
+    TensorMatcher({-1, Km1, D}).with_strides({-1, -1, 1}).with_dtype<DType>().with_device(dev).verify(cache);
     TensorMatcher({NS}).with_dtype<int64_t>().with_device(dev).verify(safe_idx);
     TensorMatcher({NS, 1, 1}).with_device(dev).verify(cache_mask);
     TensorMatcher({D, Wd}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(weight);

python/sglang/kernels/jit/csrc/inkling/draft_extend_sconv.cuh

@@ -114,7 +114,7 @@ struct DraftExtendSconvKernel {
     W1s.set_value(W1);
 
     TensorMatcher({BT, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(hidden);
-    TensorMatcher({-1, W1s, D}).with_dtype<DType>().with_device(dev).verify(cache);
+    TensorMatcher({-1, W1s, D}).with_strides({-1, -1, 1}).with_dtype<DType>().with_device(dev).verify(cache);
     TensorMatcher({B}).with_dtype<int32_t>().with_device(dev).verify(cache_indices);
     TensorMatcher({B}).with_dtype<int32_t>().with_device(dev).verify(num_accepted);
     RuntimeCheck(sizeof(DType) == 2, "draft_extend: bf16x2 kernel requires 16-bit dtype");

python/sglang/kernels/jit/csrc/inkling/fused_decode_update.cuh

@@ -149,7 +149,7 @@ struct FusedDecodeUpdateKernel {
     W1s.set_value(W - 1);
 
     TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(x);
-    TensorMatcher({-1, W1s, D}).with_dtype<DType>().with_device(dev).verify(cache);
+    TensorMatcher({-1, W1s, D}).with_strides({-1, -1, 1}).with_dtype<DType>().with_device(dev).verify(cache);
     TensorMatcher({T}).with_dtype<int32_t>().with_device(dev).verify(cache_indices);
     TensorMatcher({T}).with_device(dev).verify(cache_mask);
     TensorMatcher({D, Wd}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(weight);

python/sglang/kernels/jit/csrc/inkling/gather_scatter_sconv.cuh

@@ -75,7 +75,7 @@ struct GatherScatterSconvKernel {
     W1s.set_value(W1);
 
     TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(hidden);
-    TensorMatcher({-1, W1s, D}).with_dtype<DType>().with_device(dev).verify(cache);
+    TensorMatcher({-1, W1s, D}).with_strides({-1, -1, 1}).with_dtype<DType>().with_device(dev).verify(cache);
     TensorMatcher({B, W1s}).with_dtype<int32_t>().with_device(dev).verify(track_idx);
     TensorMatcher({B}).with_device(dev).verify(mask);
     TensorMatcher({B}).with_dtype<int64_t>().with_device(dev).verify(dst);

python/sglang/kernels/jit/csrc/inkling/inkling_ar_fused_decode.cuh

@@ -627,7 +627,7 @@ struct ArSconvNormKernel {
     TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(residual_out);
     TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(hs_out);
     TensorMatcher({D}).with_dtype<DType>().with_device(dev).verify(norm_weight);
-    TensorMatcher({-1, W1s, D}).with_dtype<DType>().with_device(dev).verify(cache);
+    TensorMatcher({-1, W1s, D}).with_strides({-1, -1, 1}).with_dtype<DType>().with_device(dev).verify(cache);
     TensorMatcher({T}).with_dtype<int32_t>().with_device(dev).verify(cache_indices);
     TensorMatcher({T}).with_device(dev).verify(cache_mask);
     TensorMatcher({D, Wd}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(conv_weight);
@@ -752,7 +752,7 @@ struct ArSconvNormVerifyKernel {
     TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(residual_out);
     TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(hs_out);
     TensorMatcher({D}).with_dtype<DType>().with_device(dev).verify(norm_weight);
-    TensorMatcher({-1, W1s, D}).with_dtype<DType>().with_device(dev).verify(cache);
+    TensorMatcher({-1, W1s, D}).with_strides({-1, -1, 1}).with_dtype<DType>().with_device(dev).verify(cache);
     TensorMatcher({B}).with_dtype<int32_t>().with_device(dev).verify(cache_indices);
     TensorMatcher({B}).with_device(dev).verify(cache_mask);
     TensorMatcher({D, Wd}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(conv_weight);

python/sglang/kernels/jit/csrc/inkling/update_sconv_cache.cuh

@@ -106,7 +106,7 @@ struct UpdateSconvCacheKernel {
     // x channel-contiguous (may be a non-contiguous row view); cache contiguous
     // [slots, W1, D]. cache_indices/qsl int32, has_state torch-bool (shape/device only).
     TensorMatcher({T, D}).with_strides({-1, 1}).with_dtype<DType>().with_device(dev).verify(x);
-    TensorMatcher({-1, W1s, D}).with_dtype<DType>().with_device(dev).verify(cache);
+    TensorMatcher({-1, W1s, D}).with_strides({-1, -1, 1}).with_dtype<DType>().with_device(dev).verify(cache);
     TensorMatcher({B}).with_dtype<int32_t>().with_device(dev).verify(cache_indices);
     TensorMatcher({B}).with_device(dev).verify(has_state);
     TensorMatcher({-1}).with_dtype<int32_t>().with_device(dev).verify(qsl);

python/sglang/srt/environ.py

@@ -530,8 +530,6 @@ class Envs:
     # fall back to the per-free eager compaction. Used for production
     # A/B and quick rollback. Default False (lazy compaction on).
     SGLANG_DISABLE_LAZY_COMPACTION = EnvBool(False)
-    # Sort the multi-ended allocator's free list after a merge (perf A/B knob).
-    SGLANG_SORT_FREE_LIST_AFTER_MERGE = EnvBool(False)
     # Periodically log lazy-compaction stats per sub-pool (observability only).
     SGLANG_LOG_LAZY_COMPACTION_STATS = EnvBool(False)
     SGLANG_LOG_LAZY_COMPACTION_STATS_INTERVAL_SEC = EnvInt(30)
@@ -1724,6 +1722,9 @@ _DEPRECATED_ENVS: Dict[str, _DeprecatedEnv] = {
     # Superseded by the unified JIT per_token_group_quant, the default CUDA path.
     "SGLANG_OPT_USE_JIT_PER_TOKEN_GROUP_QUANT": _DeprecatedEnv(),
     "SGLANG_MASKED_GEMM_FAST_ACT": _DeprecatedEnv(),
+    # The unified free list is kept unsorted between flushes by design; the
+    # sort-after-merge A/B knob never left its off default and is gone.
+    "SGLANG_SORT_FREE_LIST_AFTER_MERGE": _DeprecatedEnv(),
     "SGLANG_OPT_SWA_EVICT_DROP_PAGE_MARGIN": _DeprecatedEnv(),
     # sconv-family kernels always use the CUDA-JIT ports when supported; no toggle.
     "SGLANG_OPT_USE_CUDA_SCONV": _DeprecatedEnv(),

python/sglang/srt/layers/attention/aiter_backend.py

@@ -1030,7 +1030,6 @@ class AiterAttnBackend(AttentionBackend):
             req_pool_indices=forward_batch.req_pool_indices,
             seq_lens=forward_batch.seq_lens,
             seq_lens_sum=None if in_capture else forward_batch.seq_lens_sum,
-            encoder_lens=forward_batch.encoder_lens,
             forward_mode=forward_batch.forward_mode,
             spec_info=forward_batch.spec_info,
             seq_lens_cpu=seq_lens_cpu,
@@ -1667,7 +1666,6 @@ class AiterAttnBackend(AttentionBackend):
         req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
         seq_lens_sum: int,
-        encoder_lens: Optional[torch.Tensor],
         forward_mode: ForwardMode,
         spec_info: Optional[SpecInput],
         seq_lens_cpu: Optional[torch.Tensor],

python/sglang/srt/layers/attention/flashattention_backend.py

@@ -533,7 +533,6 @@ class FlashAttentionBackend(AttentionBackend):
                 bs=bs,
                 req_pool_indices=req_pool_indices,
                 seq_lens=seq_lens,
-                seq_lens_sum=None,
                 encoder_lens=encoder_lens,
                 forward_mode=forward_mode,
                 spec_info=spec_info,
@@ -579,7 +578,6 @@ class FlashAttentionBackend(AttentionBackend):
                 bs=bs,
                 req_pool_indices=req_pool_indices,
                 seq_lens=seq_lens,
-                seq_lens_sum=forward_batch.seq_lens_sum,
                 encoder_lens=encoder_lens,
                 forward_mode=forward_mode,
                 spec_info=spec_info,
@@ -2720,7 +2718,6 @@ class FlashAttentionBackend(AttentionBackend):
         bs: int,
         req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
-        seq_lens_sum: int,
         encoder_lens: Optional[torch.Tensor],
         forward_mode: ForwardMode,
         spec_info: Optional[SpecInput],

python/sglang/srt/mem_cache/multi_ended_allocator.py

@@ -25,7 +25,7 @@ from __future__ import annotations
 import inspect
 import logging
 import os
-from typing import Callable, Dict, List, Optional, Set, Tuple
+from typing import Callable, Dict, List, Optional, Sequence, Set, Tuple
 
 import torch
 from torch.profiler import record_function
@@ -47,10 +47,6 @@ from sglang.srt.utils.common import get_num_new_pages, next_power_of_2
 logger = logging.getLogger(__name__)
 
 
-# OFF (default): cat unsorted, `_flush` sorts once. ON: sort after each cat.
-_SORT_FREE_LIST_AFTER_MERGE = envs.SGLANG_SORT_FREE_LIST_AFTER_MERGE.get()
-
-
 import atexit
 import signal
 import time as _time_mod  # local alias so tests can patch
@@ -292,6 +288,10 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
         else:
             self.free_virtual_ids = None
         self.free_group = None
+        # Segment frees buffer page REPRESENTATIVES here, not whole token
+        # ranges: `torch.cat` of the ranges destroys the per-segment shape the
+        # stride derivation needs, forcing the position-less dedup back on.
+        self.free_page_reps_group: Optional[List[torch.Tensor]] = None
         self._inverse_history.clear()
         self._free_phys_pages = torch.empty(0, dtype=torch.int64, device=self.device)
         self._pending_reuse.clear()
@@ -452,16 +452,8 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
                     return None
 
             if n_drain > 0:
-                if _SORT_FREE_LIST_AFTER_MERGE:
-                    if self.grow_direction == "up":
-                        drained_t = self._free_phys_pages[:n_drain]
-                        self._free_phys_pages = self._free_phys_pages[n_drain:]
-                    else:
-                        drained_t = self._free_phys_pages[-n_drain:].flip(0)
-                        self._free_phys_pages = self._free_phys_pages[:-n_drain]
-                else:
-                    drained_t = self._free_phys_pages[:n_drain]
-                    self._free_phys_pages = self._free_phys_pages[n_drain:]
+                drained_t = self._free_phys_pages[:n_drain]
+                self._free_phys_pages = self._free_phys_pages[n_drain:]
             else:
                 drained_t = None
 
@@ -958,10 +950,17 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
 
     # -- free with eager compaction --
 
-    def free(self, free_index: torch.Tensor) -> None:
+    def free(
+        self, free_index: torch.Tensor, *, _pages: Optional[torch.Tensor] = None
+    ) -> None:
         """Free virtual TOKEN ids: recover virtual PAGE ids, un-map v2p/p2v,
         (if id-owner) recycle the page ids, trigger eager compaction.
 
+        `_pages` carries virtual PAGE ids already derived by `free_segment`
+        from `start_pos` arithmetic; when given, the data-dependent dedup is
+        skipped. Dropped on the free-group path, which has its own
+        representative buffer.
+
         `free_index` is token-granular and need not be page-aligned. EAGER mode
         drops one `wait_stream(forward_stream)` barrier so v2p/p2v writes and the
         compaction move serialize with the in-flight forward. LAZY mode needs no
@@ -976,7 +975,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
                 self.free_group.append(self._copy_for_free_group(free_index))
                 return
             if self.lazy_compaction:
-                self._free_lazy(free_index)
+                self._free_lazy(free_index, pages=_pages)
                 return
             # --- EAGER path ---
             # Near-no-op in normal mode (sampling's CPU sync already drained
@@ -986,8 +985,12 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
                 with record_function("MultiEndedAlloc.free.wait_stream"):
                     torch.cuda.current_stream().wait_stream(self.forward_stream)
             with record_function("MultiEndedAlloc.free.v2p_lookup"):
-                free_v_pages = torch.unique(
-                    free_index.detach().to(torch.int64) // self.page_size
+                free_v_pages = (
+                    _pages
+                    if _pages is not None
+                    else torch.unique(
+                        free_index.detach().to(torch.int64) // self.page_size
+                    )
                 )
                 freed_p_pages = self.virtual_to_physical[free_v_pages]
             with record_function("MultiEndedAlloc.free.sync_check"):
@@ -996,12 +999,51 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
                     self._raise_stale_slot_assertion(
                         free_v=free_v_pages, freed_p=freed_p_pages
                     )
-            self.virtual_to_physical[free_v_pages] = -1
+            self.virtual_to_physical.index_fill_(0, free_v_pages, -1)
             if self.is_id_owner:
                 self.free_virtual_ids = torch.cat([self.free_virtual_ids, free_v_pages])
             self._compact_pending(freed_p_pages)
 
-    def _free_lazy(self, free_index: torch.Tensor) -> None:
+    def _page_reps_pieces(
+        self, free_index: torch.Tensor, start_pos: int
+    ) -> Tuple[torch.Tensor, ...]:
+        """Page-representative TOKEN slices of one kv-row segment.
+
+        Mirrors `PagedTokenToKVPoolAllocator.free_segment`: a page's tokens sit
+        consecutively in the kv row, so with `start_pos` known on the host the
+        representatives are stride slices -- no `torch.unique`, whose
+        data-dependent output shape forces a device sync.
+
+        Exact for any segment shape: a partial head page is the `[:1]` term, a
+        partial tail page the final stride step.
+        """
+        ps = self.page_size
+        offset = start_pos % ps
+        if offset == 0:
+            return (free_index[::ps],)
+        return (free_index[:1], free_index[ps - offset :: ps])
+
+    def free_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None:
+        """Fixed-shape counterpart of `free()`; see `_page_reps_pieces`.
+
+        Contract: see base; a page must be freed by only one call per group.
+        """
+        if free_index is None or free_index.numel() == 0:
+            return
+        if self.page_size == 1:
+            # token == page: nothing to dedup, the plain path is already exact.
+            self.free(free_index)
+            return
+        pieces = self._page_reps_pieces(free_index.detach().to(torch.int64), start_pos)
+        if self.free_page_reps_group is None:
+            reps = pieces[0] if len(pieces) == 1 else torch.cat(pieces)
+            self.free(reps, _pages=reps // self.page_size)
+        else:
+            self.free_page_reps_group.extend(pieces)
+
+    def _free_lazy(
+        self, free_index: torch.Tensor, pages: Optional[torch.Tensor] = None
+    ) -> None:
         """Lazy free path: disjoint-element scatters + ONE `torch.cat` onto
         `_free_phys_pages`. No sort, no boundary absorb, no watermark mutation,
         no D2H sync. Boundary absorption is deferred to `_flush`.
@@ -1015,24 +1057,31 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
         with record_function("MultiEndedAlloc._free_lazy"):
             with record_function("MultiEndedAlloc._free_lazy.v2p_lookup"):
                 free_v_pages_raw = free_index.detach().to(torch.int64)
-                if self.page_size == 1:
+                if pages is not None:
+                    # `free_segment` already derived these by stride slicing.
+                    free_v_pages = pages
+                elif self.page_size == 1:
                     free_v_pages = free_v_pages_raw
                 else:
                     free_v_pages = torch.unique(free_v_pages_raw // self.page_size)
                 freed_p_pages = self.virtual_to_physical[free_v_pages]
             # Disjoint-element scatters — no barrier (a freed v has no live reader;
             # per-element scatter writes are atomic).
-            self.virtual_to_physical[free_v_pages] = -1
-            self.physical_to_virtual[freed_p_pages] = -1
+            # `index_fill_`, NOT `t[idx] = -1`: the scalar form makes torch
+            # materialise -1 as a CPU tensor and copy it H2D, and a pageable
+            # H2D copy is host-BLOCKING -- the scheduler parks behind the
+            # in-flight forward until the stream drains (~16 ms per free on an
+            # 8192-token prefill). `index_fill_` takes the scalar through the
+            # ATen Scalar overload: one device kernel, no host sync.
+            self.virtual_to_physical.index_fill_(0, free_v_pages, -1)
+            self.physical_to_virtual.index_fill_(0, freed_p_pages, -1)
             if self.is_id_owner:
                 self.free_virtual_ids = torch.cat([self.free_virtual_ids, free_v_pages])
             self._free_phys_pages = torch.cat([self._free_phys_pages, freed_p_pages])
-            if _SORT_FREE_LIST_AFTER_MERGE:
-                self._free_phys_pages, _ = torch.sort(self._free_phys_pages)
             self.live_page_count -= int(freed_p_pages.shape[0])
 
     def _release_phys_pages_batch(self, pages: torch.Tensor) -> None:
-        """Cat `pages` onto `_free_phys_pages` (+ optional sort). Called by `_flush`
+        """Cat `pages` onto `_free_phys_pages`. Called by `_flush`
         at END to merge event-fired compaction-srcs (`released_fired`) AFTER the
         trailing dst-slice, keeping `_free_phys_pages == holes_cpu` during the walk.
 
@@ -1044,8 +1093,6 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
         self._stats_n_release_batch += 1
         with record_function("MultiEndedAlloc._release_phys_pages_batch"):
             self._free_phys_pages = torch.cat([self._free_phys_pages, pages])
-            if _SORT_FREE_LIST_AFTER_MERGE:
-                self._free_phys_pages, _ = torch.sort(self._free_phys_pages)
 
     def _compact_pending(self, freed_physical_pages: torch.Tensor) -> None:
         """Eager compaction over the freed PHYSICAL pages: move survivors from the
@@ -1221,7 +1268,6 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
         self._stats_peak_pending_pages = max(
             self._stats_peak_pending_pages, cur_pending
         )
-        sort_tag = "ON" if _SORT_FREE_LIST_AFTER_MERGE else "OFF"
         logger.info(
             f"[lazy-stats sub={self.sub_pool_name!r}] "
             f"free_lazy={self._stats_n_free_lazy} "
@@ -1230,7 +1276,6 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
             f"moves={self._stats_n_flush_moves} "
             f"abs={self._stats_n_pages_absorbed}) "
             f"drain={self._stats_n_drain_did_work}/{self._stats_n_drain_calls} "
-            f"sort={sort_tag} "
             f"peak_holes={self._stats_peak_free_list_len} "
             f"peak_pending={self._stats_peak_pending_pages} "
             f"cur_holes={cur_holes} cur_pending={cur_pending} "
@@ -1254,7 +1299,6 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
             self._stats_peak_pending_pages = max(
                 self._stats_peak_pending_pages, cur_pending
             )
-            sort_tag = "ON" if _SORT_FREE_LIST_AFTER_MERGE else "OFF"
             self._stats_final_emitted = True
             logger.info(
                 f"[lazy-stats FINAL sub={self.sub_pool_name!r} reason={reason}] "
@@ -1264,7 +1308,6 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
                 f"moves={self._stats_n_flush_moves} "
                 f"abs={self._stats_n_pages_absorbed}) "
                 f"drain={self._stats_n_drain_did_work}/{self._stats_n_drain_calls} "
-                f"sort={sort_tag} "
                 f"peak_holes={self._stats_peak_free_list_len} "
                 f"peak_pending={self._stats_peak_pending_pages} "
                 f"cur_holes={cur_holes} cur_pending={cur_pending} "
@@ -1312,8 +1355,6 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
                 self._stats_n_drained_pages_total += sum(
                     t.numel() for t in ready_tensors
                 )
-                if _SORT_FREE_LIST_AFTER_MERGE:
-                    self._free_phys_pages, _ = torch.sort(self._free_phys_pages)
 
     def maybe_drain_pending_reuse(self) -> None:
         """Public scheduler hook (once per step): flow fired compaction-src pages
@@ -1452,8 +1493,8 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
         with record_function("MultiEndedAlloc._flush"):
             self._drain_pending_reuse(urgent=urgent)
 
-            # Sort ASCENDING (skip if the env knob keeps the list always-sorted).
-            if not _SORT_FREE_LIST_AFTER_MERGE and self._free_phys_pages.numel() > 1:
+            # Sort ASCENDING.
+            if self._free_phys_pages.numel() > 1:
                 self._free_phys_pages, _ = torch.sort(self._free_phys_pages)
 
             all_cpu = self._free_phys_pages.tolist()  # one batched D2H sync
@@ -1486,8 +1527,8 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
             # deferred to AFTER the trailing dst-slice, keeping `_free_phys_pages`
             # byte-identical to `holes_cpu` for the whole walk. That invariant is
             # what makes the directional dst-slice correct in both directions
-            # (catting srcs mid-flush would chop the wrong end / scramble under
-            # sort=ON, leaving ghost p2v=-1 pages + double-bound dsts). Event-
+            # (catting srcs mid-flush would chop the wrong end, leaving ghost
+            # p2v=-1 pages + double-bound dsts). Event-
             # PENDING srcs still route to `_pending_reuse` (read-race gating).
             released_fired: List[torch.Tensor] = []
 
@@ -1655,7 +1696,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
             # ONE bulk remap (single-writer on schedule_stream).
             self.virtual_to_physical[v_moveds_t] = dst_pages_t
             self.physical_to_virtual[dst_pages_t] = v_moveds_t
-            self.physical_to_virtual[src_pages_t] = -1
+            self.physical_to_virtual.index_fill_(0, src_pages_t, -1)
             self._inverse_history.append((src_pages_t, dst_pages_t, v_moveds_t))
             # Src disposition — ONE entry per batch. `src_pages_t` is reused as the
             # `_pending_reuse` GPU tensor (no second H2D at drain).
@@ -1695,6 +1736,19 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
             f"Caller: {callers}."
         )
 
+    # -- free-group --
+
+    def free_group_begin(self) -> None:
+        super().free_group_begin()
+        self.free_page_reps_group = []
+
+    def free_group_end(self) -> None:
+        pending, self.free_page_reps_group = self.free_page_reps_group, None
+        super().free_group_end()
+        if pending:
+            reps = torch.cat(pending)
+            self.free(reps, _pages=reps // self.page_size)
+
 
 class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
     """Composite allocator for the MHA (full-attn) + Mamba hybrid pair.
@@ -1764,6 +1818,7 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         # write locations are resolved in the attention metadata.
 
         self.free_group = None
+        self.free_page_reps_group: Optional[List[torch.Tensor]] = None
         # Base init left these None; we use watermark math, not free-lists.
         self.free_pages = torch.empty(0, dtype=torch.int64, device=device)
         self.release_pages = torch.empty(0, dtype=torch.int64, device=device)
@@ -1962,6 +2017,47 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         self.mamba_allocator.clear()
         self.free_group = None
 
+    def free_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None:
+        """Fixed-shape counterpart of `free()`; see
+        `MultiEndedAllocator._page_reps_pieces`. The mamba sub-pool is
+        slot-granular and untouched by a token free, so only the full side
+        needs the representatives.
+        """
+        if free_index is None or free_index.numel() == 0:
+            return
+        if self.page_size == 1:
+            self.free(free_index)
+            return
+        pieces = self.full_attn_allocator._page_reps_pieces(
+            free_index.detach().to(torch.int64), start_pos
+        )
+        if self.free_page_reps_group is None:
+            self._release_page_reps(pieces)
+        else:
+            self.free_page_reps_group.extend(pieces)
+
+    def _release_page_reps(self, pieces: Sequence[torch.Tensor]) -> None:
+        reps = pieces[0] if len(pieces) == 1 else torch.cat(tuple(pieces))
+        self.full_attn_allocator.free(reps, _pages=reps // self.page_size)
+        self.full_attn_allocator.clear_inverse_history()
+        self.mamba_allocator.clear_inverse_history()
+
+    def free_group_begin(self) -> None:
+        super().free_group_begin()
+        self.free_page_reps_group = []
+
+    def free_group_end(self) -> None:
+        pending, self.free_page_reps_group = self.free_page_reps_group, None
+        super().free_group_end()
+        if pending:
+            self._release_page_reps(pending)
+
+    def clear(self) -> None:
+        self.full_attn_allocator.clear()
+        self.mamba_allocator.clear()
+        self.free_group = None
+        self.free_page_reps_group = None
+
     # -- Lazy compaction hooks --
 
     def set_latest_forward_done_event(self, event: Optional[torch.cuda.Event]) -> None:
@@ -2098,6 +2194,7 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
         )
 
         self.free_group = None
+        self.free_page_reps_group: Optional[List[torch.Tensor]] = None
         # Empty (not None) for the leak checker.
         self.free_pages = torch.empty(0, dtype=torch.int64, device=device)
         self.release_pages = torch.empty(0, dtype=torch.int64, device=device)
@@ -2462,10 +2559,54 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
         # Paired with set_full_to_swa_mapping: shared mode has no mapping tensor.
         return
 
+    # -- free-group --
+
+    def free_group_begin(self) -> None:
+        super().free_group_begin()
+        self.free_page_reps_group = []
+
+    def free_group_end(self) -> None:
+        pending, self.free_page_reps_group = self.free_page_reps_group, None
+        super().free_group_end()
+        if pending:
+            self._release_page_reps(pending)
+
+    def free_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None:
+        """Fixed-shape counterpart of `free()`; see
+        `MultiEndedAllocator._page_reps_pieces`. Both sides share one
+        derivation -- neither repeats the position-less dedup.
+        """
+        if free_index is None or free_index.numel() == 0:
+            return
+        if self.page_size == 1:
+            self.free(free_index)
+            return
+        pieces = self.full_attn_allocator._page_reps_pieces(
+            free_index.detach().to(torch.int64), start_pos
+        )
+        if self.free_page_reps_group is None:
+            self._release_page_reps(pieces)
+        else:
+            self.free_page_reps_group.extend(pieces)
+
+    def _release_page_reps(self, pieces: Sequence[torch.Tensor]) -> None:
+        reps = pieces[0] if len(pieces) == 1 else torch.cat(tuple(pieces))
+        v_pages = reps // self.page_size
+        # Same tombstone filter as `free`, but at PAGE granularity (page_size
+        # times smaller): `> 0` strict -- -1 = tombstoned, 0 = padding sink.
+        swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages]
+        live_pages = v_pages[swa_v2p_pages > 0]
+        if live_pages.numel() > 0:
+            self.swa_attn_allocator.free(live_pages * self.page_size, _pages=live_pages)
+        self.full_attn_allocator.free(reps, _pages=v_pages)
+        self.full_attn_allocator.clear_inverse_history()
+        self.swa_attn_allocator.clear_inverse_history()
+
     def clear(self) -> None:
         self.full_attn_allocator.clear()
         self.swa_attn_allocator.clear()
         self.free_group = None
+        self.free_page_reps_group = None
 
     # -- Lazy compaction hooks --
 

python/sglang/srt/mem_cache/swa_memory_pool.py

@@ -209,6 +209,15 @@ class SWAKVPool(BaseSWAKVPool):
         else:
             return self.full_kv_pool.get_value_buffer(layer_id_pool)
 
+    def get_v_head_dim(self):
+        # The FULL side's dim, as HybridLinearKVPool.get_v_head_dim(): a caller
+        # asking a pool for "the" v_head_dim wants the full-attention geometry.
+        # `start_layer`, not 0, so pipeline parallelism (start_layer > 0) works,
+        # and because layer 0 need not be a full-attention layer.
+        return self.full_kv_pool.get_value_buffer(self.full_kv_pool.start_layer).shape[
+            -1
+        ]
+
     def get_kv_buffer(self, layer_id: int):
         self._wait_for_layer(layer_id)
         layer_id_pool, is_swa_layer = self.layers_mapping[layer_id]

python/sglang/srt/mem_cache/unified_cache/components/swa_component.py

@@ -178,6 +178,74 @@ class SWAComponent(TreeComponent):
             full_indices
         )
 
+    def _unified_allocator(self):
+        """The unified SWA composite, or None when running on the static pool."""
+        from sglang.srt.mem_cache.multi_ended_allocator import (
+            UnifiedSWATokenToKVPoolAllocator,
+        )
+
+        allocator = self.cache.token_to_kv_pool_allocator
+        if isinstance(allocator, UnifiedSWATokenToKVPoolAllocator):
+            return allocator
+        return None
+
+    def _page_pairs(
+        self, full_value: torch.Tensor, incoming_full_value: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Page ids of two token ranges that address the SAME logical tokens.
+
+        Dedupes by FIRST OCCURRENCE with one shared mask rather than
+        `torch.unique`: unique sorts by id value, and allocation hands out
+        virtual ids in no particular order, so sorting would pair page k of one
+        range with an unrelated page of the other. One mask keeps the pairing
+        positional, hence logical.
+        """
+        page_size = self.tree_core.page_size
+        kept = full_value.detach().to(torch.int64) // page_size
+        incoming = incoming_full_value.detach().to(torch.int64) // page_size
+        assert kept.numel() == incoming.numel(), (
+            f"locked-full recovery needs a 1:1 token correspondence, got "
+            f"{kept.numel()} kept vs {incoming.numel()} incoming"
+        )
+        starts = torch.ones_like(kept, dtype=torch.bool)
+        starts[1:] = kept[1:] != kept[:-1]
+        incoming_starts = torch.ones_like(incoming, dtype=torch.bool)
+        incoming_starts[1:] = incoming[1:] != incoming[:-1]
+        assert torch.equal(starts, incoming_starts), (
+            "the two ranges break into pages at different offsets, so no "
+            "page-granular ownership transfer expresses the token mapping"
+        )
+        return kept[starts], incoming[starts]
+
+    def _transfer_swa_pages(
+        self,
+        allocator,
+        full_value: torch.Tensor,
+        incoming_full_value: torch.Tensor,
+    ) -> None:
+        """Move swa page OWNERSHIP from the incoming ids onto the node's ids.
+
+        The static recipe re-points the node's locked full ids at the incoming
+        swa pages through `full_to_swa_index_mapping`. Under the unified pool
+        the swa sub-pool's v2p IS that mapping, so the same move is a rebind:
+        give the node's virtual pages the incoming pages' physical pages, then
+        tombstone the incoming ones. No page is allocated or freed, so no
+        capacity changes — only ownership does.
+        """
+        swa = allocator.swa_attn_allocator
+        kept_pages, incoming_pages = self._page_pairs(full_value, incoming_full_value)
+        physical = swa.virtual_to_physical[incoming_pages]
+        # `> 0` strict: -1 = tombstoned, 0 = the padding sink. The incoming ids
+        # were just allocated by the in-flight request, so every page must be
+        # live; a violation means we would hand the node the sink and serve
+        # zeros, which is worth a hard failure rather than silent corruption.
+        assert bool(
+            (physical > 0).all()
+        ), f"incoming swa pages must all be live, got {physical.tolist()}"
+        swa.bind(kept_pages, physical)
+        swa.virtual_to_physical.index_fill_(0, incoming_pages, -1)
+        swa.clear_inverse_history()
+
     def refresh_lru(
         self,
         phase: LRURefreshPhase,
@@ -1280,8 +1348,24 @@ class SWAComponent(TreeComponent):
                 alloc.set_full_to_swa_mapping(full, swa)
             return
         if isinstance(action, RecoverSWAWithLockedFull):
-            # Keep the locked full; remap it onto the incoming full's SWA translation,
+            # Keep the locked full; hand the node the INCOMING ids' swa pages,
             # freeing only the incoming full, then store the swa on the node.
+            unified = self._unified_allocator()
+            if unified is not None:
+                # No `full_to_swa_index_mapping` here: the swa sub-pool's v2p IS
+                # the mapping. Rebind page ownership, then free through the
+                # composite -- its `swa_v2p_pages > 0` filter skips the
+                # just-tombstoned swa side, releasing only the full one.
+                self._transfer_swa_pages(
+                    unified, action.kept_full, action.incoming_full
+                )
+                unified.free(action.incoming_full)
+                self.tree_core.set_component_device_value(
+                    action.node_id,
+                    self.component_type,
+                    self._translate_full_to_swa(action.kept_full),
+                )
+                return
             swa_value = self._translate_full_to_swa(action.incoming_full)
             alloc.set_full_to_swa_mapping(action.kept_full, swa_value)
             alloc.clear_full_to_swa_mapping(action.incoming_full)