[GDN][KDA][mem_cache] int8 checkpoint pool for the linear-attn prefix cache (#28185)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

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

@@ -64,9 +64,20 @@ class MambaAttnBackendBase(AttentionBackend):
             forward_batch.mamba_cow_src_indices is not None
             and len(forward_batch.mamba_cow_src_indices) > 0
         ):
-            self.req_to_token_pool.mamba_pool.copy_from(
-                forward_batch.mamba_cow_src_indices, forward_batch.mamba_cow_dst_indices
-            )
+            ckpt_pool = getattr(self.req_to_token_pool, "mamba_ckpt_pool", None)
+            if ckpt_pool is not None:
+                # int8 checkpoints: dequantize the cached state (src = int8 ckpt slot)
+                # into the request's active bf16 slot (dst).
+                ckpt_pool.load_to_active(
+                    self.req_to_token_pool.mamba_pool,
+                    forward_batch.mamba_cow_src_indices,
+                    forward_batch.mamba_cow_dst_indices,
+                )
+            else:
+                self.req_to_token_pool.mamba_pool.copy_from(
+                    forward_batch.mamba_cow_src_indices,
+                    forward_batch.mamba_cow_dst_indices,
+                )
         forward_batch.mamba_clear_indices = None
         forward_batch.mamba_cow_src_indices = None
         forward_batch.mamba_cow_dst_indices = None

python/sglang/srt/managers/scheduler_components/invariant_checker.py

@@ -114,6 +114,9 @@ class SchedulerInvariantChecker:
         )
 
     def _check_mamba_pool(self, ps: PoolStats) -> Tuple[bool, str]:
+        ckpt_pool = getattr(self.req_to_token_pool, "mamba_ckpt_pool", None)
+        if ckpt_pool is not None:
+            return self._check_mamba_pool_with_int8(ps, ckpt_pool)
         leak, msg = self._check_pool_invariant(
             "mamba",
             ps.mamba_available_size,
@@ -150,6 +153,37 @@ class SchedulerInvariantChecker:
             )
         return leak, msg
 
+    def _check_mamba_pool_with_int8(self, ps: PoolStats, ckpt_pool) -> Tuple[bool, str]:
+        """Two-pool invariant for int8 mamba checkpoints.
+
+        The radix-cached states live in the int8 checkpoint pool, NOT the active
+        bf16 pool. So the single-pool equation (active.available + radix_cached ==
+        active.size) is wrong -- it double-counts the radix states against a pool
+        that does not hold them. Instead check the two pools independently:
+
+          * active bf16 pool: backs running requests only; the radix owns ZERO
+            active slots. Checked at idle (in-flight == 0) -> available == total.
+          * int8 checkpoint pool: backs the radix-cached states; its occupancy is
+            exactly the radix evictable + protected counts.
+        """
+        active_leak, active_msg = self._check_pool_invariant(
+            "mamba-active",
+            ps.mamba_available_size,
+            ps.mamba_evictable_size,  # 0 in int8 mode (radix owns no active slots)
+            0,
+            self.pool_stats_observer.session_held_mamba_slots(),
+            self.req_to_token_pool.mamba_pool.size,
+        )
+        int8_leak, int8_msg = self._check_pool_invariant(
+            "mamba-int8",
+            ckpt_pool.available_size(),
+            self.tree_cache.mamba_evictable_size(),
+            self.tree_cache.mamba_protected_size(),
+            0,
+            ckpt_pool.num_slots,
+        )
+        return active_leak or int8_leak, active_msg + "\n" + int8_msg
+
     def _get_total_uncached_sizes(
         self,
     ) -> Tuple[int, int]:

python/sglang/srt/managers/scheduler_components/pool_stats_observer.py

@@ -247,8 +247,19 @@ class SchedulerPoolStatsObserver:
             self.tree_cache.full_evictable_size() if is_mamba_radix_cache else 0
         )
         mamba_available_size = self.req_to_token_pool.mamba_allocator.available_size()
