[AMD][DSV4] Reland unified-KV pool sizing and SWA ring accounting, fully gated

python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh

@@ -50,6 +50,7 @@ struct Prefill0Params {
   /// \brief Trailing tokens the write plan keeps resident in the compress state ring.
   /// Derived from the ring in `plan_compress_prefill`; see the bound there.
   int32_t mtp_pad;
+  bool use_req_ring;
 };
 
 struct Prefill1Params {
@@ -67,6 +68,7 @@ struct Prefill1Params {
   int32_t swa_page_size;
   int32_t ring_size;
   int32_t compress_ratio;
+  bool use_req_ring;
 };
 
 struct DecodeParams {
@@ -80,6 +82,7 @@ struct DecodeParams {
   int32_t swa_page_size;
   int32_t ring_size;
   int32_t compress_ratio;
+  bool use_req_ring;
 };
 
 struct Prefill1ParamsLegacy {
@@ -203,7 +206,7 @@ __global__ __launch_bounds__(1024, 1)  //
       const int32_t last_c_pos = (sl / cr) * cr;
       const int32_t first_w_pos = min(last_c_pos - (is_overlap ? cr : 0), sl - params.mtp_pad);
       bool do_write = position >= first_w_pos;
-      if (!do_write && is_overlap) do_write = (position % sps) >= (sps - cr);
+      if (!do_write && is_overlap && !params.use_req_ring) do_write = (position % sps) >= (sps - cr);
       if (do_write) {
         const uint32_t out_idx = atomicAdd(&counter_w, 1u);
         params.plan_w[out_idx] = pack_w(ragged_id, batch_id, position + 1);
@@ -236,7 +239,7 @@ __global__ __launch_bounds__(1024, 1)  //
         }
 
         bool do_write = position >= first_w_pos;
-        if (!do_write && is_overlap) do_write = (position % sps) >= (sps - cr);
+        if (!do_write && is_overlap && !params.use_req_ring) do_write = (position % sps) >= (sps - cr);
         if (do_write) {
           const uint32_t out_idx = atomicAdd(&counter_w, 1u);
           params.plan_w[out_idx] = pack_w(ragged_id, static_cast<uint32_t>(batch_id), position + 1);
@@ -270,7 +273,7 @@ __global__ void plan_compress_prefill_kernel_1(const Prefill1Params params) {
     const auto ring_offset = swa_loc % params.ring_size;
     return swa_page * params.ring_size + ring_offset;
   };
-  const auto compute_c128_loc = [&](int64_t rid, int32_t position) {
+  const auto compute_req_ring_loc = [&](int64_t rid, int32_t position) {
     return static_cast<int32_t>(rid * params.ring_size + position % params.ring_size);
   };
 
@@ -283,9 +286,9 @@ __global__ void plan_compress_prefill_kernel_1(const Prefill1Params params) {
       const auto position_1 = static_cast<int32_t>(plan_c.seq_len - 1);
       // only used for c4, harmless for c128
       const auto position_0 = max(position_1 - params.compress_ratio, 0);
-      if (params.compress_ratio == 128) {
-        plan_c.read_page_0 = compute_c128_loc(rid, position_0) / 128;
-        plan_c.read_page_1 = compute_c128_loc(rid, position_1) / 128;
+      if (params.compress_ratio == 128 || params.use_req_ring) {
+        plan_c.read_page_0 = compute_req_ring_loc(rid, position_0) / params.compress_ratio;
+        plan_c.read_page_1 = compute_req_ring_loc(rid, position_1) / params.compress_ratio;
       } else {
         const auto raw_loc_0 = mapping[position_0];
         const auto raw_loc_1 = mapping[position_1];
@@ -307,8 +310,8 @@ __global__ void plan_compress_prefill_kernel_1(const Prefill1Params params) {
     // `seq_len` (`write_loc`) may not be aligned here
     const auto position = static_cast<int32_t>(plan_w.write_loc - 1);
     plan_w.ragged_id = ragged_id;
-    if (params.compress_ratio == 128) {
-      plan_w.write_loc = compute_c128_loc(rid, position);
+    if (params.compress_ratio == 128 || params.use_req_ring) {
+      plan_w.write_loc = compute_req_ring_loc(rid, position);
     } else {
       const auto raw_loc = mapping[position];
       plan_w.write_loc = compute_loc(params.f2s_ptr[raw_loc]);
@@ -329,7 +332,7 @@ __global__ void plan_compress_decode_kernel(const DecodeParams params) {
     const auto ring_offset = swa_loc % params.ring_size;
     return swa_page * params.ring_size + ring_offset;
   };
-  const auto compute_c128_loc = [&](int64_t rid, int32_t position) {
+  const auto compute_req_ring_loc = [&](int64_t rid, int32_t position) {
     return static_cast<int32_t>(rid * params.ring_size + position % params.ring_size);
   };
   const auto seq_len = static_cast<int32_t>(params.seq_ptr[idx]);
@@ -338,10 +341,10 @@ __global__ void plan_compress_decode_kernel(const DecodeParams params) {
   int32_t write_loc;
   int32_t read_page_0;
   int32_t read_page_1;
-  if (params.compress_ratio == 128) {
-    write_loc = compute_c128_loc(rid, position_1);
-    read_page_0 = compute_c128_loc(rid, position_0) / 128;
-    read_page_1 = compute_c128_loc(rid, position_1) / 128;
+  if (params.compress_ratio == 128 || params.use_req_ring) {
+    write_loc = compute_req_ring_loc(rid, position_1);
+    read_page_0 = compute_req_ring_loc(rid, position_0) / params.compress_ratio;
+    read_page_1 = compute_req_ring_loc(rid, position_1) / params.compress_ratio;
   } else {
     const auto raw_loc_0 = mapping[position_0];
     const auto raw_loc_1 = mapping[position_1];
@@ -461,6 +464,7 @@ inline PrefillPlan plan_compress_prefill(
     const int32_t compress_ratio,
     const int32_t swa_page_size,
     const int32_t ring_size,
+    const bool use_req_ring,
     const bool use_cuda_graph) {
   auto B = SymbolicSize{"batch_size"};
   auto N = SymbolicSize{"num_q_tokens"};
@@ -503,6 +507,7 @@ inline PrefillPlan plan_compress_prefill(
   const auto batch_size = static_cast<uint32_t>(B.unwrap());
   constexpr auto kMaxTokens = static_cast<uint32_t>(std::numeric_limits<uint16_t>::max());
   RuntimeCheck(compress_ratio == 4 || compress_ratio == 128);
+  RuntimeCheck(!use_req_ring || compress_ratio == 4);
   RuntimeCheck(batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens);
   // `swa_page_size` >= `ring_size` >= `compress_ratio`
   RuntimeCheck(swa_page_size % ring_size == 0 && ring_size % compress_ratio == 0);
@@ -537,6 +542,7 @@ inline PrefillPlan plan_compress_prefill(
         .compress_ratio = compress_ratio,
         .swa_page_size = swa_page_size,
         .mtp_pad = mtp_pad,
+        .use_req_ring = use_req_ring,
     };
     LaunchKernel(1, kMaxPrefillBatchSize, device)(plan_compress_prefill_kernel0, params0);
     // kernel_1 sees the already-padded buffers, so num_c == num_w == num_padded == num_q_tokens.
@@ -555,6 +561,7 @@ inline PrefillPlan plan_compress_prefill(
         .swa_page_size = swa_page_size,
         .ring_size = ring_size,
         .compress_ratio = compress_ratio,
+        .use_req_ring = use_req_ring,
     };
     const auto block_size_1 = 256;
     const auto num_blocks_1 = div_ceil(params1.num_work, block_size_1);
@@ -582,7 +589,7 @@ inline PrefillPlan plan_compress_prefill(
     RuntimeCheck(0 < extend_len && extend_len <= seq_len);
     const auto should_write = [=](int32_t position) {
       if (position >= first_w_pos) return true;
-      return is_overlap && position % swa_page_size >= (swa_page_size - compress_ratio);
+      return is_overlap && !use_req_ring && position % swa_page_size >= (swa_page_size - compress_ratio);
     };
     for (const auto j : irange(extend_len)) {
       const int32_t position = prefix_len + j;
@@ -631,6 +638,7 @@ inline PrefillPlan plan_compress_prefill(
       .swa_page_size = swa_page_size,
       .ring_size = ring_size,
       .compress_ratio = compress_ratio,
+      .use_req_ring = use_req_ring,
   };
   const auto block_size = 256;
   const auto num_blocks = div_ceil(params.num_work, block_size);
@@ -645,7 +653,8 @@ inline tvm::ffi::Tensor plan_compress_decode(
     const tvm::ffi::TensorView seq_lens,          // CPU/GPU
     const int32_t compress_ratio,
     const int32_t swa_page_size,
-    const int32_t ring_size) {
+    const int32_t ring_size,
+    const bool use_req_ring) {
   auto B = SymbolicSize{"batch_size"};
   auto device_ = SymbolicDevice{};
   device_.set_options<kDLGPU>();
@@ -667,6 +676,7 @@ inline tvm::ffi::Tensor plan_compress_decode(
       .with_device(device_)
       .verify(seq_lens);
 
+  RuntimeCheck(!use_req_ring || compress_ratio == 4);
   const auto batch_size = static_cast<uint32_t>(B.unwrap());
   const auto device = device_.unwrap();
   auto D = ffi::empty({batch_size, sizeof(PlanD)}, kDLUInt8, device);
@@ -681,6 +691,7 @@ inline tvm::ffi::Tensor plan_compress_decode(
       .swa_page_size = swa_page_size,
       .ring_size = ring_size,
       .compress_ratio = compress_ratio,
+      .use_req_ring = use_req_ring,
   };
   const auto block_size = 256;
   const auto num_blocks = div_ceil(batch_size, block_size);

python/sglang/kernels/ops/attention/dsv4/attn.py

@@ -100,6 +100,7 @@ def create_paged_compress_data_kernel(
     stride_out_1_1: tl.constexpr,
     compress_ratio: tl.constexpr,
     is_overlap: tl.constexpr,
+    use_req_ring: tl.constexpr,
     swa_page_size: tl.constexpr,
     ring_size: tl.constexpr,
     BLOCK: tl.constexpr,
@@ -133,7 +134,7 @@ def create_paged_compress_data_kernel(
         else:
             pos = write_overlap_pos
         pos = tl.maximum(pos, 0)
-        if compress_ratio == 128:
+        if compress_ratio == 128 or use_req_ring:
             state_loc = rid * ring_size + (pos % ring_size)
         else:
             loc = tl.load(
@@ -182,6 +183,7 @@ def triton_create_paged_compress_data(
     extend_seq_lens: torch.Tensor,
     req_to_token: torch.Tensor,
     full_to_swa_index_mapping: torch.Tensor,
+    use_req_ring: bool = False,
     block: int = 128,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     batch_size = req_pool_indices.shape[0]
@@ -205,6 +207,7 @@ def triton_create_paged_compress_data(
         stride_out_1_1=out_1.stride(1),  # type: ignore
         compress_ratio=compress_ratio,  # type: ignore
         is_overlap=1 if is_overlap else 0,  # type: ignore
+        use_req_ring=1 if use_req_ring else 0,  # type: ignore
         swa_page_size=swa_page_size,  # type: ignore
         ring_size=ring_size,  # type: ignore
         BLOCK=block,  # type: ignore

python/sglang/kernels/ops/attention/dsv4/compress.py

@@ -162,14 +162,15 @@ def generate(
         seq_lens: torch.Tensor,
         swa_page_size: int,
         ring_size: int,
+        use_req_ring: bool = False,
     ) -> CompressorDecodePlan:
         if _is_xpu:
             fn = plan_compress_decode
         else:
             module = _jit_compress_plan_module()
             fn = module.plan_decode
 
-        plan_d = fn(
+        args = (
             req_pool_indices,
             req_to_token,
             full_to_state,
@@ -178,6 +179,10 @@ def generate(
             int(swa_page_size),
             int(ring_size),
         )
+        assert not (_is_xpu and use_req_ring), (
+            "use_req_ring is not supported by the XPU compress plan builder"
+        )
+        plan_d = fn(*args) if _is_xpu else fn(*args, bool(use_req_ring))
         return CompressorDecodePlan(compress_ratio, torch.from_dlpack(plan_d))
 
     @staticmethod
@@ -247,6 +252,7 @@ def generate(
         ring_size: int,
         num_q_tokens: int,
         use_cuda_graph: bool = False,
+        use_req_ring: bool = False,
     ) -> CompressorPrefillPlan:
         is_gpu_input = seq_lens.device.type in ["cuda", "xpu"]
         pin_buffer = torch.empty(
@@ -274,7 +280,7 @@ def generate(
             module = _jit_compress_plan_module()
             fn = module.plan_prefill
 
-        plan_c, plan_w = fn(
+        args = (
             req_pool_indices,
             req_to_token,
             full_to_state,
@@ -285,7 +291,14 @@ def generate(
             int(compress_ratio),
             int(swa_page_size),
             int(ring_size),
-            bool(use_cuda_graph),
+        )
+        assert not (_is_xpu and use_req_ring), (
+            "use_req_ring is not supported by the XPU compress plan builder"
+        )
+        plan_c, plan_w = (
+            fn(*args, bool(use_cuda_graph))
+            if _is_xpu
+            else fn(*args, bool(use_req_ring), bool(use_cuda_graph))
         )
         return CompressorPrefillPlan(
             compress_ratio,

python/sglang/srt/disaggregation/decode.py

@@ -27,7 +27,7 @@
 from concurrent.futures import Future
 from dataclasses import dataclass
 from http import HTTPStatus
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple
 
 import numpy as np
 import torch
@@ -74,6 +74,7 @@
 from sglang.srt.managers.schedule_policy import match_prefix_for_req
 from sglang.srt.managers.utils import GenerationBatchResult
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+from sglang.srt.mem_cache.allocator.swa import is_swa_req_ring
 from sglang.srt.mem_cache.base_prefix_cache import (
     BasePrefixCache,
     DecLockRefParams,
@@ -139,6 +140,9 @@ class DecodeReqToTokenPool:
     #running <= 8, #pre-allocated + #transfer <= pre_alloc_size, so we can use the free memory to pre-allocate requests to unblock prefill.
     """
 
+    # Mirrors ReqToTokenPool.register_on_alloc_rows.
+    _on_alloc_rows: Optional[Callable[[List[int]], None]] = None
+
     def __init__(
         self,
         size: int,
@@ -204,6 +208,8 @@ def alloc(self, reqs: List[Req]) -> Optional[List[int]]:
             return None
         select_index = self.free_slots[:need_size]
         self.free_slots = self.free_slots[need_size:]
+        if self._on_alloc_rows is not None and select_index:
+            self._on_alloc_rows(select_index)
         offset = 0
         for r in reqs:
             if not r.kv.holds_kv:
@@ -221,6 +227,10 @@ def clear(self):
         self.free_slots = list(range(1, self._alloc_size))
         self.req_generation.zero_()
 
+    def register_on_alloc_rows(self, hook: Callable[[List[int]], None]) -> None:
+        assert self._on_alloc_rows is None
+        self._on_alloc_rows = hook
+
 
 class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):
     def __init__(
@@ -1711,7 +1721,13 @@ def _swa_tail_allocatable_token_budget(
         window_size = self.scheduler.sliding_window_size or 0
         swa_total = self.token_to_kv_pool_allocator.size_swa
         swa_available = self.token_to_kv_pool_allocator.swa_available_size()
-        swa_evictable = self.tree_cache.swa_evictable_size()
+        # Per-request SWA ring: cached prefixes still report swa_evictable, but
+        # evicting them frees no ring space.
+        swa_evictable = (
+            0
+            if is_swa_req_ring(self.token_to_kv_pool_allocator)
+            else self.tree_cache.swa_evictable_size()
+        )
         swa_used = swa_total - swa_available - swa_evictable
         swa_growth_potential = max(0, n_active * window_size - swa_used)
         swa_reserved_tokens = min(reserved_tokens, swa_growth_potential)

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

@@ -118,6 +118,9 @@ def compress_extend_paged(
             assert isinstance(backend, DeepseekV4HipRadixBackend)
         token_to_kv_pool = backend.token_to_kv_pool
         assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool)
+        req_ring_state = self.ratio == 128 or (
+            self.ratio == 4 and token_to_kv_pool._unified_kv
+        )
 
         state_pool = self._get_state_pool(backend)
         prefix_lens = forward_batch.extend_prefix_lens_cpu
@@ -144,7 +147,7 @@ def compress_extend_paged(
             pre_state_indices = self.compute_state_len_indices(
                 seq_len=prefix_lens[i], ratio=self.ratio
             ).to(device)
-            if self.ratio == 128:
+            if req_ring_state:
                 state_loc = state_pool.translate_from_req_position_to_state_loc(
                     req_pool_indices[i], pre_state_indices
                 )
@@ -166,7 +169,7 @@ def compress_extend_paged(
             post_state_len = post_state_indices.size(0)
 
             assert post_state_len <= valid_kv_len
-            if self.ratio == 128:
+            if req_ring_state:
                 post_state_loc = state_pool.translate_from_req_position_to_state_loc(
                     req_pool_indices[i], post_state_indices
                 )
@@ -260,6 +263,9 @@ def compress_decode_paged(
         state_pool = self._get_state_pool(attn_backend)
         token_to_kv_pool = attn_backend.token_to_kv_pool
         assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool)
+        req_ring_state = self.ratio == 128 or (
+            self.ratio == 4 and token_to_kv_pool._unified_kv
+        )
         req_pool_indices = forward_batch.req_pool_indices
         req_to_token = attn_backend.req_to_token_pool.req_to_token
         seq_lens = forward_batch.seq_lens
@@ -271,7 +277,7 @@ def compress_decode_paged(
             seq_lens = seq_lens_2d.view(-1)
             req_pool_indices = req_pool_indices.repeat_interleave(draft_tokens)
 
-        if self.ratio == 128:
+        if req_ring_state:
             state_locs = state_pool.translate_from_req_position_to_state_loc(
                 req_pool_indices, seq_lens - 1
             )
@@ -286,7 +292,7 @@ def compress_decode_paged(
             -compress_bulk_len, 0, device=seq_lens.device
         )
         compress_indices.clamp_(min=-1)
-        if self.ratio == 128:
+        if req_ring_state:
             compress_indices_state = (
                 state_pool.translate_from_req_position_to_state_loc(
                     req_pool_indices[:, None], compress_indices