+        # `mamba_usage`/`mamba_num_used` track the ACTIVE bf16 pool occupancy (running
+        # requests) -- this feeds throttle decisions (get_max_pool_usage) which asserts
+        # usage >= 0. With int8 checkpoints the radix-cached states live in a SEPARATE
+        # int8 pool, so they own ZERO active slots: report evictable=0 against the active
+        # pool (otherwise active.size - (available + radix_cached) goes negative). The
+        # int8 cache pool's own occupancy is validated separately in the invariant check.
+        has_int8_ckpt = (
+            getattr(self.req_to_token_pool, "mamba_ckpt_pool", None) is not None
+        )
         mamba_evictable_size = (
-            self.tree_cache.mamba_evictable_size() if is_mamba_radix_cache else 0
+            self.tree_cache.mamba_evictable_size()
+            if (is_mamba_radix_cache and not has_int8_ckpt)
+            else 0
         )
         full_num_used = self.token_to_kv_pool_allocator.size - (
             full_available_size + full_evictable_size

python/sglang/srt/mem_cache/mamba_checkpoint_pool.py

@@ -0,0 +1,373 @@
+"""
+Copyright 2023-2026 SGLang Team
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+MambaCheckpointPool — the radix prefix cache's int8-compressed store for cached
+linear-attention (KDA / GDN / Mamba2 gated-delta-rule) recurrent states.
+
+It decouples the *cached* states (radix-owned, idle, compressed) from the *active*
+``MambaPool`` (running requests, full precision, kernel-facing). The radix stores
+one cached state per node HERE; on a prefix-cache hit it is dequantized back into
+a fresh active slot (copy-on-write).
+
+Per cached slot it holds:
+  * the SSM temporal state in **int8** (per-(head,k-channel) symmetric), via the
+    embedded ``Int8CheckpointStore`` — ~2x more cached states than bf16,
+    quality-safe (quantized once on store, dequantized once on a hit; never
+    re-enters the recurrence as a quant->dequant loop).
+  * the conv1d window state at its native dtype (tiny, W-1 tokens; not worth
+    quantizing).
+
+Why int8 (not fp8): a cached checkpoint is loaded ONCE on a cache hit, then
+decoding continues at full precision, so the only error is a single rounding of
+S. The temporal state is roughly uniformly distributed, so int8-per-(head,
+k-channel) beats fp8-e4m3 at the same 1 byte (fp8 wastes bits on the exponent).
+The scale axis (reduces over d_v) matches the per-k-channel decay diag(alpha), so
+the large state entries keep ~bf16 precision and the error concentrates on small
+entries that barely affect the readout. Storing cached states int8 gives ~2x the
+cached-prefix capacity at fixed memory, and composes with host-offload
+(HiMambaRadixCache) which it also halves.
+
+This is strategy-agnostic: whether the active slot to be cached was produced by
+the ``no_buffer`` donate (copy_from) or the ``extra_buffer`` ping-pong track
+buffer (spec path), both converge on "an active slot becomes the cached
+``mamba_value``" — which is exactly the (store_from_active) hook here. Slot
+lifecycle is owned by the caller via the embedded ``MambaSlotAllocator``.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import List, Optional
+
+import torch
+
+from sglang.srt.mem_cache.allocator.mamba import MambaSlotAllocator
+
+logger = logging.getLogger(__name__)
+
+
+class Int8CheckpointStore:
+    """int8 store for cached multi-layer linear-attn states.
+
+    Tensors (slot index handed out by the caller's allocator):
+        qdata : [L, num_slots, H, d_v, d_k]  int8        (the quantized state)
+        scale : [L, num_slots, H, 1,   d_k]  scale_dtype  (per layer,slot,head,k-chan)
+
+    A "state" spans all L mamba layers for one cached point (matching how the
+    radix caches one full state per node). The reduction axis for the scale is
+    d_v (dim=-2), so each (head, k-channel) gets its own scale — aligned with the
+    per-k-channel decay diag(alpha).
+
+    ``scale_dtype`` should match the source state's dtype (bf16 / fp16 / fp32) so
+    that quantize and dequantize use the identical scale — it is NOT required to
+    be bf16.
+    """
+
+    QMAX = 127
+
+    def __init__(
+        self,
+        *,
+        num_layers: int,
+        num_slots: int,
+        num_heads: int,
+        head_v_dim: int,
+        head_k_dim: int,
+        device: str,
+        scale_dtype: torch.dtype = torch.bfloat16,
+    ):
+        self.num_layers = num_layers
+        self.num_slots = num_slots
+        self.H = num_heads
+        self.d_v = head_v_dim
+        self.d_k = head_k_dim
+        self.device = device
+        self.qdata = torch.empty(
+            num_layers,
+            num_slots,
+            num_heads,
+            head_v_dim,
+            head_k_dim,
+            dtype=torch.int8,
+            device=device,
+        )
+        self.scale = torch.empty(
+            num_layers,
+            num_slots,
+            num_heads,
+            1,
+            head_k_dim,
+            dtype=scale_dtype,
+            device=device,
+        )
+
+    # ---- (de)quant math (also usable standalone for probes/tests) ----
+
+    @classmethod
+    def quantize(cls, state: torch.Tensor):
+        """state [..., H, d_v, d_k] -> (qint8, scale[..., H, 1, d_k]).
+
+        amax / scale / round are computed in float32 so quantizing a low-precision
+        state doesn't lose precision in the intermediate (symmetric with
+        ``dequantize``, which is already float32). The scale is rounded to the
+        state dtype (its storage precision) BEFORE the division, so quantize and
+        dequantize use the identical scale."""
+        state_fp32 = state.to(torch.float32)
+        amax = state_fp32.abs().amax(dim=-2, keepdim=True).clamp(min=1e-8)
+        scale = (amax / cls.QMAX).to(state.dtype)
+        q = (
+            torch.round(state_fp32 / scale.to(torch.float32))
+            .clamp(-cls.QMAX, cls.QMAX)
+            .to(torch.int8)
+        )
+        return q, scale
+
+    @staticmethod
+    def dequantize(q: torch.Tensor, scale: torch.Tensor, out_dtype: torch.dtype):
+        return (q.to(torch.float32) * scale.to(torch.float32)).to(out_dtype)
+
+    # ---- store / load (caller supplies slot indices) ----
+
+    def store(self, slots: torch.Tensor, state: torch.Tensor) -> None:
+        """Quantize and write states. state: [L, N, H, d_v, d_k] for the N slots
+        (or [L, H, d_v, d_k] when slots is a scalar/len-1)."""
+        if state.dim() == 4:
+            state = state.unsqueeze(1)
+        q, scale = self.quantize(state)
+        self.qdata[:, slots] = q
+        self.scale[:, slots] = scale.to(self.scale.dtype)
+
+    def load(self, slots: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor:
+        """Dequantize states at slots -> [L, N, H, d_v, d_k] in out_dtype."""
+        return self.dequantize(self.qdata[:, slots], self.scale[:, slots], out_dtype)
+
+    def copy_to_pool(
+        self,
+        dst_temporal: torch.Tensor,
+        src_slots: torch.Tensor,
+        dst_slots: torch.Tensor,
+    ) -> None:
+        """Dequantize checkpoints at ``src_slots`` directly into the active pool
+        tensor ``dst_temporal`` [L, pool_slots, H, d_v, d_k] at ``dst_slots`` (the
+        copy-on-write on a cache hit). Output dtype follows ``dst_temporal``."""
+        dst_temporal[:, dst_slots] = self.load(src_slots, dst_temporal.dtype)
+
+    def store_from_pool(
+        self,
+        src_temporal: torch.Tensor,
+        src_slots: torch.Tensor,
+        dst_slots: torch.Tensor,
+    ) -> None:
+        """Quantize states from an active pool tensor into checkpoint slots (cache
+        store / donate)."""
+        self.store(dst_slots, src_temporal[:, src_slots])
+
+    def mem_usage_bytes(self) -> int:
+        return (
+            self.qdata.numel() * self.qdata.element_size()
+            + self.scale.numel() * self.scale.element_size()
+        )
+
+    def bytes_per_slot(self) -> int:
+        return self.mem_usage_bytes() // max(1, self.num_slots)
+
+
+class MambaCheckpointPool:
+    def __init__(
+        self,
+        *,
+        num_layers: int,
+        num_slots: int,
+        num_heads: int,
+        head_v_dim: int,
+        head_k_dim: int,
+        conv_shapes: List[tuple],
+        conv_dtype: torch.dtype,
+        device: str,
+        temporal_dtype: Optional[torch.dtype] = None,
+    ):
+        self.num_slots = num_slots
+        self.device = device
+        self.temporal = Int8CheckpointStore(
+            num_layers=num_layers,
+            num_slots=num_slots + 1,  # slot 0 reserved (matches MambaSlotAllocator)
+            num_heads=num_heads,
+            head_v_dim=head_v_dim,
+            head_k_dim=head_k_dim,
+            device=device,
+            # store the scale in the temporal state's own dtype so quantize and
+            # dequantize use the identical scale (not hard-coded to bf16)
+            scale_dtype=(
+                temporal_dtype if temporal_dtype is not None else torch.bfloat16
+            ),
+        )
+        # conv windows stay at their native dtype (small); one buffer per conv
+        # tensor in the State
+        self.conv = [
+            torch.empty(
+                (num_layers, num_slots + 1) + tuple(shape),
+                dtype=conv_dtype,
+                device=device,
+            )
+            for shape in conv_shapes
+        ]
+        self.allocator = MambaSlotAllocator(size=num_slots, device=device)
+
+    # ---- lifecycle (delegates to the embedded allocator) ----
+
+    def alloc(self, n: int = 1):
+        return self.allocator.alloc(n)
+
+    def free(self, slots: torch.Tensor):
+        self.allocator.free(slots)
+
+    def available_size(self) -> int:
+        return self.allocator.available_size()
+
+    def clear(self) -> None:
+        """Release every checkpoint slot (radix flush/reset). The int8 qdata is
+        left as-is; slots are reused/overwritten on the next store."""
+        self.allocator.clear()
+
+    # ---- state transfer between the active MambaPool and this store ----
+
+    def store_from_active(self, active_mamba_pool, active_slots, ckpt_slots) -> None:
+        """Quantize temporal + copy conv from the active pool into checkpoint slots
+        (the radix donate / cache-store)."""
+        cache = active_mamba_pool.mamba_cache
+        self.temporal.store_from_pool(cache.temporal, active_slots, ckpt_slots)
+        for i, c in enumerate(self.conv):
+            c[:, ckpt_slots] = cache.conv[i][:, active_slots]
+
+    def load_to_active(self, active_mamba_pool, ckpt_slots, active_slots) -> None:
+        """Dequantize temporal + copy conv from checkpoint slots into the active pool
+        (the cache-hit copy-on-write)."""
+        cache = active_mamba_pool.mamba_cache
+        self.temporal.copy_to_pool(cache.temporal, ckpt_slots, active_slots)
+        for i, c in enumerate(self.conv):
+            cache.conv[i][:, active_slots] = c[:, ckpt_slots].to(cache.conv[i].dtype)
+
+    @staticmethod
+    def estimate_mem_usage_bytes(
+        *,
+        num_layers: int,
+        num_slots: int,
+        num_heads: int,
+        head_v_dim: int,
+        head_k_dim: int,
+        conv_shapes: List[tuple],
+        conv_dtype: torch.dtype,
+        temporal_dtype: torch.dtype,
+    ) -> dict:
+        """Estimate the pool's HBM footprint (bytes) WITHOUT allocating, so a
+        caller can check it against free memory before construction. Mirrors the
+        real layout: int8 qdata + per-(head,k) scale + bf16 conv windows, including
+        the reserved slot 0."""
+        slots = num_slots + 1  # slot 0 reserved (matches MambaSlotAllocator)
+        scale_isz = torch.empty((), dtype=temporal_dtype).element_size()
+        conv_isz = torch.empty((), dtype=conv_dtype).element_size()
+        qdata = num_layers * slots * num_heads * head_v_dim * head_k_dim  # int8 = 1B
+        scale = num_layers * slots * num_heads * head_k_dim * scale_isz
+        conv = 0
+        for shape in conv_shapes:
+            n = 1
+            for s in shape:
+                n *= int(s)
+            conv += num_layers * slots * n * conv_isz
+        return {
+            "qdata": qdata,
+            "scale": scale,
+            "conv": conv,
+            "total": qdata + scale + conv,
+        }
+
+    def mem_usage_bytes(self) -> int:
+        conv_bytes = sum(c.numel() * c.element_size() for c in self.conv)
+        return self.temporal.mem_usage_bytes() + conv_bytes
+
+
+def maybe_init_int8_mamba_checkpoint_pool(
+    *,
+    mamba_size: int,
+    cache_params,
+    mamba_layer_ids: List[int],
+    device: str,
+) -> Optional[MambaCheckpointPool]:
+    """Build the optional int8 ``MambaCheckpointPool`` when
+    ``--enable-int8-mamba-checkpoint`` is set (and a global server-args context
+    exists), else return ``None``. The radix caches states here (int8) instead of
+    in the active bf16 pool -> ~2x cached-prefix capacity at fixed memory.
+
+    Estimates the pool's HBM footprint and checks it against free memory BEFORE
+    allocating, so an oversized ``--int8-mamba-ckpt-size`` fails with an actionable
+    message instead of a cryptic mid-allocation CUDA OOM.
+    """
+    from sglang.srt.server_args import get_global_server_args
+
+    try:
+        _sa = get_global_server_args()
+    except ValueError:
+        # Some unit-test / mock runners construct HybridReqToTokenPool directly
+        # without a global server-args context. The int8 checkpoint pool is opt-in
+        # via a CLI flag, so an unset context unambiguously means it is off.
+        _sa = None
+    if not getattr(_sa, "enable_int8_mamba_checkpoint", False):
+        return None
+
+    GB = 1 << 30
+    H, d_v, d_k = cache_params.shape.temporal
+    ckpt_size = _sa.int8_mamba_ckpt_size or (2 * mamba_size)
+    kwargs = dict(
+        num_layers=len(mamba_layer_ids),
+        num_slots=ckpt_size,
+        num_heads=H,
+        head_v_dim=d_v,
+        head_k_dim=d_k,
+        conv_shapes=list(cache_params.shape.conv),
+        conv_dtype=cache_params.dtype.conv,
+        temporal_dtype=cache_params.dtype.temporal,
+    )
+
+    est = MambaCheckpointPool.estimate_mem_usage_bytes(**kwargs)
+    free_bytes = None
+    if isinstance(device, str) and device.startswith("cuda"):
+        try:
+            free_bytes, _ = torch.cuda.mem_get_info(device)
+        except Exception:
+            free_bytes = None
+    logger.info(
+        f"int8 mamba checkpoint pool: {ckpt_size} slots, "
+        f"{est['total'] / GB:.2f}GB (qdata {est['qdata'] / GB:.2f} + scale "
+        f"{est['scale'] / GB:.2f} + conv {est['conv'] / GB:.2f}); active mamba "
+        f"pool {mamba_size} slots"
+        + (f"; free HBM {free_bytes / GB:.2f}GB" if free_bytes is not None else "")
+    )
+    if free_bytes is not None and est["total"] >= free_bytes:
+        raise RuntimeError(
+            f"int8 mamba checkpoint pool needs ~{est['total'] / GB:.2f}GB but only "
+            f"{free_bytes / GB:.2f}GB HBM is free. Lower --int8-mamba-ckpt-size "
+            f"(currently {ckpt_size}) or --mem-fraction-static."
+        )
+
+    pool = MambaCheckpointPool(device=device, **kwargs)
+    # NOTE: this pool's HBM is NOT subtracted from the KV-cache budget
+    # (max_total_num_tokens); it is allocated from --mem-fraction-static headroom.
+    # The estimate check above guards against an oversized pool; accounting it in
+    # the KV budget is a follow-up.
+    logger.warning(
+        f"int8 mamba checkpoint pool ({est['total'] / GB:.2f}GB) is allocated from "
+        f"--mem-fraction-static headroom and is not reflected in "
+        f"max_total_num_tokens; ensure headroom covers it."
+    )
+    return pool

python/sglang/srt/mem_cache/mamba_radix_cache.py

@@ -562,22 +562,29 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
                 mamba_ping_pong_track_buffer_to_keep = (
                     self.req_to_token_pool.get_mamba_ping_pong_keep_idx(req)
                 )
-                mamba_value = (
-                    req.mamba_ping_pong_track_buffer[
-                        mamba_ping_pong_track_buffer_to_keep
-                    ]
-                    .unsqueeze(-1)
-                    .clone()
-                )
-                assert mamba_value.item() != -1, (
+                src_active = req.mamba_ping_pong_track_buffer[
+                    mamba_ping_pong_track_buffer_to_keep
+                ].unsqueeze(-1)
+                assert src_active.item() != -1, (
                     f"Cached mamba slot is -1: keep_idx={mamba_ping_pong_track_buffer_to_keep}, "
                     f"buf={req.mamba_ping_pong_track_buffer.tolist()}, "
                     f"next_track_idx={req.mamba_next_track_idx}, "
                     f"last_track_seqlen={req.mamba_last_track_seqlen}, "
                     f"rid={req.rid}"
                 )
+                if self.int8_ckpt_pool is not None:
+                    mamba_value = self._commit_int8_checkpoint(src_active)
+                    # quantized -> no ping-pong slot needs keeping
+                    mamba_ping_pong_track_buffer_to_keep = None
+                else:
+                    mamba_value = src_active.clone()
             else:
-                mamba_value = req.mamba_pool_idx.unsqueeze(-1).clone()
+                if self.int8_ckpt_pool is not None:
+                    mamba_value = self._commit_int8_checkpoint(
+                        req.mamba_pool_idx.unsqueeze(-1)
+                    )
+                else:
+                    mamba_value = req.mamba_pool_idx.unsqueeze(-1).clone()
                 mamba_ping_pong_track_buffer_to_keep = None
 
             result = self.insert(
@@ -589,6 +596,9 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
                 )
             )
             mamba_exist = result.mamba_exist
+            if mamba_exist and self.int8_ckpt_pool is not None:
+                # state already cached -> the int8 slot we just allocated is a duplicate
+                self.int8_ckpt_pool.free(mamba_value)
         else:
             self.token_to_kv_pool_allocator.free(kv_indices[req.cache_protected_len :])
             mamba_exist = True
@@ -596,7 +606,13 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
         if mamba_exist:
             mamba_ping_pong_track_buffer_to_keep = None
 
-        free_mamba_cache = True if self.enable_mamba_extra_buffer else mamba_exist
+        # With int8 checkpoints the radix owns an int8 slot (not the request's active
+        # slot), so the active mamba slot must always be returned to the active pool.
+        free_mamba_cache = (
+            True
+            if (self.enable_mamba_extra_buffer or self.int8_ckpt_pool is not None)
+            else mamba_exist
+        )
 
         if free_mamba_cache:
             self.req_to_token_pool.free_mamba_cache(
@@ -649,7 +665,21 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
 
         # Donate the mamba index to the radix cache instead of copying.
         # This avoids a data copy that would race with the forward stream.
-        if self.enable_mamba_extra_buffer:
+        if self.int8_ckpt_pool is not None:
+            # int8 path: quantize the to-be-cached active state into an int8 slot
+            # (strategy-agnostic donate hook).
+            if self.enable_mamba_extra_buffer:
+                new_slot = self._alloc_mamba_slot()
+                src_active = self.req_to_token_pool.donate_mamba_ping_pong_slot(
+                    req, new_slot
+                )
+                mamba_value_donated = self._commit_int8_checkpoint(src_active)
+                self.req_to_token_pool.mamba_allocator.free(src_active)
+            else:
+                mamba_value_donated = self._commit_int8_checkpoint(
+                    req.mamba_pool_idx.view(-1)
+                )
+        elif self.enable_mamba_extra_buffer:
             new_slot = self._alloc_mamba_slot()
             mamba_value_donated = self.req_to_token_pool.donate_mamba_ping_pong_slot(
                 req, new_slot
@@ -671,7 +701,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
         )
         new_prefix_len, mamba_exist = result.prefix_len, result.mamba_exist
         if mamba_exist:
-            self.req_to_token_pool.mamba_allocator.free(mamba_value_donated)
+            self._free_mamba_value(mamba_value_donated)
 
         # The prefix indices could be updated, reuse it
         match_result = self.match_prefix(
@@ -729,7 +759,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
         self._record_remove_event(x)
         self.token_to_kv_pool_allocator.free(x.value)
         full_num_evicted = len(x.value)
-        self.req_to_token_pool.mamba_allocator.free(x.mamba_value)
+        self._free_mamba_value(x.mamba_value)
         mamba_num_evicted = len(x.mamba_value)
 
         # 2. get the next node, update the lru lists
@@ -782,7 +812,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
 
             if len(x.children) > 0:
                 # 1. an internal node, free mamba tokens.
-                self.req_to_token_pool.mamba_allocator.free(x.mamba_value)
+                self._free_mamba_value(x.mamba_value)
                 mamba_num_evicted += len(x.mamba_value)
 
                 # 2. get the next node, update the lru lists
@@ -954,6 +984,41 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
             assert slot is not None, "Can not alloc mamba cache"
         return slot
 
+    @property
+    def int8_ckpt_pool(self):
+        """The int8 checkpoint pool, or None when --enable-int8-mamba-checkpoint is off.
+        When enabled, radix-cached mamba states live HERE (int8), not in the active
+        bf16 pool -> ~2x cached-prefix capacity at fixed memory."""
+        return getattr(self.req_to_token_pool, "mamba_ckpt_pool", None)
+
+    def _alloc_int8_ckpt_slot(self) -> torch.Tensor:
+        """Allocate one int8 checkpoint slot, evicting cached states if the pool is full."""
+        slot = self.int8_ckpt_pool.alloc(1)
+        if slot is None:
+            self.evict(EvictParams(num_tokens=0, mamba_num=1))
+            slot = self.int8_ckpt_pool.alloc(1)
+            assert slot is not None, "Can not alloc int8 mamba checkpoint slot"
+        return slot
+
+    def _commit_int8_checkpoint(self, active_slots: torch.Tensor) -> torch.Tensor:
+        """Quantize the active-pool state at ``active_slots`` into a fresh int8
+        checkpoint slot and return that slot. Strategy-agnostic donate hook: both
+        no_buffer (copy_from) and extra_buffer (ping-pong) converge here. The caller
+        frees ``active_slots`` separately."""
+        ckpt_slot = self._alloc_int8_ckpt_slot()
+        self.int8_ckpt_pool.store_from_active(
+            self.req_to_token_pool.mamba_pool, active_slots, ckpt_slot
+        )
+        return ckpt_slot
+
+    def _free_mamba_value(self, mamba_value: torch.Tensor) -> None:
+        """Free a node's mamba_value to the right allocator (int8 ckpt pool or the
+        active mamba allocator)."""
+        if self.int8_ckpt_pool is not None:
+            self.int8_ckpt_pool.free(mamba_value)
+        else:
+            self.req_to_token_pool.mamba_allocator.free(mamba_value)
+
     def _match_prefix_helper(
         self, key: RadixKey
     ) -> Tuple[List[torch.Tensor], TreeNode, int]:

python/sglang/srt/mem_cache/memory_pool.py

@@ -612,6 +612,20 @@ class HybridReqToTokenPool(ReqToTokenPool):
         )
         self.mamba_map = {layer_id: i for i, layer_id in enumerate(mamba_layer_ids)}
 
+        # Optional int8 checkpoint pool: the radix caches states here (int8) instead
+        # of holding them in the active bf16 pool -> ~2x cached-prefix capacity at
+        # fixed memory. Strategy-agnostic (no_buffer / extra_buffer / spec).
+        from sglang.srt.mem_cache.mamba_checkpoint_pool import (
+            maybe_init_int8_mamba_checkpoint_pool,
+        )
+
+        self.mamba_ckpt_pool = maybe_init_int8_mamba_checkpoint_pool(
+            mamba_size=mamba_size,
+            cache_params=cache_params,
+            mamba_layer_ids=mamba_layer_ids,
+            device=device,
+        )
+
         self.device = device
         req_pool_size = self.req_to_token.shape[0]
         self.req_index_to_mamba_index_mapping: torch.Tensor = torch.zeros(
@@ -821,6 +835,12 @@ class HybridReqToTokenPool(ReqToTokenPool):
         logger.info("Reset HybridReqToTokenPool")
         super().clear()
         self.mamba_allocator.clear()
+        # The int8 checkpoint pool holds radix-cached states in its own slots; a
+        # flush/reset drops the radix tree, so its slots must be released too,
+        # otherwise the (now unreferenced) slots leak and break the int8-pool
+        # invariant (int8_available + radix_cached != int8_total).
+        if self.mamba_ckpt_pool is not None:
+            self.mamba_ckpt_pool.clear()
         self.req_index_to_mamba_index_mapping.zero_()
         if self.enable_mamba_extra_buffer:
             self.req_index_to_mamba_ping_pong_track_buffer_mapping.zero_()

python/sglang/srt/server_args.py

@@ -694,6 +694,12 @@ class ServerArgs:
     mamba_full_memory_ratio: float = 0.9
     mamba_scheduler_strategy: str = "auto"
     mamba_track_interval: int = 256
+    # int8-compress radix-cached linear-attn (mamba) checkpoints -> ~2x cached
+    # prefixes at fixed memory (quality-safe; see mem_cache/mamba_checkpoint_pool.py).
+    enable_int8_mamba_checkpoint: bool = False
+    int8_mamba_ckpt_size: Optional[int] = (
+        None  # #int8 checkpoint slots; default 2x the active pool
+    )
     linear_attn_backend: str = "triton"
     linear_attn_decode_backend: Optional[str] = None
     linear_attn_prefill_backend: Optional[str] = None
@@ -1009,6 +1015,7 @@ class ServerArgs:
         self._handle_deterministic_inference()
         self._handle_attention_backend_compatibility()
         self._handle_mamba_backend()
+        self._handle_int8_mamba_checkpoint()
         self._handle_linear_attn_backend()
         self._handle_kv4_compatibility()
         self._handle_page_size()
@@ -3376,6 +3383,28 @@ class ServerArgs:
                     "FlashInfer mamba module not available, please check flashinfer installation."
                 )
 
+    def _handle_int8_mamba_checkpoint(self):
+        # The int8 mamba checkpoint pool is only wired into the built-in
+        # MambaRadixCache. The host-offload variant (HiMambaRadixCache, enabled by
+        # --enable-hierarchical-cache) and custom radix-cache backends are NOT
+        # int8-aware: they would read int8 checkpoint slots as bf16 active slots
+        # (wrong pool / out-of-range). Reject the combination up front rather than
+        # silently corrupting state.
+        if not self.enable_int8_mamba_checkpoint:
+            return
+        if self.enable_hierarchical_cache:
+            raise ValueError(
+                "--enable-int8-mamba-checkpoint is not supported together with "
+                "--enable-hierarchical-cache: the host-offload path "
+                "(HiMambaRadixCache) is not int8-aware. Disable one of them."
+            )
+        if self.radix_cache_backend is not None:
+            raise ValueError(
+                "--enable-int8-mamba-checkpoint only supports the built-in mamba "
+                f"radix cache; --radix-cache-backend={self.radix_cache_backend!r} "
+                "is not int8-aware. Omit --radix-cache-backend."
+            )
+
     def _handle_linear_attn_backend(self):
         import torch
 
@@ -6444,6 +6473,19 @@ class ServerArgs:
             default=ServerArgs.max_mamba_cache_size,
             help="The maximum size of the mamba cache.",
         )
+        parser.add_argument(
+            "--enable-int8-mamba-checkpoint",
+            action="store_true",
+            help="Store radix-cached linear-attn (mamba) states in int8 (separate "
+            "checkpoint pool) for ~2x cached-prefix capacity at fixed memory.",
+        )
+        parser.add_argument(
+            "--int8-mamba-ckpt-size",
+            type=int,
+            default=ServerArgs.int8_mamba_ckpt_size,
+            help="Number of int8 mamba checkpoint slots (default: 2x the active "
+            "mamba pool size).",
+        )
         parser.add_argument(
             "--mamba-ssm-dtype",
             type=str,