From c653f2e10ea21f2c3d844081dd9fe7deb2a7a475 Mon Sep 17 00:00:00 2001 From: yuttian1 Date: Sun, 6 Sep 2026 01:00:10 +0000 Subject: [PATCH 01/15] Reland "[AMD][DSV4] Fix unified-KV pool sizing and SWA ring accounting" (#30315) This reverts commit f5819b09bf6eb58c3f91381274e8fec994e4746a (#38163), restoring #30315 verbatim. The follow-up commit gates every remaining shared-path change behind the unified-KV switch. --- .../kernels/jit/csrc/deepseek_v4/c_plan.cuh | 119 ++++++----- .../sglang/kernels/ops/attention/dsv4/attn.py | 5 + .../kernels/ops/attention/dsv4/compress.py | 13 +- python/sglang/srt/disaggregation/decode.py | 7 + .../srt/layers/attention/dsv4/compress_hip.py | 16 +- .../srt/layers/attention/dsv4/compressor.py | 8 +- .../layers/attention/dsv4/compressor_v2.py | 3 + python/sglang/srt/managers/schedule_batch.py | 52 ++++- python/sglang/srt/managers/schedule_policy.py | 52 +++-- .../scheduler_components/invariant_checker.py | 34 ++-- .../pool_stats_observer.py | 19 +- python/sglang/srt/mem_cache/allocation.py | 14 +- .../srt/mem_cache/allocator/hisparse.py | 4 + python/sglang/srt/mem_cache/allocator/swa.py | 78 +++++++- .../srt/mem_cache/deepseek_v4_memory_pool.py | 49 ++++- .../srt/mem_cache/kv_cache_configurator.py | 60 +++--- .../srt/model_executor/pool_configurator.py | 186 ++++++++++++++---- .../kernels/ops/attention/test_c4_v2.py | 91 ++++++++- .../unit/managers/test_prefill_adder.py | 12 +- .../unit/mem_cache/test_dllm_fdfo_kv_reuse.py | 3 + .../mem_cache/test_dsv4_c4_state_lifecycle.py | 119 +++++++++++ .../unit/mem_cache/test_hisparse_allocator.py | 1 + .../test_swa_alloc_extend_page_estimation.py | 5 + 23 files changed, 791 insertions(+), 159 deletions(-) create mode 100644 test/registered/unit/mem_cache/test_dsv4_c4_state_lifecycle.py diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh index 7e7b1005ee38..e341d4ba4426 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh @@ -28,12 +28,14 @@ using R2T_T = int32_t; using F2S_T = int64_t; using IDX_T = int64_t; -/// NOTE: for the internal use, we pack the ragged and batch id, since both not exceed 65536 +/// NOTE: for the internal use, we pack the ragged and batch id, since both not +/// exceed 65536 SGL_DEVICE __host__ PlanW pack_w(uint32_t ragged_id, uint32_t batch_id, int32_t seq_len) { return {static_cast(ragged_id | batch_id << 16), seq_len}; } -/// NOTE: for the internal use, we pack the ragged and batch id, since both not exceed 65536 +/// NOTE: for the internal use, we pack the ragged and batch id, since both not +/// exceed 65536 SGL_DEVICE uint2 unpack_w(PlanW plan) { return {static_cast(plan.ragged_id), static_cast(plan.ragged_id >> 16)}; } @@ -47,9 +49,11 @@ struct Prefill0Params { uint32_t num_q_tokens; int32_t compress_ratio; int32_t swa_page_size; - /// \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. + /// \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 +71,7 @@ struct Prefill1Params { int32_t swa_page_size; int32_t ring_size; int32_t compress_ratio; + bool use_req_ring; }; struct DecodeParams { @@ -80,6 +85,7 @@ struct DecodeParams { int32_t swa_page_size; int32_t ring_size; int32_t compress_ratio; + bool use_req_ring; }; struct Prefill1ParamsLegacy { @@ -155,7 +161,8 @@ __global__ __launch_bounds__(1024, 1) // counter_w = 0; } // === Stage B: min/max(extend_len) for MTP-uniform detection === - // For min, treat threads outside `batch_size` as +inf so they don't pull the min down. + // For min, treat threads outside `batch_size` as +inf so they don't pull the + // min down. const uint32_t e_for_max = static_cast(extend_len); const uint32_t e_for_min = (tx < params.batch_size) ? e_for_max : 0xFFFFFFFFu; warp_max[warp_id] = warp::reduce_max(e_for_max); @@ -168,17 +175,19 @@ __global__ __launch_bounds__(1024, 1) // __syncthreads(); const auto num_q = params.num_q_tokens; - // MTP-uniform: every batch shares the same small extend_len `E`, so we can decompose - // a global token id `k` into (batch_id, j) = (k / E, k % E) and skip the per-batch loop. + // MTP-uniform: every batch shares the same small extend_len `E`, so we can + // decompose a global token id `k` into (batch_id, j) = (k / E, k % E) and + // skip the per-batch loop. const bool is_mtp_extend = (s_min_extend == s_max_extend) && (s_max_extend > 0) && (s_max_extend <= 32); // === Stage C: emit valid plans, slot allocation via shared-mem atomicAdd === if (is_mtp_extend) { - // Path 1: token-driven. Each global token id maps to exactly one (batch_id, j). + // Path 1: token-driven. Each global token id maps to exactly one (batch_id, + // j). const uint32_t E = s_max_extend; - // num_q is the padded buffer size (graph bucket), not the work size: cap the - // loop at the real token count so batch_id = k / E stays < batch_size on an - // underfilled replay; Stage D pads [counter, num_q) with invalid. + // num_q is the padded buffer size (graph bucket), not the work size: cap + // the loop at the real token count so batch_id = k / E stays < batch_size + // on an underfilled replay; Stage D pads [counter, num_q) with invalid. const uint32_t num_real_q = params.batch_size * E; for (uint32_t k = tx; k < num_real_q; k += block_size) { const uint32_t batch_id = k / E; @@ -203,15 +212,15 @@ __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); } } } else { - // Path 2: general prefill (long extend_len). Iterate batches in an outer loop; - // the whole block sweeps each batch's tokens in parallel. + // Path 2: general prefill (long extend_len). Iterate batches in an outer + // loop; the whole block sweeps each batch's tokens in parallel. uint32_t base_e = 0; for (uint32_t batch_id = 0; batch_id < params.batch_size; ++batch_id) { const int32_t pl = s_prefix_len[batch_id]; @@ -236,7 +245,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(batch_id), position + 1); @@ -270,7 +279,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(rid * params.ring_size + position % params.ring_size); }; @@ -283,9 +292,9 @@ __global__ void plan_compress_prefill_kernel_1(const Prefill1Params params) { const auto position_1 = static_cast(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 +316,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(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 +338,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(rid * params.ring_size + position % params.ring_size); }; const auto seq_len = static_cast(params.seq_ptr[idx]); @@ -338,10 +347,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]; @@ -366,8 +375,10 @@ __global__ void plan_compress_prefill_legacy_kernel(const Prefill1ParamsLegacy p auto plan_w = idx < params.num_w ? params.plan_w[idx] : PlanW::invalid(); /// Per-request ring buffer slot translation: - /// - c4: page = rid * 2 + (position / 4) % 2; slot = page * 4 + position % 4 - /// - c128: page = rid; slot = rid * 128 + position % 128 + /// - c4: page = rid * 2 + (position / 4) % 2; slot = page * 4 + position % + /// 4 + /// - c128: page = rid; slot = rid * 128 + position % + /// 128 const auto legacy_compute_page = [&](int32_t rid, int32_t position) { if (params.compress_ratio == 4) return rid * 2 + ((position / 4) & 1); return rid; // c128 @@ -393,7 +404,8 @@ __global__ void plan_compress_prefill_legacy_kernel(const Prefill1ParamsLegacy p if (!plan_w.is_invalid()) { const auto [ragged_id, batch_id] = unpack_w(plan_w); const auto rid = static_cast(params.rid_ptr[batch_id]); - // `write_loc` carries (position + 1) at this stage; may not be ratio-aligned + // `write_loc` carries (position + 1) at this stage; may not be + // ratio-aligned const auto position = static_cast(plan_w.write_loc) - 1; plan_w.ragged_id = ragged_id; plan_w.write_loc = legacy_compute_loc(rid, position); @@ -407,8 +419,10 @@ __global__ void plan_compress_decode_legacy_kernel(const DecodeParamsLegacy para const auto idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= params.batch_size) return; /// Per-request ring buffer slot translation: - /// - c4: page = rid * 2 + (position / 4) % 2; slot = page * 4 + position % 4 - /// - c128: page = rid; slot = rid * 128 + position % 128 + /// - c4: page = rid * 2 + (position / 4) % 2; slot = page * 4 + position % + /// 4 + /// - c128: page = rid; slot = rid * 128 + position % + /// 128 const auto legacy_compute_page = [&](int32_t rid, int32_t position) { if (params.compress_ratio == 4) return rid * 2 + ((position / 4) & 1); return rid; // c128 @@ -447,7 +461,8 @@ using PrefillPlan = tvm::ffi::Tuple; * @param compress_plan `[num_q_tokens, 16]` uint8 (output) * @param write_plan `[num_q_tokens, 8]` uint8 (output) * @param compress_ratio 4 for c4, 128 for c128 - * @param use_cuda_graph Whether the plans will be used with cuda graph (affects padding) + * @param use_cuda_graph Whether the plans will be used with cuda graph (affects + * padding) * @return (compress plan tensor, write plan tensor) */ inline PrefillPlan plan_compress_prefill( @@ -461,6 +476,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,27 +519,29 @@ inline PrefillPlan plan_compress_prefill( const auto batch_size = static_cast(B.unwrap()); constexpr auto kMaxTokens = static_cast(std::numeric_limits::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); - // Write pad: trailing tokens kept resident so a verify batch's committed tail survives - // any accept length. Zero without speculation -- nothing rolls back, and the ring is - // then exactly one window wide. Otherwise the ring bounds it: a write at `w` aliases - // onto `w - ring_size`, and the earliest position a future compression still needs is - // `prefix_len - window_size + 2` (the next batch commits >= 1 token, and `run_prefill` - // launches the compress kernel before the write kernel, so a batch's own compressions - // read the pre-write ring). Padding past the extend range is harmless: the loops only - // span `[prefix_len, seq_len)`. + // Write pad: trailing tokens kept resident so a verify batch's committed tail + // survives any accept length. Zero without speculation -- nothing rolls back, + // and the ring is then exactly one window wide. Otherwise the ring bounds it: + // a write at `w` aliases onto `w - ring_size`, and the earliest position a + // future compression still needs is `prefix_len - window_size + 2` (the next + // batch commits >= 1 token, and `run_prefill` launches the compress kernel + // before the write kernel, so a batch's own compressions read the pre-write + // ring). Padding past the extend range is harmless: the loops only span + // `[prefix_len, seq_len)`. const auto mtp_pad = ring_size > window_size ? ring_size - window_size + 2 : 0; const auto device = device_.unwrap(); const auto stream = LaunchKernel::resolve_device(device); if (cpu_or_gpu.unwrap().device_type == kDLGPU) { - // GPU input path: kernel0 builds the (CPU-loop-equivalent) plan metadata directly - // on device, padding to num_q_tokens with invalid; kernel_1 then finalizes the - // SWA-translated read/write locations. Used for MTP / cuda-graph capture where - // a host sync would be expensive. + // GPU input path: kernel0 builds the (CPU-loop-equivalent) plan metadata + // directly on device, padding to num_q_tokens with invalid; kernel_1 then + // finalizes the SWA-translated read/write locations. Used for MTP / + // cuda-graph capture where a host sync would be expensive. RuntimeCheck(batch_size <= kMaxPrefillBatchSize, "GPU plan only support batch size up to ", kMaxPrefillBatchSize); auto C = ffi::empty({num_q_tokens, sizeof(PlanC)}, kDLUInt8, device); auto W = ffi::empty({num_q_tokens, sizeof(PlanW)}, kDLUInt8, device); @@ -537,9 +555,11 @@ 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. + // kernel_1 sees the already-padded buffers, so num_c == num_w == num_padded + // == num_q_tokens. const auto params1 = Prefill1Params{ .plan_c = static_cast(C.data_ptr()), .plan_w = static_cast(W.data_ptr()), @@ -555,6 +575,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 +603,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 +652,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 +667,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(); @@ -667,6 +690,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(B.unwrap()); const auto device = device_.unwrap(); auto D = ffi::empty({batch_size, sizeof(PlanD)}, kDLUInt8, device); @@ -681,6 +705,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); diff --git a/python/sglang/kernels/ops/attention/dsv4/attn.py b/python/sglang/kernels/ops/attention/dsv4/attn.py index f973bf916290..b98dc922d161 100644 --- a/python/sglang/kernels/ops/attention/dsv4/attn.py +++ b/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, @@ -135,6 +136,8 @@ def create_paged_compress_data_kernel( pos = tl.maximum(pos, 0) if compress_ratio == 128: state_loc = rid * ring_size + (pos % ring_size) + elif use_req_ring: + state_loc = rid * ring_size + (pos % ring_size) else: loc = tl.load( req_to_token_ptr @@ -182,6 +185,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 +209,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 diff --git a/python/sglang/kernels/ops/attention/dsv4/compress.py b/python/sglang/kernels/ops/attention/dsv4/compress.py index a9e915d67887..f126ba175e1a 100644 --- a/python/sglang/kernels/ops/attention/dsv4/compress.py +++ b/python/sglang/kernels/ops/attention/dsv4/compress.py @@ -162,6 +162,7 @@ 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 @@ -169,7 +170,7 @@ def generate( 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,7 @@ def generate( int(swa_page_size), int(ring_size), ) + 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 +249,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 +277,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 +288,11 @@ def generate( int(compress_ratio), int(swa_page_size), int(ring_size), - bool(use_cuda_graph), + ) + 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, diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 03cd52a33703..06b4384b5313 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1768,11 +1768,18 @@ def _pre_alloc( if total_prefix_len is None: total_prefix_len = prefix_len + is_new_req_slot = req.kv.req_pool_idx is None req_pool_indices = self.req_to_token_pool.alloc([req]) assert req_pool_indices is not None, ( "req_pool_indices is full! There is a bug in memory estimation." ) + if is_new_req_slot: + clear_c4_req_states = getattr( + self.token_to_kv_pool, "clear_c4_req_states", None + ) + if clear_c4_req_states is not None: + clear_c4_req_states(req_pool_indices) fill_len = self._pre_alloc_fill_len(req) req.kv.kv_committed_len = fill_len diff --git a/python/sglang/srt/layers/attention/dsv4/compress_hip.py b/python/sglang/srt/layers/attention/dsv4/compress_hip.py index 225a008a3c4b..7ae66e27bb3c 100644 --- a/python/sglang/srt/layers/attention/dsv4/compress_hip.py +++ b/python/sglang/srt/layers/attention/dsv4/compress_hip.py @@ -144,7 +144,9 @@ 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 self.ratio == 128 or ( + self.ratio == 4 and getattr(token_to_kv_pool, "_unified_kv", False) + ): state_loc = state_pool.translate_from_req_position_to_state_loc( req_pool_indices[i], pre_state_indices ) @@ -166,7 +168,9 @@ def compress_extend_paged( post_state_len = post_state_indices.size(0) assert post_state_len <= valid_kv_len - if self.ratio == 128: + if self.ratio == 128 or ( + self.ratio == 4 and getattr(token_to_kv_pool, "_unified_kv", False) + ): post_state_loc = state_pool.translate_from_req_position_to_state_loc( req_pool_indices[i], post_state_indices ) @@ -271,7 +275,9 @@ 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 self.ratio == 128 or ( + self.ratio == 4 and getattr(token_to_kv_pool, "_unified_kv", False) + ): state_locs = state_pool.translate_from_req_position_to_state_loc( req_pool_indices, seq_lens - 1 ) @@ -286,7 +292,9 @@ def compress_decode_paged( -compress_bulk_len, 0, device=seq_lens.device ) compress_indices.clamp_(min=-1) - if self.ratio == 128: + if self.ratio == 128 or ( + self.ratio == 4 and getattr(token_to_kv_pool, "_unified_kv", False) + ): compress_indices_state = ( state_pool.translate_from_req_position_to_state_loc( req_pool_indices[:, None], compress_indices diff --git a/python/sglang/srt/layers/attention/dsv4/compressor.py b/python/sglang/srt/layers/attention/dsv4/compressor.py index 9b213c1bac81..d5b8c29b867a 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor.py @@ -26,9 +26,7 @@ cp_all_gather_rerange_finish, cp_all_gather_rerange_launch, ) -from sglang.srt.mem_cache.deepseek_v4_compress_state import ( - CompressStatePool, -) +from sglang.srt.mem_cache.deepseek_v4_compress_state import CompressStatePool from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.model_executor.forward_context import get_attn_backend from sglang.srt.models.deepseek_v2 import _is_hip @@ -264,6 +262,7 @@ def create_paged_compressor_data( ) -> FusedCompressMetadata: swa_page_size = token_to_kv_pool.swa_page_size ring_size = token_to_kv_pool.get_ring_size(compress_ratio=compress_ratio) + use_req_ring = compress_ratio == 4 and token_to_kv_pool._unified_kv # assert ring_size % compress_ratio == 0 def clip_down(positions: torch.Tensor) -> torch.Tensor: @@ -273,6 +272,8 @@ def get_raw_loc(positions: torch.Tensor) -> torch.Tensor: positions = positions.masked_fill(positions < 0, 0) if compress_ratio == 128: state_loc = req_pool_indices * ring_size + positions % ring_size + elif use_req_ring: + state_loc = req_pool_indices * ring_size + positions % ring_size else: loc = req_to_token[req_pool_indices, positions] swa_loc = token_to_kv_pool.translate_loc_from_full_to_swa(loc) @@ -294,6 +295,7 @@ def get_raw_loc(positions: torch.Tensor) -> torch.Tensor: extend_seq_lens=extend_lens, req_to_token=req_to_token, full_to_swa_index_mapping=token_to_kv_pool.full_to_swa_index_mapping, + use_req_ring=use_req_ring, ) plan_kwargs: dict diff --git a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py index 707efb74d7e1..c3515413119b 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py @@ -441,6 +441,7 @@ def create_paged_compressor_data( swa_page_size = token_to_kv_pool.swa_page_size ring_size = token_to_kv_pool.get_ring_size(compress_ratio=compress_ratio) + use_req_ring = compress_ratio == 4 and token_to_kv_pool._unified_kv # NOTE: This is actually a proxy, which encounter some bug with tvm-ffi. # As a workaround, we use `.detach()` to get the real tensor. full_to_swa = token_to_kv_pool.full_to_swa_index_mapping.detach() @@ -467,6 +468,7 @@ def create_paged_compressor_data( full_to_state=full_to_swa, swa_page_size=swa_page_size, ring_size=ring_size, + use_req_ring=use_req_ring, num_q_tokens=num_q_tokens, use_cuda_graph=use_prefill_cuda_graph, ) @@ -479,6 +481,7 @@ def create_paged_compressor_data( seq_lens=seq_lens.to(torch.int64), swa_page_size=swa_page_size, ring_size=ring_size, + use_req_ring=use_req_ring, ) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index db3d02c2921a..e539e91e98fd 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -57,6 +57,7 @@ import logging import re import sys +import time from array import array from concurrent.futures import Future from enum import Enum, auto @@ -100,10 +101,7 @@ from sglang.srt.managers.scheduler_components.new_token_ratio_tracker import ( NewTokenRatioTracker, ) -from sglang.srt.mem_cache.allocation import ( - alloc_for_decode, - alloc_for_extend, -) +from sglang.srt.mem_cache.allocation import alloc_for_decode, alloc_for_extend from sglang.srt.mem_cache.allocation_sizing import get_alloc_reserve_per_decode from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( @@ -161,6 +159,9 @@ logger = logging.getLogger(__name__) +# Throttle for the unified-KV SWA bottleneck diagnostic (seconds). +_last_swa_bottleneck_log = 0.0 + ReturnHiddenStatesMode = Union[bool, Literal["last"]] @@ -3077,9 +3078,50 @@ def check_decode_mem(self, selected_indices: Optional[List[int]] = None): shortfalls retract gracefully instead of tripping fail-loud alloc errors.""" num_tokens = self.new_tokens_required_next_decode(selected_indices) - return self.token_to_kv_pool_allocator.check_decode_capacity( + allocator = self.token_to_kv_pool_allocator + ok = allocator.check_decode_capacity( num_tokens=num_tokens, tree_cache=self.tree_cache ) + if not ok and getattr(allocator.get_kvcache(), "_unified_kv", False): + self._log_unified_swa_bottleneck(allocator, num_tokens, selected_indices) + return ok + + def _log_unified_swa_bottleneck(self, allocator, num_tokens, selected_indices): + """Diagnostic (unified-KV only): when check_decode_mem is short, compare + the SWA token bookkeeping against the real per-slot ring utilization. + Throttled to avoid log floods during retract storms.""" + global _last_swa_bottleneck_log + now = time.monotonic() + if now - _last_swa_bottleneck_log < 1.0: + return + _last_swa_bottleneck_log = now + try: + full_avail = allocator.full_available_size() + swa_avail = allocator.swa_available_size() + reqs = ( + self.reqs + if selected_indices is None + else [self.reqs[i] for i in selected_indices] + ) + active_slots = len( + {int(r.kv.req_pool_idx) for r in reqs if r.kv.req_pool_idx is not None} + ) + unified = getattr(allocator.get_kvcache(), "unified_kv_pool", None) + if unified is not None: + num_slots = unified.num_slots + ring_util = ( + active_slots * unified.swa_ring_size / max(unified.swa_pages, 1) + ) + else: + num_slots, ring_util = -1, -1.0 + logger.warning( + "[SWA-BOTTLENECK] check_decode_mem short: " + f"need={num_tokens}, full_avail={full_avail}, swa_avail={swa_avail}, " + f"active_slots={active_slots}/{num_slots}, " + f"ring_util_upper={ring_util:.4f}" + ) + except Exception as e: # diagnostics must never break scheduling + logger.warning(f"[SWA-BOTTLENECK] logging failed: {e}") def retract_decode(self) -> Tuple[List[Req], float, List[Req]]: """Retract the decoding requests when there is not enough memory.""" diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 023071a6709d..a4fbf5e6e231 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -5,10 +5,7 @@ from sglang.srt.environ import envs from sglang.srt.managers.prefill_delayer import PrefillDelayerSinglePassExecutor -from sglang.srt.runtime_context import ( - get_disagg, - get_schedule, -) +from sglang.srt.runtime_context import get_disagg, get_schedule from sglang.srt.utils import get_bool_env_var, is_hip _ROUTING_KEY_POLICY_DEBUG_LOG = get_bool_env_var("SGLANG_ROUTING_KEY_POLICY_DEBUG_LOG") @@ -663,8 +660,16 @@ def rem_total_tokens(self): @property def rem_swa_tokens(self): + allocator = self.token_to_kv_pool_allocator + if getattr(allocator.get_kvcache(), "_unified_kv", False): + # Unified-KV: SWA is a per-request ring, not a tree-reusable token + # pool. swa_available_size() already reports ring capacity + # (free_slots * ring_cost). tree swa_evictable is in the old linear + # token unit and freeing it does not release ring space, so exclude + # it here to keep a single consistent accounting unit. + return allocator.swa_available_size() - self.rem_swa_token_offset return ( - self.token_to_kv_pool_allocator.swa_available_size() + allocator.swa_available_size() + self.tree_cache.swa_evictable_size() - self.rem_swa_token_offset ) @@ -711,6 +716,13 @@ def _swa_budget_for_req( where alloc = min(extend, rem_chunk); the min() cap keeps the two terms from double-counting extend, so budget <= extend + max_new_tokens + page. """ + allocator = self.token_to_kv_pool_allocator + if getattr(allocator.get_kvcache(), "_unified_kv", False): + # Unified-KV: each request occupies exactly one fixed SWA ring slot, + # independent of context / chunk length; a host-hit prefix reuses the + # same ring. Budget the fixed per-slot ring cost (paired with the + # ring-based swa_available_size on the allocator). + return allocator.swa_ring_cost_tokens if self.rem_chunk_tokens is not None: alloc = min(extend_input_len, self.rem_chunk_tokens) else: @@ -838,6 +850,9 @@ def _update_prefill_budget( max_new_tokens: int, retracted_stain: bool, mamba_gap_reserve: int = 0, + host_hit_len: int = 0, + storage_hit_len: int = 0, + is_chunked_continuation: bool = False, ): # TODO(lsyin): check this workaround logic, which only ensures the prefill will not out of memory, and may be too conservative extend_input_len = self.ceil_paged_tokens(extend_input_len) @@ -861,9 +876,17 @@ def _update_prefill_budget( self.rem_input_tokens -= extend_input_len if self.is_hybrid_swa: - self.rem_swa_token_offset += self._swa_budget_for_req( - extend_input_len, max_new_tokens + # Unified-KV: SWA is a fixed per-request ring slot reserved once at + # first admission and already reflected in swa_available_size() on + # later rounds. Charging it again on a chunked continuation would + # double-count the slot and over-throttle admission, so skip it. + _unified = getattr( + self.token_to_kv_pool_allocator.get_kvcache(), "_unified_kv", False ) + if not (_unified and is_chunked_continuation): + self.rem_swa_token_offset += self._swa_budget_for_req( + extend_input_len, max_new_tokens + ) if self.dllm_config is not None: self.rem_dllm_tokens -= extend_input_len @@ -998,9 +1021,15 @@ def add_chunked_req(self, req: Req): _rem_tokens = self._get_dllm_remain_tokens() else: _rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens)) - if self.is_hybrid_swa: + if self.is_hybrid_swa and not getattr( + self.token_to_kv_pool_allocator.get_kvcache(), "_unified_kv", False + ): # alloc_extend needs extend_num_tokens + page_size per request, - # so reserve one page here to avoid OOM + # so reserve one page here to avoid OOM. + # Unified-KV: rem_swa_tokens is ring capacity (free_slots * ring + # cost), not a linear per-chunk token budget, and this request's + # ring slot is already reserved -- mixing units here would wrongly + # truncate the chunk, so skip the SWA clamp. _rem_tokens = min( _rem_tokens, int(self.rem_swa_tokens) - self.page_size ) @@ -1039,6 +1068,7 @@ def add_chunked_req(self, req: Req): ), req.retracted_stain, mamba_gap_reserve=self._mamba_gap_budget_for_req(req), + is_chunked_continuation=True, ) # Return if chunked prefill not finished @@ -1248,7 +1278,7 @@ def add_one_req( self._swa_new_tokens(req), swa_host_hit_length=req.swa_host_hit_length, ) - if swa_needed >= self.rem_swa_tokens: + if swa_needed > self.rem_swa_tokens: if not self._swa_req_never_fits( real_input_tokens, self._swa_new_tokens(req), @@ -1284,7 +1314,7 @@ def add_one_req( self._swa_new_tokens(req), swa_host_hit_length=req.swa_host_hit_length, ) - if swa_needed >= self.rem_swa_tokens: + if swa_needed > self.rem_swa_tokens: if not self._swa_req_never_fits( real_input_tokens, self._swa_new_tokens(req), diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index ed24b842220c..fc861c589077 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -3,14 +3,7 @@ import logging from collections import deque from dataclasses import dataclass, field -from typing import ( - TYPE_CHECKING, - Callable, - Deque, - List, - Optional, - Tuple, -) +from typing import TYPE_CHECKING, Callable, Deque, List, Optional, Tuple import torch @@ -32,10 +25,7 @@ scheduler_stage_method, ) from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils.common import ( - ceil_align, - raise_error_or_warn, -) +from sglang.srt.utils.common import ceil_align, raise_error_or_warn from sglang.srt.utils.watchdog import WatchdogRaw if TYPE_CHECKING: @@ -152,6 +142,23 @@ def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str] def _check_swa_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]: allocator = self.token_to_kv_pool_allocator + kv = allocator.get_kvcache() + if getattr(kv, "_unified_kv", False): + # Unified-KV DSV4: SWA is a fixed per-request ring, reused per request + # and released together with the req_pool slot (which has its own + # leak check). swa_available_size() is deliberately non-binding (it + # always reports the full ring so it never throttles admission), and + # cached radix prefixes still report swa_evictable even though the + # completed request already freed its ring slot. The token-pool + # invariant (available + evictable + protected + session == total) + # therefore does not model this pool -- skip it to avoid a spurious + # leak. Ring-slot leaks are still caught by the req_to_token check. + return False, ( + "[swa] unified ring (leak-check skipped): " + f"available={ps.swa_available_size}, " + f"evictable={ps.swa_evictable_size}, " + f"total={self.swa_tokens_per_layer}" + ) swa_available = ps.swa_available_size if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): # Tri-pool: same floating-boundary phantom as the full pool -- use the @@ -387,7 +394,10 @@ def _add_owner(req_or_slot, label, rpi, committed, allocated): # Sub-allocators to check: a flat allocator is its own single sub; a # hybrid-SWA wrapper exposes full_attn_allocator + swa_attn_allocator. + # DSV4-HiSparse nests the real SWA allocator under logical_attn_allocator, + # so unwrap first (no-op for a plain/flat allocator). alloc = self.token_to_kv_pool_allocator + alloc = getattr(alloc, "logical_attn_allocator", alloc) sub_allocs = ( [alloc] if getattr(alloc, "free_pages", None) is not None diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index e1dd565ebd6d..e8f9a669750d 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -2,14 +2,7 @@ import dataclasses from dataclasses import dataclass -from typing import ( - TYPE_CHECKING, - Any, - Callable, - List, - Optional, - Tuple, -) +from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( UnifiedMambaSWATokenToKVPoolAllocator, @@ -301,6 +294,16 @@ def _get_swa_token_info(self) -> PoolStats: swa_available_size = allocator.swa_available_size() full_evictable_size = self.tree_cache.full_evictable_size() swa_evictable_size = self.tree_cache.swa_evictable_size() + # Unified-KV DSV4: SWA is a fixed per-request ring, released with the + # req_pool slot. Cached radix prefixes still report swa_evictable even + # though the completed request already freed its ring slot, and + # swa_available_size() is non-binding (always the full ring). Counting + # that evictable here would double-count against the ring and drive + # swa_num_used / swa_token_usage negative. The ring holds nothing + # evictable, so zero it out to keep the usage stats coherent. + _swa_kv = self.token_to_kv_pool_allocator.get_kvcache() + if getattr(_swa_kv, "_unified_kv", False): + swa_evictable_size = 0 full_num_used = self.full_tokens_per_layer - ( full_available_size + full_evictable_size ) diff --git a/python/sglang/srt/mem_cache/allocation.py b/python/sglang/srt/mem_cache/allocation.py index 5e109a4c897e..4cf251db0f8d 100644 --- a/python/sglang/srt/mem_cache/allocation.py +++ b/python/sglang/srt/mem_cache/allocation.py @@ -230,6 +230,7 @@ def alloc_req_slots( req_to_token_pool: ReqToTokenPool, reqs: list[Req], tree_cache: BasePrefixCache | None, + token_to_kv_pool=None, ) -> list[int]: """Allocate request slots from the pool. @@ -260,6 +261,7 @@ def alloc_req_slots( tree_cache.evict_for_alloc( EvictParams(num_tokens=0, mamba_num=mamba_num) ) + newly_allocated = [req.kv.req_pool_idx is None for req in reqs] req_pool_indices = req_to_token_pool.alloc(reqs) if req_pool_indices is None: raise RuntimeError( @@ -267,6 +269,13 @@ def alloc_req_slots( "Please set a smaller number for `--max-running-requests`. " f"{req_to_token_pool.available_size()=}, {num_reqs=}, " ) + + new_req_pool_indices = [ + idx for idx, is_new in zip(req_pool_indices, newly_allocated) if is_new + ] + clear_c4_req_states = getattr(token_to_kv_pool, "clear_c4_req_states", None) + if new_req_pool_indices and clear_c4_req_states is not None: + clear_c4_req_states(new_req_pool_indices) return req_pool_indices @@ -311,7 +320,10 @@ def alloc_for_extend( # Allocate req slots (raises RuntimeError if the pool is exhausted) req_pool_indices = alloc_req_slots( - batch.req_to_token_pool, batch.reqs, batch.tree_cache + batch.req_to_token_pool, + batch.reqs, + batch.tree_cache, + token_to_kv_pool=batch.token_to_kv_pool_allocator.get_kvcache(), ) req_pool_indices_cpu = torch.tensor( req_pool_indices, dtype=torch.int64, pin_memory=pin_memory diff --git a/python/sglang/srt/mem_cache/allocator/hisparse.py b/python/sglang/srt/mem_cache/allocator/hisparse.py index 5647154f70b1..44047c9b5c02 100644 --- a/python/sglang/srt/mem_cache/allocator/hisparse.py +++ b/python/sglang/srt/mem_cache/allocator/hisparse.py @@ -343,6 +343,10 @@ def debug_print(self) -> str: def get_kvcache(self): return self._kvcache + @property + def swa_ring_cost_tokens(self) -> int: + return self.logical_attn_allocator.swa_ring_cost_tokens + def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor): return self.logical_attn_allocator.translate_loc_from_full_to_swa(kv_indices) diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index 4b5c61812928..c5abd1d756e7 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -1,3 +1,5 @@ +import logging + import torch from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator @@ -8,6 +10,8 @@ from sglang.srt.utils.common import get_num_new_pages from sglang.srt.utils.invariants import Bucket, Invariant, IsTrue, expect +logger = logging.getLogger(__name__) + _is_npu = is_npu() if _is_npu: @@ -37,6 +41,7 @@ def __init__( device: str, kvcache: BaseSWAKVPool, need_sort: bool, + req_to_token_pool=None, ): assert isinstance(kvcache, BaseSWAKVPool) self._size_full = size @@ -104,10 +109,46 @@ def __init__( self.swa_free_group = [] self._kvcache = kvcache + + # Unified-KV (DSV4): SWA is a per-request ring addressed by state_slot + # (== req_pool_idx) + position inside the DSV4 kernels. The paged SWA + # indices / full_to_swa_index_mapping produced here are NOT consumed on + # that path, so treating SWA as a linearly-consumed token pool + # over-throttles admission and decode retract. Instead account for it as + # a fixed per-request ring slot; the real bound is concurrency + # (num_req_slots), already enforced by req_to_token_pool / + # max_running_requests. + self._unified = getattr(kvcache, "_unified_kv", False) + self._req_to_token_pool = req_to_token_pool + if self._unified: + ring_size = getattr(kvcache, "unified_swa_ring_size", self.page_size) + self._swa_ring_cost = ( + (ring_size + self.page_size - 1) // self.page_size + ) * self.page_size + logger.info( + "[SWA-BOOKKEEPING] unified ring accounting enabled: " + f"num_slots={getattr(kvcache, 'num_req_slots', '?')}, " + f"swa_ring_size={ring_size}, " + f"ring_cost_tokens={self._swa_ring_cost}, " + f"unified_swa_pages={getattr(kvcache, 'unified_swa_pages', '?')} | " + f"legacy paged size_swa={self._size_swa} (bypassed)" + ) + else: + self._swa_ring_cost = 0 + self.clear() self._kvcache.register_mapping(self.full_to_swa_index_mapping) + @property + def swa_ring_cost_tokens(self) -> int: + """Unified: paged SWA cost of one request's ring slot (0 otherwise).""" + return self._swa_ring_cost + def available_size(self): + if self._unified: + # The SWA ring is pre-allocated per slot and reused by decode, so it + # never constrains token growth; full attention is the real limiter. + return self.full_attn_allocator.available_size() return min( self.full_attn_allocator.available_size(), self.swa_attn_allocator.available_size(), @@ -117,6 +158,12 @@ def full_available_size(self): return self.full_attn_allocator.available_size() def swa_available_size(self): + if self._unified: + # Ring-based availability: free request slots * per-slot ring cost. + # Fall back to non-binding if the req pool wasn't wired in. + if self._req_to_token_pool is None: + return self.full_attn_allocator.available_size() + return self._req_to_token_pool.available_size() * self._swa_ring_cost return self.swa_attn_allocator.available_size() # Slot-conservation views for the leak invariant. On the non-shared allocator @@ -171,11 +218,15 @@ def alloc(self, need_size: int): return alloc_full_indices def new_pages_available(self, num_full_pages: int, num_swa_pages: int) -> bool: - return ( + full_ok = ( num_full_pages <= self.full_attn_allocator.available_size() // self.page_size - and num_swa_pages - <= self.swa_attn_allocator.available_size() // self.page_size + ) + if self._unified: + # SWA ring rows are pre-allocated per slot; no per-token SWA paging. + return full_ok + return full_ok and ( + num_swa_pages <= self.swa_attn_allocator.available_size() // self.page_size ) def alloc_extend( @@ -195,6 +246,20 @@ def alloc_extend( if not self.new_pages_available(num_new_pages, num_new_pages): return None + if self._unified: + # Unified SWA ring is slot-addressed and not paged here: allocate only + # the full-attention KV and skip the vestigial SWA allocator / mapping + # (unused by the DSV4 kernels). + return self.full_attn_allocator.alloc_extend( + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + last_loc, + extend_num_tokens, + num_new_pages=num_new_pages, + ) + swa_last_loc = self.translate_loc_from_full_to_swa(last_loc) alloc_full_indices = self.full_attn_allocator.alloc_extend( @@ -291,6 +356,13 @@ def alloc_decode( last_loc: torch.Tensor, # last_loc for full layers ): assert self.page_size > 1 + if self._unified: + # See alloc_extend: unified SWA ring is slot-addressed, allocate full + # only and skip the vestigial SWA allocator / mapping. + return self.full_attn_allocator.alloc_decode( + seq_lens, seq_lens_cpu, last_loc + ) + swa_last_loc = self.translate_loc_from_full_to_swa(last_loc) alloc_full_indices = self.full_attn_allocator.alloc_decode( diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index 8287abbf8608..991a8a5fb13b 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -2,7 +2,7 @@ import logging from contextlib import nullcontext -from typing import List, Literal, NamedTuple, Optional, Tuple +from typing import List, Literal, NamedTuple, Optional, Sequence, Tuple import torch @@ -564,6 +564,11 @@ def __init__( self.c4_size = c4_size self.c4_logical_size = c4_logical_size self.c128_size = c128_size + # Keep the legacy SWA-addressed pool large enough on non-unified paths. + # Unified request-addressed sizing is set exactly after resolving the + # unified-kv gate below. + c4_ring_size = self.get_ring_size(4) + c4_state_pool_size = max(c4_state_pool_size, self.num_req_slots * c4_ring_size) self.c4_state_pool_size = c4_state_pool_size c128_ring_size = self.get_ring_size(128) if ONLINE_C128: @@ -624,6 +629,9 @@ def __init__( ) self._unified_kv = is_unified_kv_triton() + if self._unified_kv: + # Unified C4 state is request-scoped: no SWA-derived over-allocation. + self.c4_state_pool_size = self.num_req_slots * c4_ring_size if self._unified_kv: self.swa_kv_pool = None @@ -1050,6 +1058,37 @@ def get_online_c128_mtp_pending_seq_lens(self) -> torch.Tensor: assert self.online_c128_mtp_pending_seq_lens is not None return self.online_c128_mtp_pending_seq_lens + def clear_c4_req_states(self, req_pool_indices: Sequence[int]) -> None: + """Reset newly allocated unified C4 attention and indexer state rings. + + Only the request-owned rows are touched. The extra sentinel/ring padding + allocated by :class:`CompressStatePool` remains intact. + """ + if not self._unified_kv or not req_pool_indices: + return + + pools = [ + pool + for pool in self.compress_state_pools + self.indexer_compress_state_pools + if pool is not None and pool.ratio == 4 + ] + if not pools: + return + + ring_size = self.get_ring_size(4) + device = pools[0].kv_score_buffer.kv_score.device + req_indices = torch.as_tensor(req_pool_indices, dtype=torch.long, device=device) + state_locs = ( + req_indices[:, None] * ring_size + + torch.arange(ring_size, dtype=torch.long, device=device) + ).flatten() + + for pool in pools: + state = pool.kv_score_buffer.kv_score + half = state.shape[-1] // 2 + state[state_locs, :half] = 0 + state[state_locs, half:] = float("-inf") + def clear_c128_req_state(self, req_pool_idx: int) -> None: """Reset request-scoped C128 state for one req slot.""" for pool in self.compress_state_pools: @@ -1076,7 +1115,13 @@ def clear_unaccepted_c128_draft_states( accept_lens: torch.Tensor, num_draft_tokens: int, ) -> None: - """Clear offline C128 ring slots written for rejected speculative tokens.""" + """Clear offline C128 ring slots written for rejected speculative tokens. + + C4 needs no equivalent cleanup: draft states are written in position order, + and every rejected position is overwritten before it can become the prior + state of a later accepted token. C128 cleanup is required because its + compression boundary can consume a previously written draft slot directly. + """ if ONLINE_C128 or num_draft_tokens <= 1 or req_pool_indices.numel() == 0: return diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 1c8df49118b5..a4077d892caf 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -204,9 +204,7 @@ def _pp_local_per_request_bytes( from sglang.srt.model_executor.model_runner_components.spec_aux_hidden_state import ( SpecAuxHiddenStateConfig, ) - from sglang.srt.model_executor.pool_configurator import ( - MemoryPoolConfig, - ) + from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig class KVCacheConfigResult(msgspec.Struct, frozen=True, kw_only=True): @@ -336,6 +334,35 @@ def configure(self, *, pre_model_load_memory: int) -> KVCacheConfigResult: token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, ) + swa_max_total_num_tokens = sizes.swa_max_total_num_tokens + # Unified-KV DSV4: SWA is a fixed per-request ring, so the allocator + # reports ring capacity (free_req_slots * ring_cost) from + # swa_available_size(), while swa_max_total_num_tokens was sized from + # the (vestigial, unallocated) full_token-scaled SWA pool. The idle + # pool-leak invariant requires swa total == swa available, so + # reconcile the reported SWA total to the allocator's actual idle ring + # capacity. Safe: on unified_kv swa_kv_pool is None, so no real buffer + # is resized -- this only fixes token accounting / usage reporting. + if ( + self.is_hybrid_swa + and not self.is_draft_worker + and getattr(pools.token_to_kv_pool, "_unified_kv", False) + ): + alloc = pools.token_to_kv_pool_allocator + if hasattr(alloc, "swa_available_size"): + ring_capacity = int(alloc.swa_available_size()) + # Only reconcile downward to the (smaller) ring capacity. A + # value >= the current total means swa_available_size() hit a + # non-binding fallback (e.g. req_to_token pool not wired), in + # which case leave the reported total untouched. + if 0 < ring_capacity < swa_max_total_num_tokens: + logger.info( + "Unified-KV: reconciling swa_max_total_num_tokens " + f"{swa_max_total_num_tokens} -> {ring_capacity} " + "(fixed per-request SWA ring capacity)." + ) + swa_max_total_num_tokens = ring_capacity + logger.info( f"Memory pool end. " f"avail mem={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB" @@ -345,7 +372,7 @@ def configure(self, *, pre_model_load_memory: int) -> KVCacheConfigResult: max_total_num_tokens=sizes.max_total_num_tokens, max_running_requests=sizes.max_running_requests, full_max_total_num_tokens=sizes.full_max_total_num_tokens, - swa_max_total_num_tokens=sizes.swa_max_total_num_tokens, + swa_max_total_num_tokens=swa_max_total_num_tokens, req_to_token_pool=pools.req_to_token_pool, token_to_kv_pool=pools.token_to_kv_pool, token_to_kv_pool_allocator=pools.token_to_kv_pool_allocator, @@ -1006,9 +1033,7 @@ def _build_hybrid_mamba_decode_req_pool( extra_max_context_len: int, pre_alloc_size: int, ) -> ReqToTokenPool: - from sglang.srt.disaggregation.decode import ( - HybridMambaDecodeReqToTokenPool, - ) + from sglang.srt.disaggregation.decode import HybridMambaDecodeReqToTokenPool req_to_token_pool = HybridMambaDecodeReqToTokenPool( size=max_num_reqs, @@ -1296,9 +1321,7 @@ def _build_dsv4_kv_pool( assert swa_page_size == 256, "In paged swa mode, page_size must be 256." if self.is_draft_worker: - from sglang.srt.models.deepseek_v4_nextn import ( - COMPRESS_RATIO_NEXTN_LAYER, - ) + from sglang.srt.models.deepseek_v4_nextn import COMPRESS_RATIO_NEXTN_LAYER compression_ratios = [ COMPRESS_RATIO_NEXTN_LAYER @@ -1413,9 +1436,7 @@ def _build_ascend_swa_kv_pool( full_max_total_num_tokens: Optional[int], swa_max_total_num_tokens: Optional[int], ) -> KVCache: - from sglang.srt.hardware_backend.npu.memory_pool_npu import ( - NPUMHATokenToKVPool, - ) + from sglang.srt.hardware_backend.npu.memory_pool_npu import NPUMHATokenToKVPool kwargs = {} if self.is_hybrid_swa_compress: @@ -1481,9 +1502,7 @@ def _build_ascend_minimax_sparse_kv_pool( def _build_ascend_mla_kv_pool( self, *, max_total_num_tokens: int, is_dsa_model: bool ) -> KVCache: - from sglang.srt.hardware_backend.npu.memory_pool_npu import ( - NPUMLATokenToKVPool, - ) + from sglang.srt.hardware_backend.npu.memory_pool_npu import NPUMLATokenToKVPool token_to_kv_pool = NPUMLATokenToKVPool( max_total_num_tokens, @@ -1501,9 +1520,7 @@ def _build_ascend_mla_kv_pool( return token_to_kv_pool def _build_ascend_mha_kv_pool(self, *, max_total_num_tokens: int) -> KVCache: - from sglang.srt.hardware_backend.npu.memory_pool_npu import ( - NPUMHATokenToKVPool, - ) + from sglang.srt.hardware_backend.npu.memory_pool_npu import NPUMHATokenToKVPool token_to_kv_pool = NPUMHATokenToKVPool( max_total_num_tokens, @@ -1943,12 +1960,11 @@ def _build_token_to_kv_pool_allocator( device=self.device, kvcache=token_to_kv_pool, need_sort=need_sort, + req_to_token_pool=req_to_token_pool, ) else: if get_memory().enable_hisparse: - from sglang.srt.mem_cache.sparsity import ( - parse_hisparse_config, - ) + from sglang.srt.mem_cache.sparsity import parse_hisparse_config hisparse_cfg = parse_hisparse_config() token_to_kv_pool_allocator = HiSparseTokenToKVPoolAllocator( diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 3a5e7d3da623..e21b61f56746 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -408,9 +408,7 @@ def _compute_dsa_indexer_cell_size( indexer_ratio = parse_hisparse_config().host_to_device_ratio - from sglang.srt.mem_cache.kv_cache_configurator import ( - _should_elide_dsa_index_k, - ) + from sglang.srt.mem_cache.kv_cache_configurator import _should_elide_dsa_index_k if allocate_all_layers or not _should_elide_dsa_index_k( is_draft_worker=kvc.is_draft_worker @@ -847,7 +845,10 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): Splits available memory across full / swa / c4 / c128 + c4_state / c128_state pools. coeff is bytes_per_full_token (inflated by (T+D)/T when speculative - decode reserves a draft worker, mirroring dflash's cell_size scaling); bias = 0. + decode reserves a draft worker, mirroring dflash's cell_size scaling). The + bias is the sum of request-scoped fixed pools that do not scale with + full_token: the c128 state pool and, on the unified_kv path, the fixed SWA + per-request ring (bf16, see _fixed_swa_bytes). """ def __init__(self, kvc: KVCacheConfigurator): @@ -904,6 +905,27 @@ def __init__(self, kvc: KVCacheConfigurator): self.num_layers_ca4 = sum(1 for r in self.compression_ratios if r == 4) self.num_layers_ca128 = sum(1 for r in self.compression_ratios if r == 128) + # Unified-KV uses a different physical layout than the fp8 path: + # * KV is stored bf16 over the full latent (attn_head_dim * 2 bytes), + # not the fp8(nope) + bf16(rope) + scales 584-byte cell. + # * SWA is a fixed per-request ring (num_req_slots * ring_size), + # independent of full_token, so it is a fixed *bias* rather than a + # per-token term. Gate on the same switch the pool itself uses so the + # sizing and the allocation never drift apart. + from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( + is_unified_kv_triton, + ) + + self._unified = is_unified_kv_triton() + self.attn_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + # Mirror DeepSeekV4TokenToKVPool: swa_ring_size = sliding_window + + # (speculative_num_draft_tokens - 1). + spec_num_draft = get_spec().speculative_num_draft_tokens or 1 + self._swa_ring_size = self.swa_page_size + ( + (spec_num_draft - 1) if self.is_speculative else 0 + ) + self._spec_infl = 1.0 + if self.is_speculative: # Ring is sized once here, so it must serve the largest adaptive tier. self._assert_ring_serves_draft_tokens( @@ -918,7 +940,8 @@ def __init__(self, kvc: KVCacheConfigurator): # bytes_per_full_token: tokens = avail / (bpft * (T+D)/T). draft_layers = 1 target_layers = self.num_layers_total - self.bytes_per_full_token *= (target_layers + draft_layers) / target_layers + self._spec_infl = (target_layers + draft_layers) / target_layers + self.bytes_per_full_token *= self._spec_infl # Online c128 keeps a single in-progress (max, sum, kv) state per index # and assumes a strict forward-only schedule. Speculative decode (MTP) @@ -971,7 +994,11 @@ def _assert_ring_serves_draft_tokens(self, num_draft_tokens: int) -> None: ) def _get_bytes_per_full_token(self) -> float: - kv_bytes = self.qk_nope_head_dim + self.qk_rope_head_dim * 2 + 8 + if self._unified: + # Unified_kv stores the whole latent in bf16. + kv_bytes = self.attn_head_dim * 2 + else: + kv_bytes = self.qk_nope_head_dim + self.qk_rope_head_dim * 2 + 8 attn_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim c4_state_dtype_size, c128_state_dtype_size = ( @@ -995,16 +1022,40 @@ def _get_bytes_per_full_token(self) -> float: c4_frac = 1 / (4 * self.c4_shrink_factor) return ( - self.swa_ratio * kv_bytes * self.num_layers_total + # Unified_kv: SWA is a fixed per-request ring (see _fixed_swa_bytes), + # not a per-token pool, so it is excluded from the per-token coeff. + ( + 0.0 + if self._unified + else self.swa_ratio * kv_bytes * self.num_layers_total + ) + c4_frac * kv_bytes * self.num_layers_ca4 + 1 / 128 * kv_bytes * self.num_layers_ca128 + 1 / 4 * self.indexer_bytes_per_token * self.num_layers_ca4 - + self.swa_ratio * c4_state_ratio * c4_state_bytes * self.num_layers_ca4 + # Unified_kv: the c4 (attn + indexer) compress-state is a ring buffer + # addressed off the SWA slot ((swa_loc // swa_page_size) * ring_size), + # and the unified SWA pool is a fixed per-request ring + # (swa_pages = num_req_slots * swa_ring_size), so the state ring is + # request-scoped, not full_token-scoped. It is therefore a fixed bias + # (see _fixed_c4_state_bytes), not a per-token term. On the non-unified + # path the SWA pool scales with full_token, so it stays per-token. + + ( + 0.0 + if self._unified + else self.swa_ratio + * c4_state_ratio + * c4_state_bytes + * self.num_layers_ca4 + ) + c128_state_ratio * c128_state_bytes * self.num_layers_ca128 - + self.swa_ratio - * c4_state_ratio - * c4_indexer_state_bytes - * self.num_layers_ca4 + + ( + 0.0 + if self._unified + else self.swa_ratio + * c4_state_ratio + * c4_indexer_state_bytes + * self.num_layers_ca4 + ) ) def _compute_dsv4_sizes(self, full_token: int, page_size: int) -> _DSV4PoolSizes: @@ -1016,7 +1067,14 @@ def _compute_dsv4_sizes(self, full_token: int, page_size: int) -> _DSV4PoolSizes swa_max_total_num_tokens=swa_tokens, c4_max_total_num_tokens=full_token // (4 * self.c4_shrink_factor), c128_max_total_num_tokens=full_token // 128, - c4_state_pool_size=swa_tokens // self.swa_page_size * self.c4_ring_size, + # Unified_kv sizes the c4 state ring from the fixed SWA ring + # (request-scoped), finalized once max_running_requests is known -- so + # it must not scale with full_token here (mirrors c128_state below). + c4_state_pool_size=( + 0 + if self._unified + else swa_tokens // self.swa_page_size * self.c4_ring_size + ), c128_state_pool_size=0, ) @@ -1047,18 +1105,64 @@ def _get_c128_state_fixed_bytes(self, max_running_requests: int) -> int: state_rows * state_last_dim * c128_state_dtype_size * self.num_layers_ca128 ) - def _get_c128_state_fixed_bytes_for_token_capacity( - self, token_capacity: int - ) -> int: + def _unified_c4_state_pool_size(self, max_running_requests: int) -> int: + """Exact request-scoped C4 ring size for the unified address contract. + + Unified C4 state locations are + ``req_pool_idx * c4_ring_size + position % c4_ring_size``. + """ + num_req_slots = self._get_num_req_slots(max_running_requests) + return num_req_slots * self.c4_ring_size + + def _fixed_c4_state_bytes(self, max_running_requests: int) -> int: + """Unified_kv c4 (attn + indexer) compress-state is a fixed per-request + ring, sized by concurrency rather than full_token. Return its byte + footprint across all c4 layers. Returns 0 on the non-unified path (where + the c4 state pool scales with the SWA pool and is accounted per-token).""" + if not self._unified or self.num_layers_ca4 == 0: + return 0 + + c4_state_dtype_size, _ = _get_dsv4_compress_state_dtype_sizes() + attn_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + # CompressStatePool allocates `size + ring_size + 1` rows, padded to the + # compress ratio (see CompressStatePool.__init__). Mirror that here so the + # reserved bias covers the real allocation. + state_rows = self._unified_c4_state_pool_size(max_running_requests) + state_rows = ceil_div(state_rows + self.c4_ring_size + 1, 4) * 4 + # overlap c4: last_dim = 2 * (1 + overlap) * head_dim = 4 * head_dim. + core_bytes = 4 * attn_head_dim * c4_state_dtype_size + indexer_bytes = 4 * self.indexer_head_dim * c4_state_dtype_size + return state_rows * (core_bytes + indexer_bytes) * self.num_layers_ca4 + + def _resolve_max_running_requests_per_worker(self, available_bytes: int) -> int: + """Approximate ModelRunner._resolve_max_num_reqs closely enough to size + the request-scoped fixed pools (c128 state, unified SWA ring). Over- + estimating is safe: a larger fixed bias yields a smaller full_token.""" if self.requested_max_running_requests_per_worker is not None: - return self._get_c128_state_fixed_bytes( - self.requested_max_running_requests_per_worker - ) + return self.requested_max_running_requests_per_worker - estimated = int(token_capacity / self.context_len * 512) + full_token = int(available_bytes / self.bytes_per_full_token) + estimated = int(full_token / self.context_len * 512) estimated = max(min(estimated, 4096), 2048) - max_running_requests = min(estimated, token_capacity // 2) - return self._get_c128_state_fixed_bytes(max_running_requests) + return min(estimated, full_token // 2) + + def _fixed_swa_bytes(self, max_running_requests: int) -> int: + """Unified_kv SWA is a fixed per-request ring, sized by concurrency + (num_req_slots) rather than by full_token. Return its bf16 byte + footprint across all full layers, inflated for the draft worker the same + way as the per-token coeff. Returns 0 on the non-unified path (where SWA + is already accounted per-token).""" + if not self._unified: + return 0 + num_req_slots = self._get_num_req_slots(max_running_requests) + ring_bytes = ( + num_req_slots + * self._swa_ring_size + * self.attn_head_dim + * 2 # bf16 + * self.num_layers_total + ) + return int(ring_bytes * self._spec_infl) def _to_config(self, sizes: _DSV4PoolSizes) -> MemoryPoolConfig: full = sizes.full_max_total_num_tokens @@ -1089,6 +1193,13 @@ def finalize_with_max_running_requests( config.c128_state_pool_size = num_req_slots else: config.c128_state_pool_size = num_req_slots * self.c128_ring_size + # Unified_kv: the c4 state ring is request-scoped (fixed SWA pool), so + # finalize it here from the now-known concurrency. On the non-unified path + # it was already sized from full_token in _compute_dsv4_sizes. + if self._unified and self.num_layers_ca4 > 0: + config.c4_state_pool_size = self._unified_c4_state_pool_size( + config.max_running_requests + ) return config def calculate_pool_sizes( @@ -1098,25 +1209,34 @@ def calculate_pool_sizes( "page_size must be multiple of 128 for compressed attention" ) - if self.requested_max_running_requests_per_worker is not None: - c128_state_fixed_bytes = self._get_c128_state_fixed_bytes( - self.requested_max_running_requests_per_worker - ) - else: - full_token = int(available_bytes / self.bytes_per_full_token) - c128_state_fixed_bytes = ( - self._get_c128_state_fixed_bytes_for_token_capacity(full_token) - ) + max_running_requests_per_worker = self._resolve_max_running_requests_per_worker( + available_bytes + ) + c128_state_fixed_bytes = self._get_c128_state_fixed_bytes( + max_running_requests_per_worker + ) + swa_ring_fixed_bytes = self._fixed_swa_bytes(max_running_requests_per_worker) + c4_state_fixed_bytes = self._fixed_c4_state_bytes( + max_running_requests_per_worker + ) - available_bytes_for_tokens = max(available_bytes - c128_state_fixed_bytes, 0) + available_bytes_for_tokens = max( + available_bytes + - c128_state_fixed_bytes + - swa_ring_fixed_bytes + - c4_state_fixed_bytes, + 0, + ) full_token = int(available_bytes_for_tokens / self.bytes_per_full_token) sizes = self._compute_dsv4_sizes(full_token, page_size) logger.info( - f"DSV4 memory calculation: " + f"DSV4 memory calculation: unified={self._unified}, " f"bytes_per_full_token={self.bytes_per_full_token:.2f}, " f"available_bytes={available_bytes / (1 << 30):.2f} GB, " f"c128_state_fixed={c128_state_fixed_bytes / (1 << 30):.2f} GB, " + f"swa_ring_fixed={swa_ring_fixed_bytes / (1 << 30):.2f} GB, " + f"c4_state_fixed={c4_state_fixed_bytes / (1 << 30):.2f} GB, " f"full_token={sizes.full_max_total_num_tokens}" ) return self._to_config(sizes) diff --git a/test/registered/kernels/ops/attention/test_c4_v2.py b/test/registered/kernels/ops/attention/test_c4_v2.py index f05aaa9039ad..cd419ef76db6 100644 --- a/test/registered/kernels/ops/attention/test_c4_v2.py +++ b/test/registered/kernels/ops/attention/test_c4_v2.py @@ -7,7 +7,11 @@ import torch import triton -from sglang.kernels.ops.attention.dsv4 import compress_forward +from sglang.kernels.ops.attention.dsv4 import ( + CompressorDecodePlan, + CompressorPrefillPlan, + compress_forward, +) from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.kernels.deepseek_v4.common import ( @@ -122,6 +126,91 @@ def _make_inputs( # ----------------------------------------------------------------------------- +@pytest.mark.parametrize("ring_size", [8, 16]) +@pytest.mark.parametrize( + ("gpu_inputs", "use_cuda_graph"), + [(False, False), (True, False), (True, True)], +) +def test_unified_request_ring_plans_ignore_full_to_state( + ring_size: int, gpu_inputs: bool, use_cuda_graph: bool +) -> None: + """C4 plans must address state by request slot, not the full-cache map.""" + device = torch.device(get_device()) + req_pool_indices = torch.tensor([2, 5], dtype=torch.int64, device=device) + req_to_token = torch.zeros((6, 16), dtype=torch.int32, device=device) + full_to_state = torch.zeros(1, dtype=torch.int64, device=device) + seq_lens = torch.tensor([8, 12], dtype=torch.int64) + extend_lens = torch.tensor([4, 4], dtype=torch.int64) + if gpu_inputs: + seq_lens = seq_lens.to(device) + extend_lens = extend_lens.to(device) + + prefill = CompressorPrefillPlan.generate( + compress_ratio=RATIO, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + extend_lens=extend_lens, + req_to_token=req_to_token, + full_to_state=full_to_state, + swa_page_size=256, + ring_size=ring_size, + num_q_tokens=8, + use_cuda_graph=use_cuda_graph, + use_req_ring=True, + ) + plan_c = prefill.plan_c.view(torch.int32).reshape(-1, 4).cpu() + plan_w = prefill.plan_w.view(torch.int32).reshape(-1, 2).cpu() + + valid_c = plan_c[plan_c[:, 2] >= 0] + got_reads = { + int(row[1].item()) & 0xFFFF: (int(row[2].item()), int(row[3].item())) + for row in valid_c + } + expected_reads = { + 3: ( + (2 * ring_size + 3 % ring_size) // RATIO, + (2 * ring_size + 7 % ring_size) // RATIO, + ), + 7: ( + (5 * ring_size + 7 % ring_size) // RATIO, + (5 * ring_size + 11 % ring_size) // RATIO, + ), + } + assert got_reads == expected_reads + + valid_w = plan_w[plan_w[:, 1] >= 0] + got_writes = {int(row[0].item()): int(row[1].item()) for row in valid_w} + expected_writes = { + **{j: 2 * ring_size + (4 + j) % ring_size for j in range(4)}, + **{4 + j: 5 * ring_size + (8 + j) % ring_size for j in range(4)}, + } + assert got_writes == expected_writes + assert {got_writes[j] for j in range(4)}.isdisjoint( + {got_writes[j] for j in range(4, 8)} + ) + + decode = CompressorDecodePlan.generate( + compress_ratio=RATIO, + req_pool_indices=req_pool_indices, + req_to_token=req_to_token, + full_to_state=full_to_state, + seq_lens=torch.tensor([8, 12], dtype=torch.int64, device=device), + swa_page_size=256, + ring_size=ring_size, + use_req_ring=True, + ) + got_decode = decode.plan_d.view(torch.int32).reshape(-1, 4).cpu() + expected_decode = torch.tensor( + [ + [8, 2 * ring_size + 7 % ring_size, *expected_reads[3]], + [12, 5 * ring_size + 11 % ring_size, *expected_reads[7]], + ], + dtype=torch.int32, + ) + assert torch.equal(got_decode, expected_decode) + assert got_decode[0, 1] != got_decode[1, 1] + + @pytest.mark.parametrize("mode", ["legacy", "paged"]) @pytest.mark.parametrize("seq_len", [4, 8, 32, 256, 1024]) def test_prefill_no_context(mode: str, seq_len: int) -> None: diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index d74db97d6cbb..bd60d6a68f51 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -9,10 +9,7 @@ PrefillAdder, estimate_prefill_extend_tile_metrics, ) -from sglang.srt.mem_cache.base_prefix_cache import ( - DecLockRefResult, - IncLockRefResult, -) +from sglang.srt.mem_cache.base_prefix_cache import DecLockRefResult, IncLockRefResult from sglang.srt.runtime_context import get_context from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler from sglang.srt.utils.common import Range @@ -71,6 +68,13 @@ def create_token_allocator( allocator.swa_available_size.return_value = swa_available_size allocator.available_size.return_value = available_size allocator.size_swa = size_swa + # get_kvcache().[_unified_kv] gates the unified-KV SWA-ring accounting + # path in schedule_policy.add_chunked_req / rem_swa_tokens. A bare + # MagicMock auto-creates any attribute access as a truthy Mock, so + # without this the getattr(..., "_unified_kv", False) default never + # triggers and these tests silently exercise the unified-KV branch + # instead of the standard hybrid-SWA one they intend to cover. + allocator.get_kvcache.return_value._unified_kv = False return allocator def create_running_batch(self, reqs=None) -> MagicMock: diff --git a/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py b/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py index 6152e59cc4d4..54440519b62c 100644 --- a/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py +++ b/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py @@ -23,6 +23,9 @@ def __init__(self, base=1000, page_size=1): self.alloc_calls = [] self.extend_calls = [] + def get_kvcache(self): + return None + def available_size(self): return 1 << 30 diff --git a/test/registered/unit/mem_cache/test_dsv4_c4_state_lifecycle.py b/test/registered/unit/mem_cache/test_dsv4_c4_state_lifecycle.py new file mode 100644 index 000000000000..24ec0e83dd2c --- /dev/null +++ b/test/registered/unit/mem_cache/test_dsv4_c4_state_lifecycle.py @@ -0,0 +1,119 @@ +"""CPU/mock tests for unified DSV4 C4 request-state lifecycle.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock + +import torch + +from sglang.srt.mem_cache.allocation import alloc_req_slots +from sglang.srt.mem_cache.deepseek_v4_compress_state import KVAndScore +from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool +from sglang.srt.mem_cache.memory_pool import ReqToTokenPool +from sglang.srt.model_executor.pool_configurator import DSV4PoolConfigurator +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _request(req_pool_idx=None, *, reused=False): + return SimpleNamespace( + kv=SimpleNamespace( + req_pool_idx=req_pool_idx, + kv_committed_len=1 if reused else 0, + kv_allocated_len=1 if reused else 0, + holds_kv=reused, + ), + inflight_middle_chunks=1 if reused else 0, + ) + + +def _c4_pool(rows: int, width: int, ring_size: int): + return SimpleNamespace( + ratio=4, + ring_size=ring_size, + kv_score_buffer=KVAndScore(torch.full((rows, width), 7.0)), + ) + + +class TestUnifiedC4StateLifecycle(unittest.TestCase): + def test_pool_size_is_exact_request_ring_product(self): + configurator = object.__new__(DSV4PoolConfigurator) + configurator.disaggregation_mode = "decode" + configurator.disaggregation_decode_extra_slots = 3 + configurator.c4_ring_size = 16 + + self.assertEqual(configurator._unified_c4_state_pool_size(10), 14 * 16) + + def test_clear_resets_only_selected_request_rings(self): + ring_size = 8 + logical_rows = 4 * ring_size + physical_rows = logical_rows + ring_size + 4 + attn = _c4_pool(physical_rows, width=12, ring_size=ring_size) + indexer = _c4_pool(physical_rows, width=8, ring_size=ring_size) + c128 = SimpleNamespace( + ratio=128, + ring_size=128, + kv_score_buffer=KVAndScore(torch.full((physical_rows, 8), 9.0)), + ) + + token_pool = object.__new__(DeepSeekV4TokenToKVPool) + token_pool._unified_kv = True + token_pool.compress_state_pools = [attn, c128] + token_pool.indexer_compress_state_pools = [indexer, None] + token_pool.get_ring_size = MagicMock(return_value=ring_size) + + token_pool.clear_c4_req_states([1, 3]) + + selected = torch.tensor(list(range(8, 16)) + list(range(24, 32))) + untouched = torch.tensor(list(range(0, 8)) + list(range(16, 24))) + for pool in (attn, indexer): + state = pool.kv_score_buffer.kv_score + half = state.shape[-1] // 2 + self.assertTrue( + torch.equal( + state[selected, :half], torch.zeros_like(state[selected, :half]) + ) + ) + self.assertTrue(torch.isneginf(state[selected, half:]).all()) + self.assertTrue((state[untouched] == 7).all()) + self.assertTrue((state[logical_rows:] == 7).all()) + self.assertTrue((c128.kv_score_buffer.kv_score == 9).all()) + + def test_alloc_clears_new_slots_but_not_reused_slots(self): + req_pool = ReqToTokenPool(3, 16, "cpu", enable_memory_saver=False) + token_pool = MagicMock() + reused = _request() + + # First admission: a brand-new slot, so its C4 ring must be cleared. + (reused_idx,) = alloc_req_slots( + req_pool, [reused], None, token_to_kv_pool=token_pool + ) + token_pool.clear_c4_req_states.assert_called_once_with([reused_idx]) + + # Chunked continuation reuses the same slot -- clearing it here would + # wipe the state captured by the previous chunk. + token_pool.clear_c4_req_states.reset_mock() + reused.kv.req_pool_idx = reused_idx + reused.kv.kv_committed_len = 1 + reused.kv.kv_allocated_len = 1 + reused.kv.holds_kv = True + reused.inflight_middle_chunks = 1 + self.assertEqual( + alloc_req_slots(req_pool, [reused], None, token_to_kv_pool=token_pool), + [reused_idx], + ) + token_pool.clear_c4_req_states.assert_not_called() + + # Mixed batch: only the newly allocated slot is cleared. + fresh = _request() + indices = alloc_req_slots( + req_pool, [reused, fresh], None, token_to_kv_pool=token_pool + ) + self.assertEqual(indices[0], reused_idx) + self.assertNotEqual(indices[1], reused_idx) + token_pool.clear_c4_req_states.assert_called_once_with([indices[1]]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_hisparse_allocator.py b/test/registered/unit/mem_cache/test_hisparse_allocator.py index 7ea10c1f4f0b..ebc173a4ce59 100644 --- a/test/registered/unit/mem_cache/test_hisparse_allocator.py +++ b/test/registered/unit/mem_cache/test_hisparse_allocator.py @@ -131,6 +131,7 @@ def write(self, indices, values): queue = DecodePreallocQueue.__new__(DecodePreallocQueue) queue.req_to_token_pool = req_to_token_pool queue.token_to_kv_pool_allocator = allocator + queue.token_to_kv_pool = None queue.tree_cache = SimpleNamespace( evictable_size=MagicMock(return_value=0), protected_size=MagicMock(return_value=0), diff --git a/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py b/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py index 7cef341de813..e3de143f2e21 100644 --- a/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py +++ b/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py @@ -36,6 +36,11 @@ def set_full_to_swa_mapping( return SimpleNamespace( page_size=page_size, + # alloc_extend branches on self._unified to skip the vestigial paged SWA + # allocator on the unified-KV path. This stub exercises the standard + # hybrid-SWA path, so pin it False rather than letting the attribute go + # missing (SimpleNamespace raises instead of defaulting). + _unified=False, full_attn_allocator=SimpleNamespace( available_size=lambda: full_available, alloc_extend=MagicMock(return_value=full_indices), From 7ce72c5aa8e9965444b2161d2aed296328a23212 Mon Sep 17 00:00:00 2001 From: yuttian1 Date: Sun, 6 Sep 2026 04:39:05 +0000 Subject: [PATCH 02/15] [AMD][DSV4] Gate every shared-path change behind the unified-KV switch The reverted #30315 claimed "no behavior change on the non-unified (fp8) path", but three changes reached shared code with no gate at all: 1. deepseek_v4_memory_pool.py: `c4_state_pool_size = max(caller, num_req_slots * ring)` ran *before* `_unified_kv` was resolved, so the fp8 path also paid the request-addressed floor -- raising its footprint and lowering max_total_num_tokens. The gate is now resolved before any sizing, and unified sets the size exactly instead of via max(); fp8 keeps the caller-supplied SWA-addressed size byte-for-byte. 2. schedule_policy.py: the `swa_needed >= rem_swa_tokens` -> `>` relaxation at both admission sites applied to every hybrid-SWA model. It is correct only where rem_swa_tokens is an exact ring-slot capacity, so it is now conditional; the legacy SWA-token path keeps its conservative `>=`. 3. invariant_checker.py: the unconditional `logical_attn_allocator` unwrap changed which object the invariant was asserted against on every DSV4-HiSparse deployment, not just unified ones. Also tightens the two duck-typed `clear_c4_req_states` call sites (allocation.py, disaggregation/decode.py) to check `_unified_kv` explicitly rather than relying on the hook's internal early return, and skips the `newly_allocated` list comprehension entirely off the unified path. The gate is identity-compared (`is True`) so a duck-typed test stub that auto-creates attributes cannot accidentally select the unified path. Drops the unused `host_hit_len` / `storage_hit_len` parameters #30315 added to `_update_prefill_budget`; no call site ever passed them. test_dsv4_c4_state_lifecycle.py: the unified lifecycle test built a bare MagicMock and relied on duck-typing alone, so it now declares `_unified_kv` like the pool it stands in for. Adds the matching negative test that the reset never fires off the unified path. --- python/sglang/srt/disaggregation/decode.py | 6 ++- python/sglang/srt/managers/schedule_policy.py | 40 +++++++++++++------ .../scheduler_components/invariant_checker.py | 10 +++-- python/sglang/srt/mem_cache/allocation.py | 25 ++++++++---- .../srt/mem_cache/deepseek_v4_memory_pool.py | 24 +++++------ .../mem_cache/test_dsv4_c4_state_lifecycle.py | 20 ++++++++++ 6 files changed, 88 insertions(+), 37 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 06b4384b5313..a9ae8d8661b0 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1774,7 +1774,11 @@ def _pre_alloc( assert req_pool_indices is not None, ( "req_pool_indices is full! There is a bug in memory estimation." ) - if is_new_req_slot: + if ( + is_new_req_slot + and getattr(self.token_to_kv_pool, "_unified_kv", False) is True + ): + # Unified-KV DSV4 only: reset the stale C4 ring on a fresh slot. clear_c4_req_states = getattr( self.token_to_kv_pool, "clear_c4_req_states", None ) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index a4fbf5e6e231..a32300a79cb5 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -499,6 +499,16 @@ def __init__( self.prefill_tile_block_m = prefill_tile_block_m self.tree_cache = tree_cache self.token_to_kv_pool_allocator = token_to_kv_pool_allocator + # Unified-KV DSV4 sizes SWA as a fixed per-request ring slot, not a + # token budget. `is True`: a duck-typed stub must not select that path. + self._unified_kv = ( + getattr( + getattr(token_to_kv_pool_allocator, "get_kvcache", lambda: None)(), + "_unified_kv", + False, + ) + is True + ) self.running_batch = running_batch self.new_token_ratio = new_token_ratio self.rem_input_tokens = rem_input_tokens - num_mixed_decode_tokens @@ -661,7 +671,7 @@ def rem_total_tokens(self): @property def rem_swa_tokens(self): allocator = self.token_to_kv_pool_allocator - if getattr(allocator.get_kvcache(), "_unified_kv", False): + if self._unified_kv: # Unified-KV: SWA is a per-request ring, not a tree-reusable token # pool. swa_available_size() already reports ring capacity # (free_slots * ring_cost). tree swa_evictable is in the old linear @@ -717,7 +727,7 @@ def _swa_budget_for_req( from double-counting extend, so budget <= extend + max_new_tokens + page. """ allocator = self.token_to_kv_pool_allocator - if getattr(allocator.get_kvcache(), "_unified_kv", False): + if self._unified_kv: # Unified-KV: each request occupies exactly one fixed SWA ring slot, # independent of context / chunk length; a host-hit prefix reuses the # same ring. Budget the fixed per-slot ring cost (paired with the @@ -850,8 +860,6 @@ def _update_prefill_budget( max_new_tokens: int, retracted_stain: bool, mamba_gap_reserve: int = 0, - host_hit_len: int = 0, - storage_hit_len: int = 0, is_chunked_continuation: bool = False, ): # TODO(lsyin): check this workaround logic, which only ensures the prefill will not out of memory, and may be too conservative @@ -880,10 +888,7 @@ def _update_prefill_budget( # first admission and already reflected in swa_available_size() on # later rounds. Charging it again on a chunked continuation would # double-count the slot and over-throttle admission, so skip it. - _unified = getattr( - self.token_to_kv_pool_allocator.get_kvcache(), "_unified_kv", False - ) - if not (_unified and is_chunked_continuation): + if not (self._unified_kv and is_chunked_continuation): self.rem_swa_token_offset += self._swa_budget_for_req( extend_input_len, max_new_tokens ) @@ -1021,9 +1026,7 @@ def add_chunked_req(self, req: Req): _rem_tokens = self._get_dllm_remain_tokens() else: _rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens)) - if self.is_hybrid_swa and not getattr( - self.token_to_kv_pool_allocator.get_kvcache(), "_unified_kv", False - ): + if self.is_hybrid_swa and not self._unified_kv: # alloc_extend needs extend_num_tokens + page_size per request, # so reserve one page here to avoid OOM. # Unified-KV: rem_swa_tokens is ring capacity (free_slots * ring @@ -1278,7 +1281,14 @@ def add_one_req( self._swa_new_tokens(req), swa_host_hit_length=req.swa_host_hit_length, ) - if swa_needed > self.rem_swa_tokens: + # Unified-KV: rem_swa_tokens is an exact ring-slot capacity, so a + # request needing exactly what is left still fits. The legacy + # SWA-token path keeps its original conservative `>=`. + if ( + swa_needed > self.rem_swa_tokens + if self._unified_kv + else swa_needed >= self.rem_swa_tokens + ): if not self._swa_req_never_fits( real_input_tokens, self._swa_new_tokens(req), @@ -1314,7 +1324,11 @@ def add_one_req( self._swa_new_tokens(req), swa_host_hit_length=req.swa_host_hit_length, ) - if swa_needed > self.rem_swa_tokens: + if ( + swa_needed > self.rem_swa_tokens + if self._unified_kv + else swa_needed >= self.rem_swa_tokens + ): if not self._swa_req_never_fits( real_input_tokens, self._swa_new_tokens(req), diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index fc861c589077..ec33656725c9 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -394,10 +394,14 @@ def _add_owner(req_or_slot, label, rpi, committed, allocated): # Sub-allocators to check: a flat allocator is its own single sub; a # hybrid-SWA wrapper exposes full_attn_allocator + swa_attn_allocator. - # DSV4-HiSparse nests the real SWA allocator under logical_attn_allocator, - # so unwrap first (no-op for a plain/flat allocator). alloc = self.token_to_kv_pool_allocator - alloc = getattr(alloc, "logical_attn_allocator", alloc) + # Unified-KV DSV4-HiSparse nests the real SWA allocator one level + # down; elsewhere the wrapper is the object this invariant asserts on. + if ( + getattr(getattr(alloc, "get_kvcache", lambda: None)(), "_unified_kv", False) + is True + ): + alloc = getattr(alloc, "logical_attn_allocator", alloc) sub_allocs = ( [alloc] if getattr(alloc, "free_pages", None) is not None diff --git a/python/sglang/srt/mem_cache/allocation.py b/python/sglang/srt/mem_cache/allocation.py index 4cf251db0f8d..78affd3f3654 100644 --- a/python/sglang/srt/mem_cache/allocation.py +++ b/python/sglang/srt/mem_cache/allocation.py @@ -261,7 +261,18 @@ def alloc_req_slots( tree_cache.evict_for_alloc( EvictParams(num_tokens=0, mamba_num=mamba_num) ) - newly_allocated = [req.kv.req_pool_idx is None for req in reqs] + # Unified-KV DSV4 only: a freshly allocated req slot carries stale C4 ring + # state. Resolve the hook before alloc() overwrites req_pool_idx. + clear_c4_req_states = ( + getattr(token_to_kv_pool, "clear_c4_req_states", None) + if getattr(token_to_kv_pool, "_unified_kv", False) is True + else None + ) + newly_allocated = ( + [req.kv.req_pool_idx is None for req in reqs] + if clear_c4_req_states is not None + else None + ) req_pool_indices = req_to_token_pool.alloc(reqs) if req_pool_indices is None: raise RuntimeError( @@ -270,12 +281,12 @@ def alloc_req_slots( f"{req_to_token_pool.available_size()=}, {num_reqs=}, " ) - new_req_pool_indices = [ - idx for idx, is_new in zip(req_pool_indices, newly_allocated) if is_new - ] - clear_c4_req_states = getattr(token_to_kv_pool, "clear_c4_req_states", None) - if new_req_pool_indices and clear_c4_req_states is not None: - clear_c4_req_states(new_req_pool_indices) + if clear_c4_req_states is not None: + new_req_pool_indices = [ + idx for idx, is_new in zip(req_pool_indices, newly_allocated) if is_new + ] + if new_req_pool_indices: + clear_c4_req_states(new_req_pool_indices) return req_pool_indices diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index 991a8a5fb13b..ae4d00b16b32 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -564,11 +564,18 @@ def __init__( self.c4_size = c4_size self.c4_logical_size = c4_logical_size self.c128_size = c128_size - # Keep the legacy SWA-addressed pool large enough on non-unified paths. - # Unified request-addressed sizing is set exactly after resolving the - # unified-kv gate below. + from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( + is_unified_kv_triton, + ) + + # Resolve the unified-kv gate before any sizing so the two cannot drift. + self._unified_kv = is_unified_kv_triton() c4_ring_size = self.get_ring_size(4) - c4_state_pool_size = max(c4_state_pool_size, self.num_req_slots * c4_ring_size) + if self._unified_kv: + # Unified C4 state is request-addressed: one ring per req slot, + # so the caller-supplied, SWA-scaled size does not apply here. + c4_state_pool_size = self.num_req_slots * c4_ring_size + # Non-unified (fp8) keeps the caller-supplied, SWA-addressed size. self.c4_state_pool_size = c4_state_pool_size c128_ring_size = self.get_ring_size(128) if ONLINE_C128: @@ -624,15 +631,6 @@ def __init__( c4_page_size = page_size // 4 c128_page_size = page_size // 128 - from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( - is_unified_kv_triton, - ) - - self._unified_kv = is_unified_kv_triton() - if self._unified_kv: - # Unified C4 state is request-scoped: no SWA-derived over-allocation. - self.c4_state_pool_size = self.num_req_slots * c4_ring_size - if self._unified_kv: self.swa_kv_pool = None self.c4_kv_pool = None diff --git a/test/registered/unit/mem_cache/test_dsv4_c4_state_lifecycle.py b/test/registered/unit/mem_cache/test_dsv4_c4_state_lifecycle.py index 24ec0e83dd2c..11102b4e5378 100644 --- a/test/registered/unit/mem_cache/test_dsv4_c4_state_lifecycle.py +++ b/test/registered/unit/mem_cache/test_dsv4_c4_state_lifecycle.py @@ -83,6 +83,9 @@ def test_clear_resets_only_selected_request_rings(self): def test_alloc_clears_new_slots_but_not_reused_slots(self): req_pool = ReqToTokenPool(3, 16, "cpu", enable_memory_saver=False) token_pool = MagicMock() + # A bare MagicMock auto-creates a truthy `_unified_kv`, so the stub + # must declare it; the gate identity-compares against True. + token_pool._unified_kv = True reused = _request() # First admission: a brand-new slot, so its C4 ring must be cleared. @@ -114,6 +117,23 @@ def test_alloc_clears_new_slots_but_not_reused_slots(self): self.assertNotEqual(indices[1], reused_idx) token_pool.clear_c4_req_states.assert_called_once_with([indices[1]]) + def test_alloc_does_not_clear_c4_state_off_the_unified_path(self): + """The reset must not reach the non-unified (fp8) path. + + A duck-typed `hasattr(pool, "clear_c4_req_states")` check is not enough: + the attribute exists on every DeepSeekV4TokenToKVPool, unified or not. + """ + for unified in (False, None, 1, "yes"): + with self.subTest(unified_kv=unified): + # Fresh pool per subtest: each iteration consumes a req slot. + req_pool = ReqToTokenPool(1, 16, "cpu", enable_memory_saver=False) + token_pool = MagicMock() + token_pool._unified_kv = unified + alloc_req_slots( + req_pool, [_request()], None, token_to_kv_pool=token_pool + ) + token_pool.clear_c4_req_states.assert_not_called() + if __name__ == "__main__": unittest.main() From 4e7916b8e43de570e89b25c343bfaf26bff0927a Mon Sep 17 00:00:00 2001 From: yuttian1 Date: Sun, 6 Sep 2026 07:10:19 +0000 Subject: [PATCH 03/15] [AMD][DSV4] Drop incidental churn from the reland The reverted PR carried edits that change no behavior and only widen the review surface. Removing them shrinks the diff from 113 hunks / +860/-166 to 92 hunks / +815/-97 with no functional delta. - Restore 14 import blocks that were collapsed to one line. The repo's own ruff-format (v0.15.1, magic trailing comma) keeps the multi-line form, so the collapse did not come from repo tooling. - Collapse two byte-identical Triton branches. `elif use_req_ring:` had the same body as `if compress_ratio == 128:` in attn.py and compressor.py; both operands are tl.constexpr, so `or` folds at compile time. This is already the form compress_hip.py uses in the same series. - Drop a local `attn_head_dim` in _fixed_c4_state_bytes that duplicates the `self.attn_head_dim` the series itself introduces. - Restore seven comments in c_plan.cuh that were re-wrapped to ~80 columns. The governing .clang-format sets ColumnLimit: 120 and clang-format 20.1.7 reports both the old and the new text as clean, so the re-wrap was gratuitous; two of them broke a column-aligned address table. --- .../kernels/jit/csrc/deepseek_v4/c_plan.cuh | 41 +++++++------------ .../sglang/kernels/ops/attention/dsv4/attn.py | 4 +- .../srt/layers/attention/dsv4/compressor.py | 8 ++-- python/sglang/srt/managers/schedule_batch.py | 5 ++- python/sglang/srt/managers/schedule_policy.py | 5 ++- .../scheduler_components/invariant_checker.py | 14 ++++++- .../pool_stats_observer.py | 9 +++- .../srt/mem_cache/kv_cache_configurator.py | 24 ++++++++--- .../srt/model_executor/pool_configurator.py | 7 ++-- .../unit/managers/test_prefill_adder.py | 5 ++- 10 files changed, 74 insertions(+), 48 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh index e341d4ba4426..5e23da74dfb8 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh @@ -28,14 +28,12 @@ using R2T_T = int32_t; using F2S_T = int64_t; using IDX_T = int64_t; -/// NOTE: for the internal use, we pack the ragged and batch id, since both not -/// exceed 65536 +/// NOTE: for the internal use, we pack the ragged and batch id, since both not exceed 65536 SGL_DEVICE __host__ PlanW pack_w(uint32_t ragged_id, uint32_t batch_id, int32_t seq_len) { return {static_cast(ragged_id | batch_id << 16), seq_len}; } -/// NOTE: for the internal use, we pack the ragged and batch id, since both not -/// exceed 65536 +/// NOTE: for the internal use, we pack the ragged and batch id, since both not exceed 65536 SGL_DEVICE uint2 unpack_w(PlanW plan) { return {static_cast(plan.ragged_id), static_cast(plan.ragged_id >> 16)}; } @@ -161,8 +159,7 @@ __global__ __launch_bounds__(1024, 1) // counter_w = 0; } // === Stage B: min/max(extend_len) for MTP-uniform detection === - // For min, treat threads outside `batch_size` as +inf so they don't pull the - // min down. + // For min, treat threads outside `batch_size` as +inf so they don't pull the min down. const uint32_t e_for_max = static_cast(extend_len); const uint32_t e_for_min = (tx < params.batch_size) ? e_for_max : 0xFFFFFFFFu; warp_max[warp_id] = warp::reduce_max(e_for_max); @@ -175,19 +172,17 @@ __global__ __launch_bounds__(1024, 1) // __syncthreads(); const auto num_q = params.num_q_tokens; - // MTP-uniform: every batch shares the same small extend_len `E`, so we can - // decompose a global token id `k` into (batch_id, j) = (k / E, k % E) and - // skip the per-batch loop. + // MTP-uniform: every batch shares the same small extend_len `E`, so we can decompose + // a global token id `k` into (batch_id, j) = (k / E, k % E) and skip the per-batch loop. const bool is_mtp_extend = (s_min_extend == s_max_extend) && (s_max_extend > 0) && (s_max_extend <= 32); // === Stage C: emit valid plans, slot allocation via shared-mem atomicAdd === if (is_mtp_extend) { - // Path 1: token-driven. Each global token id maps to exactly one (batch_id, - // j). + // Path 1: token-driven. Each global token id maps to exactly one (batch_id, j). const uint32_t E = s_max_extend; - // num_q is the padded buffer size (graph bucket), not the work size: cap - // the loop at the real token count so batch_id = k / E stays < batch_size - // on an underfilled replay; Stage D pads [counter, num_q) with invalid. + // num_q is the padded buffer size (graph bucket), not the work size: cap the + // loop at the real token count so batch_id = k / E stays < batch_size on an + // underfilled replay; Stage D pads [counter, num_q) with invalid. const uint32_t num_real_q = params.batch_size * E; for (uint32_t k = tx; k < num_real_q; k += block_size) { const uint32_t batch_id = k / E; @@ -375,10 +370,8 @@ __global__ void plan_compress_prefill_legacy_kernel(const Prefill1ParamsLegacy p auto plan_w = idx < params.num_w ? params.plan_w[idx] : PlanW::invalid(); /// Per-request ring buffer slot translation: - /// - c4: page = rid * 2 + (position / 4) % 2; slot = page * 4 + position % - /// 4 - /// - c128: page = rid; slot = rid * 128 + position % - /// 128 + /// - c4: page = rid * 2 + (position / 4) % 2; slot = page * 4 + position % 4 + /// - c128: page = rid; slot = rid * 128 + position % 128 const auto legacy_compute_page = [&](int32_t rid, int32_t position) { if (params.compress_ratio == 4) return rid * 2 + ((position / 4) & 1); return rid; // c128 @@ -404,8 +397,7 @@ __global__ void plan_compress_prefill_legacy_kernel(const Prefill1ParamsLegacy p if (!plan_w.is_invalid()) { const auto [ragged_id, batch_id] = unpack_w(plan_w); const auto rid = static_cast(params.rid_ptr[batch_id]); - // `write_loc` carries (position + 1) at this stage; may not be - // ratio-aligned + // `write_loc` carries (position + 1) at this stage; may not be ratio-aligned const auto position = static_cast(plan_w.write_loc) - 1; plan_w.ragged_id = ragged_id; plan_w.write_loc = legacy_compute_loc(rid, position); @@ -419,10 +411,8 @@ __global__ void plan_compress_decode_legacy_kernel(const DecodeParamsLegacy para const auto idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= params.batch_size) return; /// Per-request ring buffer slot translation: - /// - c4: page = rid * 2 + (position / 4) % 2; slot = page * 4 + position % - /// 4 - /// - c128: page = rid; slot = rid * 128 + position % - /// 128 + /// - c4: page = rid * 2 + (position / 4) % 2; slot = page * 4 + position % 4 + /// - c128: page = rid; slot = rid * 128 + position % 128 const auto legacy_compute_page = [&](int32_t rid, int32_t position) { if (params.compress_ratio == 4) return rid * 2 + ((position / 4) & 1); return rid; // c128 @@ -461,8 +451,7 @@ using PrefillPlan = tvm::ffi::Tuple; * @param compress_plan `[num_q_tokens, 16]` uint8 (output) * @param write_plan `[num_q_tokens, 8]` uint8 (output) * @param compress_ratio 4 for c4, 128 for c128 - * @param use_cuda_graph Whether the plans will be used with cuda graph (affects - * padding) + * @param use_cuda_graph Whether the plans will be used with cuda graph (affects padding) * @return (compress plan tensor, write plan tensor) */ inline PrefillPlan plan_compress_prefill( diff --git a/python/sglang/kernels/ops/attention/dsv4/attn.py b/python/sglang/kernels/ops/attention/dsv4/attn.py index b98dc922d161..8996bb5226f6 100644 --- a/python/sglang/kernels/ops/attention/dsv4/attn.py +++ b/python/sglang/kernels/ops/attention/dsv4/attn.py @@ -134,9 +134,7 @@ def create_paged_compress_data_kernel( else: pos = write_overlap_pos pos = tl.maximum(pos, 0) - if compress_ratio == 128: - state_loc = rid * ring_size + (pos % ring_size) - elif use_req_ring: + if compress_ratio == 128 or use_req_ring: state_loc = rid * ring_size + (pos % ring_size) else: loc = tl.load( diff --git a/python/sglang/srt/layers/attention/dsv4/compressor.py b/python/sglang/srt/layers/attention/dsv4/compressor.py index d5b8c29b867a..2e420f65367a 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor.py @@ -26,7 +26,9 @@ cp_all_gather_rerange_finish, cp_all_gather_rerange_launch, ) -from sglang.srt.mem_cache.deepseek_v4_compress_state import CompressStatePool +from sglang.srt.mem_cache.deepseek_v4_compress_state import ( + CompressStatePool, +) from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.model_executor.forward_context import get_attn_backend from sglang.srt.models.deepseek_v2 import _is_hip @@ -270,9 +272,7 @@ def clip_down(positions: torch.Tensor) -> torch.Tensor: def get_raw_loc(positions: torch.Tensor) -> torch.Tensor: positions = positions.masked_fill(positions < 0, 0) - if compress_ratio == 128: - state_loc = req_pool_indices * ring_size + positions % ring_size - elif use_req_ring: + if compress_ratio == 128 or use_req_ring: state_loc = req_pool_indices * ring_size + positions % ring_size else: loc = req_to_token[req_pool_indices, positions] diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index e539e91e98fd..95db5c369a5a 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -101,7 +101,10 @@ from sglang.srt.managers.scheduler_components.new_token_ratio_tracker import ( NewTokenRatioTracker, ) -from sglang.srt.mem_cache.allocation import alloc_for_decode, alloc_for_extend +from sglang.srt.mem_cache.allocation import ( + alloc_for_decode, + alloc_for_extend, +) from sglang.srt.mem_cache.allocation_sizing import get_alloc_reserve_per_decode from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index a32300a79cb5..7a1bac2e3e62 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -5,7 +5,10 @@ from sglang.srt.environ import envs from sglang.srt.managers.prefill_delayer import PrefillDelayerSinglePassExecutor -from sglang.srt.runtime_context import get_disagg, get_schedule +from sglang.srt.runtime_context import ( + get_disagg, + get_schedule, +) from sglang.srt.utils import get_bool_env_var, is_hip _ROUTING_KEY_POLICY_DEBUG_LOG = get_bool_env_var("SGLANG_ROUTING_KEY_POLICY_DEBUG_LOG") diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index ec33656725c9..5e19c6e5a115 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -3,7 +3,14 @@ import logging from collections import deque from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Callable, Deque, List, Optional, Tuple +from typing import ( + TYPE_CHECKING, + Callable, + Deque, + List, + Optional, + Tuple, +) import torch @@ -25,7 +32,10 @@ scheduler_stage_method, ) from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils.common import ceil_align, raise_error_or_warn +from sglang.srt.utils.common import ( + ceil_align, + raise_error_or_warn, +) from sglang.srt.utils.watchdog import WatchdogRaw if TYPE_CHECKING: diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index e8f9a669750d..18d730ff169a 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -2,7 +2,14 @@ import dataclasses from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple +from typing import ( + TYPE_CHECKING, + Any, + Callable, + List, + Optional, + Tuple, +) from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( UnifiedMambaSWATokenToKVPoolAllocator, diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index a4077d892caf..0143eee5568c 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -204,7 +204,9 @@ def _pp_local_per_request_bytes( from sglang.srt.model_executor.model_runner_components.spec_aux_hidden_state import ( SpecAuxHiddenStateConfig, ) - from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig + from sglang.srt.model_executor.pool_configurator import ( + MemoryPoolConfig, + ) class KVCacheConfigResult(msgspec.Struct, frozen=True, kw_only=True): @@ -1033,7 +1035,9 @@ def _build_hybrid_mamba_decode_req_pool( extra_max_context_len: int, pre_alloc_size: int, ) -> ReqToTokenPool: - from sglang.srt.disaggregation.decode import HybridMambaDecodeReqToTokenPool + from sglang.srt.disaggregation.decode import ( + HybridMambaDecodeReqToTokenPool, + ) req_to_token_pool = HybridMambaDecodeReqToTokenPool( size=max_num_reqs, @@ -1321,7 +1325,9 @@ def _build_dsv4_kv_pool( assert swa_page_size == 256, "In paged swa mode, page_size must be 256." if self.is_draft_worker: - from sglang.srt.models.deepseek_v4_nextn import COMPRESS_RATIO_NEXTN_LAYER + from sglang.srt.models.deepseek_v4_nextn import ( + COMPRESS_RATIO_NEXTN_LAYER, + ) compression_ratios = [ COMPRESS_RATIO_NEXTN_LAYER @@ -1436,7 +1442,9 @@ def _build_ascend_swa_kv_pool( full_max_total_num_tokens: Optional[int], swa_max_total_num_tokens: Optional[int], ) -> KVCache: - from sglang.srt.hardware_backend.npu.memory_pool_npu import NPUMHATokenToKVPool + from sglang.srt.hardware_backend.npu.memory_pool_npu import ( + NPUMHATokenToKVPool, + ) kwargs = {} if self.is_hybrid_swa_compress: @@ -1502,7 +1510,9 @@ def _build_ascend_minimax_sparse_kv_pool( def _build_ascend_mla_kv_pool( self, *, max_total_num_tokens: int, is_dsa_model: bool ) -> KVCache: - from sglang.srt.hardware_backend.npu.memory_pool_npu import NPUMLATokenToKVPool + from sglang.srt.hardware_backend.npu.memory_pool_npu import ( + NPUMLATokenToKVPool, + ) token_to_kv_pool = NPUMLATokenToKVPool( max_total_num_tokens, @@ -1520,7 +1530,9 @@ def _build_ascend_mla_kv_pool( return token_to_kv_pool def _build_ascend_mha_kv_pool(self, *, max_total_num_tokens: int) -> KVCache: - from sglang.srt.hardware_backend.npu.memory_pool_npu import NPUMHATokenToKVPool + from sglang.srt.hardware_backend.npu.memory_pool_npu import ( + NPUMHATokenToKVPool, + ) token_to_kv_pool = NPUMHATokenToKVPool( max_total_num_tokens, diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index e21b61f56746..95d1187510fb 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -408,7 +408,9 @@ def _compute_dsa_indexer_cell_size( indexer_ratio = parse_hisparse_config().host_to_device_ratio - from sglang.srt.mem_cache.kv_cache_configurator import _should_elide_dsa_index_k + from sglang.srt.mem_cache.kv_cache_configurator import ( + _should_elide_dsa_index_k, + ) if allocate_all_layers or not _should_elide_dsa_index_k( is_draft_worker=kvc.is_draft_worker @@ -1123,14 +1125,13 @@ def _fixed_c4_state_bytes(self, max_running_requests: int) -> int: return 0 c4_state_dtype_size, _ = _get_dsv4_compress_state_dtype_sizes() - attn_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim # CompressStatePool allocates `size + ring_size + 1` rows, padded to the # compress ratio (see CompressStatePool.__init__). Mirror that here so the # reserved bias covers the real allocation. state_rows = self._unified_c4_state_pool_size(max_running_requests) state_rows = ceil_div(state_rows + self.c4_ring_size + 1, 4) * 4 # overlap c4: last_dim = 2 * (1 + overlap) * head_dim = 4 * head_dim. - core_bytes = 4 * attn_head_dim * c4_state_dtype_size + core_bytes = 4 * self.attn_head_dim * c4_state_dtype_size indexer_bytes = 4 * self.indexer_head_dim * c4_state_dtype_size return state_rows * (core_bytes + indexer_bytes) * self.num_layers_ca4 diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index bd60d6a68f51..69f951dee34e 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -9,7 +9,10 @@ PrefillAdder, estimate_prefill_extend_tile_metrics, ) -from sglang.srt.mem_cache.base_prefix_cache import DecLockRefResult, IncLockRefResult +from sglang.srt.mem_cache.base_prefix_cache import ( + DecLockRefResult, + IncLockRefResult, +) from sglang.srt.runtime_context import get_context from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler from sglang.srt.utils.common import Range From cb6614d5324d9cd2d1af29a176e1a93773cd510e Mon Sep 17 00:00:00 2001 From: yuttian1 Date: Sun, 6 Sep 2026 07:30:38 +0000 Subject: [PATCH 04/15] [AMD][DSV4] Gate alloc_extend_swa_tail on the unified-KV switch alloc_extend, alloc_decode and new_pages_available all take a unified early return; alloc_extend_swa_tail did not, while new_pages_available stopped consulting num_swa_pages on that path. The vestigial paged SWA allocator is therefore called with no capacity gate, and a bare assert turns its exhaustion into a scheduler crash instead of a None backoff. Reachable via DecodePreallocQueue, which routes DSV4 page_size > 1 preallocation through this method. --- python/sglang/srt/mem_cache/allocator/swa.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index c5abd1d756e7..dc681f426ce7 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -310,6 +310,21 @@ def alloc_extend_swa_tail( if not self.new_pages_available(num_full_pages, num_swa_pages): return None + if self._unified: + # Unified SWA ring is slot-addressed and not paged here, so the tail + # needs no SWA allocation and no full->swa mapping, as in + # alloc_extend. new_pages_available already ignored num_swa_pages, + # so the paged SWA allocator below has no capacity gate left. + return self.full_attn_allocator.alloc_extend( + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + last_loc, + extend_num_tokens, + num_new_pages=num_full_pages, + ) + alloc_full_indices = self.full_attn_allocator.alloc_extend( prefix_lens, prefix_lens_cpu, From 9ad761ce640912994b3fcda2533824d2cee922e9 Mon Sep 17 00:00:00 2001 From: yuttian1 Date: Sun, 6 Sep 2026 08:15:33 +0000 Subject: [PATCH 05/15] [AMD][DSV4] Make unified-KV gates identity-compare, fix SWA debug_print Every unified-KV gate now spells the check `... is True` instead of relying on truthiness. `_unified_kv` is a plain bool on the real pool, so this is a no-op in production, but a duck-typed stub (a bare MagicMock auto-creates a truthy attribute) would otherwise select the unified path in tests. Six sites still used the truthy form: the allocator's master switch in SWATokenToKVPoolAllocator.__init__, the four CompressorHip ring-address branches, the pool-stats observer, the SWA leak-invariant skip, the swa_max_total_num_tokens reconciliation, and the decode-bottleneck diagnostic. SWATokenToKVPoolAllocator.debug_print() reported swa_attn_allocator's available size, which on the unified path is the vestigial paged allocator and contradicts swa_available_size(). It now calls swa_available_size(), so the two agree; off the unified path the two are the same expression. Also condenses four comment blocks added earlier in this series that ran past the two-line limit in .claude/rules/comment-style.md. --- .../srt/layers/attention/dsv4/compress_hip.py | 10 ++++++---- python/sglang/srt/managers/schedule_batch.py | 2 +- .../scheduler_components/invariant_checker.py | 13 +++---------- .../pool_stats_observer.py | 11 +++-------- python/sglang/srt/mem_cache/allocator/swa.py | 14 ++++---------- .../srt/mem_cache/kv_cache_configurator.py | 19 ++++++------------- 6 files changed, 23 insertions(+), 46 deletions(-) diff --git a/python/sglang/srt/layers/attention/dsv4/compress_hip.py b/python/sglang/srt/layers/attention/dsv4/compress_hip.py index 7ae66e27bb3c..a6822b867c60 100644 --- a/python/sglang/srt/layers/attention/dsv4/compress_hip.py +++ b/python/sglang/srt/layers/attention/dsv4/compress_hip.py @@ -145,7 +145,8 @@ def compress_extend_paged( seq_len=prefix_lens[i], ratio=self.ratio ).to(device) if self.ratio == 128 or ( - self.ratio == 4 and getattr(token_to_kv_pool, "_unified_kv", False) + self.ratio == 4 + and getattr(token_to_kv_pool, "_unified_kv", False) is True ): state_loc = state_pool.translate_from_req_position_to_state_loc( req_pool_indices[i], pre_state_indices @@ -169,7 +170,8 @@ def compress_extend_paged( assert post_state_len <= valid_kv_len if self.ratio == 128 or ( - self.ratio == 4 and getattr(token_to_kv_pool, "_unified_kv", False) + self.ratio == 4 + and getattr(token_to_kv_pool, "_unified_kv", False) is True ): post_state_loc = state_pool.translate_from_req_position_to_state_loc( req_pool_indices[i], post_state_indices @@ -276,7 +278,7 @@ def compress_decode_paged( req_pool_indices = req_pool_indices.repeat_interleave(draft_tokens) if self.ratio == 128 or ( - self.ratio == 4 and getattr(token_to_kv_pool, "_unified_kv", False) + self.ratio == 4 and getattr(token_to_kv_pool, "_unified_kv", False) is True ): state_locs = state_pool.translate_from_req_position_to_state_loc( req_pool_indices, seq_lens - 1 @@ -293,7 +295,7 @@ def compress_decode_paged( ) compress_indices.clamp_(min=-1) if self.ratio == 128 or ( - self.ratio == 4 and getattr(token_to_kv_pool, "_unified_kv", False) + self.ratio == 4 and getattr(token_to_kv_pool, "_unified_kv", False) is True ): compress_indices_state = ( state_pool.translate_from_req_position_to_state_loc( diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 95db5c369a5a..1d843f14f20e 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -3085,7 +3085,7 @@ def check_decode_mem(self, selected_indices: Optional[List[int]] = None): ok = allocator.check_decode_capacity( num_tokens=num_tokens, tree_cache=self.tree_cache ) - if not ok and getattr(allocator.get_kvcache(), "_unified_kv", False): + if not ok and getattr(allocator.get_kvcache(), "_unified_kv", False) is True: self._log_unified_swa_bottleneck(allocator, num_tokens, selected_indices) return ok diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index 5e19c6e5a115..a640344661f5 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -153,16 +153,9 @@ def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str] def _check_swa_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]: allocator = self.token_to_kv_pool_allocator kv = allocator.get_kvcache() - if getattr(kv, "_unified_kv", False): - # Unified-KV DSV4: SWA is a fixed per-request ring, reused per request - # and released together with the req_pool slot (which has its own - # leak check). swa_available_size() is deliberately non-binding (it - # always reports the full ring so it never throttles admission), and - # cached radix prefixes still report swa_evictable even though the - # completed request already freed its ring slot. The token-pool - # invariant (available + evictable + protected + session == total) - # therefore does not model this pool -- skip it to avoid a spurious - # leak. Ring-slot leaks are still caught by the req_to_token check. + if getattr(kv, "_unified_kv", False) is True: + # Unified-KV DSV4: a per-request SWA ring does not satisfy the token-pool + # invariant; ring-slot leaks are caught by the req_to_token check instead. return False, ( "[swa] unified ring (leak-check skipped): " f"available={ps.swa_available_size}, " diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index 18d730ff169a..532513aa1230 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -301,15 +301,10 @@ def _get_swa_token_info(self) -> PoolStats: swa_available_size = allocator.swa_available_size() full_evictable_size = self.tree_cache.full_evictable_size() swa_evictable_size = self.tree_cache.swa_evictable_size() - # Unified-KV DSV4: SWA is a fixed per-request ring, released with the - # req_pool slot. Cached radix prefixes still report swa_evictable even - # though the completed request already freed its ring slot, and - # swa_available_size() is non-binding (always the full ring). Counting - # that evictable here would double-count against the ring and drive - # swa_num_used / swa_token_usage negative. The ring holds nothing - # evictable, so zero it out to keep the usage stats coherent. + # Unified-KV DSV4: the SWA ring is released with the req slot, but cached + # radix prefixes still report swa_evictable; counting it drives usage negative. _swa_kv = self.token_to_kv_pool_allocator.get_kvcache() - if getattr(_swa_kv, "_unified_kv", False): + if getattr(_swa_kv, "_unified_kv", False) is True: swa_evictable_size = 0 full_num_used = self.full_tokens_per_layer - ( full_available_size + full_evictable_size diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index dc681f426ce7..5cfbc5949369 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -110,15 +110,9 @@ def __init__( self._kvcache = kvcache - # Unified-KV (DSV4): SWA is a per-request ring addressed by state_slot - # (== req_pool_idx) + position inside the DSV4 kernels. The paged SWA - # indices / full_to_swa_index_mapping produced here are NOT consumed on - # that path, so treating SWA as a linearly-consumed token pool - # over-throttles admission and decode retract. Instead account for it as - # a fixed per-request ring slot; the real bound is concurrency - # (num_req_slots), already enforced by req_to_token_pool / - # max_running_requests. - self._unified = getattr(kvcache, "_unified_kv", False) + # Unified-KV DSV4: the DSV4 kernels address SWA as a per-request ring, so + # the paged indices built here are unused and the bound is num_req_slots. + self._unified = getattr(kvcache, "_unified_kv", False) is True self._req_to_token_pool = req_to_token_pool if self._unified: ring_size = getattr(kvcache, "unified_swa_ring_size", self.page_size) @@ -189,7 +183,7 @@ def size_full(self): def debug_print(self) -> str: msg = "" - msg += f"#swa-available-size: {self.swa_attn_allocator.available_size()}, " + msg += f"#swa-available-size: {self.swa_available_size()}, " msg += ( f"#full-attn-available-size: {self.full_attn_allocator.available_size()}, " ) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 0143eee5568c..4ea3bb7e7759 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -337,26 +337,19 @@ def configure(self, *, pre_model_load_memory: int) -> KVCacheConfigResult: ) swa_max_total_num_tokens = sizes.swa_max_total_num_tokens - # Unified-KV DSV4: SWA is a fixed per-request ring, so the allocator - # reports ring capacity (free_req_slots * ring_cost) from - # swa_available_size(), while swa_max_total_num_tokens was sized from - # the (vestigial, unallocated) full_token-scaled SWA pool. The idle - # pool-leak invariant requires swa total == swa available, so - # reconcile the reported SWA total to the allocator's actual idle ring - # capacity. Safe: on unified_kv swa_kv_pool is None, so no real buffer - # is resized -- this only fixes token accounting / usage reporting. + # Unified-KV DSV4: swa_max_total_num_tokens was sized from the vestigial + # SWA pool; reconcile it to real ring capacity so the idle-leak invariant + # holds. Safe: swa_kv_pool is None here, so no buffer is resized. if ( self.is_hybrid_swa and not self.is_draft_worker - and getattr(pools.token_to_kv_pool, "_unified_kv", False) + and getattr(pools.token_to_kv_pool, "_unified_kv", False) is True ): alloc = pools.token_to_kv_pool_allocator if hasattr(alloc, "swa_available_size"): ring_capacity = int(alloc.swa_available_size()) - # Only reconcile downward to the (smaller) ring capacity. A - # value >= the current total means swa_available_size() hit a - # non-binding fallback (e.g. req_to_token pool not wired), in - # which case leave the reported total untouched. + # Only reconcile downward: a value >= the current total means + # swa_available_size() hit its non-binding fallback. if 0 < ring_capacity < swa_max_total_num_tokens: logger.info( "Unified-KV: reconciling swa_max_total_num_tokens " From ac57b0b7c032b5f0f5b0845ffb9f24f5708a8880 Mon Sep 17 00:00:00 2001 From: yuttian1 Date: Sun, 6 Sep 2026 09:38:48 +0000 Subject: [PATCH 06/15] [AMD][DSV4] Finish the identity-compare sweep, assert on XPU + req ring The two compressor entry points still gated use_req_ring on bare truthiness. Both take a DeepSeekV4TokenToKVPool, so read _unified_kv directly; the same holds for the four compress_hip sites, which sit after an isinstance assert. The XPU compress plan builder has no use_req_ring parameter, so a unified-KV plan would silently fall back to the SWA-paged layout. Assert instead. Co-Authored-By: Claude Opus 5 --- python/sglang/kernels/ops/attention/dsv4/compress.py | 10 ++++++++++ .../sglang/srt/layers/attention/dsv4/compress_hip.py | 10 ++++------ python/sglang/srt/layers/attention/dsv4/compressor.py | 2 +- .../sglang/srt/layers/attention/dsv4/compressor_v2.py | 2 +- python/sglang/srt/mem_cache/kv_cache_configurator.py | 4 +++- 5 files changed, 19 insertions(+), 9 deletions(-) diff --git a/python/sglang/kernels/ops/attention/dsv4/compress.py b/python/sglang/kernels/ops/attention/dsv4/compress.py index f126ba175e1a..cfa4778827a0 100644 --- a/python/sglang/kernels/ops/attention/dsv4/compress.py +++ b/python/sglang/kernels/ops/attention/dsv4/compress.py @@ -179,6 +179,12 @@ def generate( int(swa_page_size), int(ring_size), ) + # The XPU plan builder has no use_req_ring parameter, so the unified-KV + # request-ring addressing cannot be expressed there. Fail loudly rather + # than silently planning against the SWA-paged layout. + 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)) @@ -289,6 +295,10 @@ def generate( int(swa_page_size), int(ring_size), ) + # See plan_decode: XPU cannot express the unified-KV request ring. + 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 diff --git a/python/sglang/srt/layers/attention/dsv4/compress_hip.py b/python/sglang/srt/layers/attention/dsv4/compress_hip.py index a6822b867c60..1c05e92b4f0d 100644 --- a/python/sglang/srt/layers/attention/dsv4/compress_hip.py +++ b/python/sglang/srt/layers/attention/dsv4/compress_hip.py @@ -145,8 +145,7 @@ def compress_extend_paged( seq_len=prefix_lens[i], ratio=self.ratio ).to(device) if self.ratio == 128 or ( - self.ratio == 4 - and getattr(token_to_kv_pool, "_unified_kv", False) is True + self.ratio == 4 and token_to_kv_pool._unified_kv is True ): state_loc = state_pool.translate_from_req_position_to_state_loc( req_pool_indices[i], pre_state_indices @@ -170,8 +169,7 @@ def compress_extend_paged( assert post_state_len <= valid_kv_len if self.ratio == 128 or ( - self.ratio == 4 - and getattr(token_to_kv_pool, "_unified_kv", False) is True + self.ratio == 4 and token_to_kv_pool._unified_kv is True ): post_state_loc = state_pool.translate_from_req_position_to_state_loc( req_pool_indices[i], post_state_indices @@ -278,7 +276,7 @@ def compress_decode_paged( req_pool_indices = req_pool_indices.repeat_interleave(draft_tokens) if self.ratio == 128 or ( - self.ratio == 4 and getattr(token_to_kv_pool, "_unified_kv", False) is True + self.ratio == 4 and token_to_kv_pool._unified_kv is True ): state_locs = state_pool.translate_from_req_position_to_state_loc( req_pool_indices, seq_lens - 1 @@ -295,7 +293,7 @@ def compress_decode_paged( ) compress_indices.clamp_(min=-1) if self.ratio == 128 or ( - self.ratio == 4 and getattr(token_to_kv_pool, "_unified_kv", False) is True + self.ratio == 4 and token_to_kv_pool._unified_kv is True ): compress_indices_state = ( state_pool.translate_from_req_position_to_state_loc( diff --git a/python/sglang/srt/layers/attention/dsv4/compressor.py b/python/sglang/srt/layers/attention/dsv4/compressor.py index 2e420f65367a..684fe1131341 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor.py @@ -264,7 +264,7 @@ def create_paged_compressor_data( ) -> FusedCompressMetadata: swa_page_size = token_to_kv_pool.swa_page_size ring_size = token_to_kv_pool.get_ring_size(compress_ratio=compress_ratio) - use_req_ring = compress_ratio == 4 and token_to_kv_pool._unified_kv + use_req_ring = compress_ratio == 4 and token_to_kv_pool._unified_kv is True # assert ring_size % compress_ratio == 0 def clip_down(positions: torch.Tensor) -> torch.Tensor: diff --git a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py index c3515413119b..04f33ba45c47 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py @@ -441,7 +441,7 @@ def create_paged_compressor_data( swa_page_size = token_to_kv_pool.swa_page_size ring_size = token_to_kv_pool.get_ring_size(compress_ratio=compress_ratio) - use_req_ring = compress_ratio == 4 and token_to_kv_pool._unified_kv + use_req_ring = compress_ratio == 4 and token_to_kv_pool._unified_kv is True # NOTE: This is actually a proxy, which encounter some bug with tvm-ffi. # As a workaround, we use `.detach()` to get the real tensor. full_to_swa = token_to_kv_pool.full_to_swa_index_mapping.detach() diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 4ea3bb7e7759..c7d476fedcb8 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -1969,7 +1969,9 @@ def _build_token_to_kv_pool_allocator( ) else: if get_memory().enable_hisparse: - from sglang.srt.mem_cache.sparsity import parse_hisparse_config + from sglang.srt.mem_cache.sparsity import ( + parse_hisparse_config, + ) hisparse_cfg = parse_hisparse_config() token_to_kv_pool_allocator = HiSparseTokenToKVPoolAllocator( From 718e9e95a40f04d09d57b57a5c57f39235d25a44 Mon Sep 17 00:00:00 2001 From: yuttian1 Date: Sun, 6 Sep 2026 09:38:48 +0000 Subject: [PATCH 07/15] [AMD][DSV4] Fix test_pool_configurator fixture for the _unified attribute _compute_dsv4_sizes now reads self._unified, but the white-box fixture builds DSV4PoolConfigurator via object.__new__ and only set five attributes, so test_dsv4_accepts_pool_above_floor raised AttributeError. Set it, and cover the unified branch where the c4 state pool is sized per request instead. Co-Authored-By: Claude Opus 5 --- .../unit/model_executor/test_pool_configurator.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index b76f5420dd7a..21d4e22dcd7b 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -1042,9 +1042,9 @@ def test_hybrid_swa_accepts_pool_above_floor(self): ) self.assertEqual(config.swa_max_total_num_tokens, 3072) - def _dsv4_sizes(self, max_tokens, page_size): + def _dsv4_sizes(self, max_tokens, page_size, unified=False): """Exercise the DSV4 size arithmetic without a full V4 model fixture: - _compute_dsv4_sizes reads only these five attributes.""" + _compute_dsv4_sizes reads only these six attributes.""" from sglang.srt.model_executor.pool_configurator import DSV4PoolConfigurator cfg = object.__new__(DSV4PoolConfigurator) @@ -1053,6 +1053,7 @@ def _dsv4_sizes(self, max_tokens, page_size): cfg.swa_page_size = 128 cfg.c4_ring_size = 8 cfg.c4_shrink_factor = 1 + cfg._unified = unified return cfg._compute_dsv4_sizes(max_tokens, page_size) def test_dsv4_rejects_single_page_pool(self): @@ -1066,6 +1067,16 @@ def test_dsv4_accepts_pool_above_floor(self): sizes = self._dsv4_sizes(max_tokens=32768, page_size=256) self.assertEqual(sizes.full_max_total_num_tokens, 32768) self.assertEqual(sizes.swa_max_total_num_tokens, 3072) + # Non-unified: the c4 state pool scales with the paged SWA pool. + self.assertEqual(sizes.c4_state_pool_size, 3072 // 128 * 8) + + def test_dsv4_unified_c4_state_not_token_scaled(self): + # Unified-KV sizes the c4 state ring from max_running_requests in + # finalize_with_max_running_requests, so it must not scale here. + sizes = self._dsv4_sizes(max_tokens=32768, page_size=256, unified=True) + self.assertEqual(sizes.full_max_total_num_tokens, 32768) + self.assertEqual(sizes.swa_max_total_num_tokens, 3072) + self.assertEqual(sizes.c4_state_pool_size, 0) if __name__ == "__main__": From a3c5d4c32c12c0d4868363821d8fee9a0b6625fa Mon Sep 17 00:00:00 2001 From: yuttian1 Date: Sun, 6 Sep 2026 10:02:43 +0000 Subject: [PATCH 08/15] [AMD][DSV4] Assert the token-cap monotonicity the fixed-pool bias relies on calculate_pool_sizes subtracts the three request-scoped fixed pools from the byte budget; calculate_pool_sizes_from_max_tokens cannot, because it takes a token count and subtracting bytes there would double-count. That is only safe while every constraint in config_from_budget is a min(), which nothing checked. Assert it, state the caller contract, and add a regression test that caps a budget-derived token count and requires the total footprint to shrink. Co-Authored-By: Claude Opus 5 --- .../srt/mem_cache/kv_cache_configurator.py | 7 +++ .../srt/model_executor/pool_configurator.py | 7 +++ .../model_executor/test_pool_configurator.py | 62 +++++++++++++++++++ 3 files changed, 76 insertions(+) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index c7d476fedcb8..1709db211e4d 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -2262,6 +2262,13 @@ def config_from_budget( max_tokens = self._apply_token_constraints(config.max_total_num_tokens) if cap_tokens is not None: max_tokens = min(max_tokens, cap_tokens) + # calculate_pool_sizes_from_max_tokens takes a token count, not a byte + # budget, so it cannot re-subtract the request-scoped fixed pools; it is + # only safe because every constraint above is a min(). Assert that here. + assert max_tokens <= config.max_total_num_tokens, ( + f"token constraints must not raise capacity: {max_tokens} > " + f"{config.max_total_num_tokens}" + ) if max_tokens != config.max_total_num_tokens: # Token-capped re-derivation: the profiled budget no longer # applies; the recalced config's unified_total_bytes stays None diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 95d1187510fb..7665b176a8b6 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -1245,6 +1245,13 @@ def calculate_pool_sizes( def calculate_pool_sizes_from_max_tokens( self, max_total_num_tokens: int, page_size: int ) -> MemoryPoolConfig: + """Caller contract: max_total_num_tokens must not exceed the value + calculate_pool_sizes derived from the same budget (config_from_budget + asserts this). No fixed-pool bias is subtracted here -- the input is a + token count, not a byte budget, so subtracting bytes would double-count + what calculate_pool_sizes already removed. The request-scoped pools are + re-derived from the capped token count by + finalize_with_max_running_requests, so they shrink with it.""" assert page_size % 128 == 0, ( "page_size must be multiple of 128 for compressed attention" ) diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 21d4e22dcd7b..0878fce0c2be 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -1070,6 +1070,68 @@ def test_dsv4_accepts_pool_above_floor(self): # Non-unified: the c4 state pool scales with the paged SWA pool. self.assertEqual(sizes.c4_state_pool_size, 3072 // 128 * 8) + def test_dsv4_token_cap_never_grows_total_footprint(self): + """Regression: the token-cap path subtracts no fixed-pool bias, which is + only safe while every constraint in config_from_budget is a min(). Cap + the budget-derived token count and check the total footprint shrinks.""" + cfg = self._dsv4_configurator_for_budget() + page_size = 128 + budget = 256 * (1 << 30) + base = cfg.calculate_pool_sizes(budget, page_size) + base_bytes = self._dsv4_total_bytes(cfg, base.max_total_num_tokens) + self.assertLessEqual(base_bytes, budget) + for numerator in (999, 900, 500, 100, 1): + capped_tokens = ( + base.max_total_num_tokens * numerator // 1000 // page_size * page_size + ) + if capped_tokens <= 0: + continue + capped = cfg.calculate_pool_sizes_from_max_tokens(capped_tokens, page_size) + capped_bytes = self._dsv4_total_bytes(cfg, capped.max_total_num_tokens) + with self.subTest(numerator=numerator): + self.assertLessEqual(capped_bytes, base_bytes) + + def _dsv4_configurator_for_budget(self): + """DSV4 configurator with a 671B-class shape, built white-box so the + byte arithmetic runs without a real model fixture.""" + from sglang.srt.model_executor.pool_configurator import DSV4PoolConfigurator + + cfg = object.__new__(DSV4PoolConfigurator) + cfg.qk_nope_head_dim, cfg.qk_rope_head_dim = 128, 64 + cfg.attn_head_dim = 192 + cfg.indexer_head_dim = 128 + cfg.num_layers_total = 61 + cfg.num_layers_ca4 = 61 + cfg.num_layers_ca128 = 61 + cfg.c4_ring_size = 8 + cfg.c128_ring_size = 128 + cfg._swa_ring_size = 128 + cfg._spec_infl = 1.0 + cfg.context_len = 65536 + cfg.bytes_per_full_token = 576.0 + cfg.requested_max_running_requests_per_worker = None + cfg.swa_ratio = 0.1 + cfg.sliding_window_size = 4096 + cfg.swa_page_size = 128 + cfg.c4_shrink_factor = 1 + cfg.online_c128_mtp_max_draft_tokens = 0 + cfg.disaggregation_mode = None + cfg.disaggregation_decode_extra_slots = 0 + cfg._unified = True + return cfg + + def _dsv4_total_bytes(self, cfg, tokens): + """Token pool plus the three request-scoped fixed pools, sized from the + concurrency that resolve_max_num_reqs derives from this token count.""" + estimated = max(min(int(tokens / cfg.context_len * 512), 4096), 2048) + max_running_requests = min(estimated, tokens // 2) + return int( + tokens * cfg.bytes_per_full_token + + cfg._fixed_swa_bytes(max_running_requests) + + cfg._fixed_c4_state_bytes(max_running_requests) + + cfg._get_c128_state_fixed_bytes(max_running_requests) + ) + def test_dsv4_unified_c4_state_not_token_scaled(self): # Unified-KV sizes the c4 state ring from max_running_requests in # finalize_with_max_running_requests, so it must not scale here. From 8882029500c6dafbd3226e21d7d14cc9c1c9ba05 Mon Sep 17 00:00:00 2001 From: yuttian1 Date: Sun, 6 Sep 2026 10:56:56 +0000 Subject: [PATCH 09/15] chore(dsv4): trim comment prose on the unified-KV C4 request-ring branch Apply .claude/rules/comment-style.md to the branch's added comments: drop function-top preambles whose facts are already stated at the lines they constrain, drop docstrings on private helpers, and condense multi-line rationale to the single non-recoverable fact. Kept: cross-file mirrors (swa_ring_size, CompressStatePool row math), layout and unit facts, magic-number derivations, the kernel write-pad aliasing bound, and the token-count-vs-byte-budget caller contract. Co-Authored-By: Claude Opus 5 --- .../kernels/jit/csrc/deepseek_v4/c_plan.cuh | 22 +++---- .../kernels/ops/attention/dsv4/compress.py | 5 +- python/sglang/srt/managers/schedule_batch.py | 5 +- python/sglang/srt/managers/schedule_policy.py | 30 +++------ python/sglang/srt/mem_cache/allocator/swa.py | 14 ++-- .../srt/mem_cache/deepseek_v4_memory_pool.py | 16 ++--- .../srt/mem_cache/kv_cache_configurator.py | 6 +- .../srt/model_executor/pool_configurator.py | 65 +++++-------------- .../unit/managers/test_prefill_adder.py | 8 +-- .../test_swa_alloc_extend_page_estimation.py | 6 +- .../model_executor/test_pool_configurator.py | 12 ++-- 11 files changed, 58 insertions(+), 131 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh index 5e23da74dfb8..c87f23e89f76 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh @@ -48,8 +48,7 @@ struct Prefill0Params { int32_t compress_ratio; int32_t swa_page_size; /// \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. + /// ring; the bound is derived in `plan_compress_prefill`. int32_t mtp_pad; bool use_req_ring; }; @@ -513,24 +512,19 @@ inline PrefillPlan plan_compress_prefill( // `swa_page_size` >= `ring_size` >= `compress_ratio` RuntimeCheck(swa_page_size % ring_size == 0 && ring_size % compress_ratio == 0); // Write pad: trailing tokens kept resident so a verify batch's committed tail - // survives any accept length. Zero without speculation -- nothing rolls back, - // and the ring is then exactly one window wide. Otherwise the ring bounds it: - // a write at `w` aliases onto `w - ring_size`, and the earliest position a - // future compression still needs is `prefix_len - window_size + 2` (the next - // batch commits >= 1 token, and `run_prefill` launches the compress kernel - // before the write kernel, so a batch's own compressions read the pre-write - // ring). Padding past the extend range is harmless: the loops only span - // `[prefix_len, seq_len)`. + // survives any accept length. A write at `w` aliases onto `w - ring_size`, and + // the earliest position a future compression still needs is + // `prefix_len - window_size + 2` -- the next batch commits >= 1 token, and + // `run_prefill` launches compress before write, so a batch's own compressions + // read the pre-write ring. const auto mtp_pad = ring_size > window_size ? ring_size - window_size + 2 : 0; const auto device = device_.unwrap(); const auto stream = LaunchKernel::resolve_device(device); if (cpu_or_gpu.unwrap().device_type == kDLGPU) { - // GPU input path: kernel0 builds the (CPU-loop-equivalent) plan metadata - // directly on device, padding to num_q_tokens with invalid; kernel_1 then - // finalizes the SWA-translated read/write locations. Used for MTP / - // cuda-graph capture where a host sync would be expensive. + // GPU input path for MTP / cuda-graph capture, where a host sync would be + // expensive: kernel0 builds the plan metadata, kernel_1 translates SWA locs. RuntimeCheck(batch_size <= kMaxPrefillBatchSize, "GPU plan only support batch size up to ", kMaxPrefillBatchSize); auto C = ffi::empty({num_q_tokens, sizeof(PlanC)}, kDLUInt8, device); auto W = ffi::empty({num_q_tokens, sizeof(PlanW)}, kDLUInt8, device); diff --git a/python/sglang/kernels/ops/attention/dsv4/compress.py b/python/sglang/kernels/ops/attention/dsv4/compress.py index cfa4778827a0..8dadcadc75bd 100644 --- a/python/sglang/kernels/ops/attention/dsv4/compress.py +++ b/python/sglang/kernels/ops/attention/dsv4/compress.py @@ -179,9 +179,8 @@ def generate( int(swa_page_size), int(ring_size), ) - # The XPU plan builder has no use_req_ring parameter, so the unified-KV - # request-ring addressing cannot be expressed there. Fail loudly rather - # than silently planning against the SWA-paged layout. + # The XPU plan builder has no use_req_ring parameter, so unified-KV + # request-ring addressing cannot be expressed there. assert not (_is_xpu and use_req_ring), ( "use_req_ring is not supported by the XPU compress plan builder" ) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 1d843f14f20e..857eca34c1ef 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -3090,9 +3090,8 @@ def check_decode_mem(self, selected_indices: Optional[List[int]] = None): return ok def _log_unified_swa_bottleneck(self, allocator, num_tokens, selected_indices): - """Diagnostic (unified-KV only): when check_decode_mem is short, compare - the SWA token bookkeeping against the real per-slot ring utilization. - Throttled to avoid log floods during retract storms.""" + """Compare the SWA token bookkeeping against real per-slot ring + utilization; throttled so retract storms cannot flood the log.""" global _last_swa_bottleneck_log now = time.monotonic() if now - _last_swa_bottleneck_log < 1.0: diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 7a1bac2e3e62..d1e3eb427ec4 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -675,11 +675,8 @@ def rem_total_tokens(self): def rem_swa_tokens(self): allocator = self.token_to_kv_pool_allocator if self._unified_kv: - # Unified-KV: SWA is a per-request ring, not a tree-reusable token - # pool. swa_available_size() already reports ring capacity - # (free_slots * ring_cost). tree swa_evictable is in the old linear - # token unit and freeing it does not release ring space, so exclude - # it here to keep a single consistent accounting unit. + # swa_available_size() already reports ring capacity; tree + # swa_evictable is in linear token units and frees no ring space. return allocator.swa_available_size() - self.rem_swa_token_offset return ( allocator.swa_available_size() @@ -731,10 +728,8 @@ def _swa_budget_for_req( """ allocator = self.token_to_kv_pool_allocator if self._unified_kv: - # Unified-KV: each request occupies exactly one fixed SWA ring slot, - # independent of context / chunk length; a host-hit prefix reuses the - # same ring. Budget the fixed per-slot ring cost (paired with the - # ring-based swa_available_size on the allocator). + # One fixed ring slot per request, independent of context or chunk + # length; pairs with the ring-based swa_available_size. return allocator.swa_ring_cost_tokens if self.rem_chunk_tokens is not None: alloc = min(extend_input_len, self.rem_chunk_tokens) @@ -887,10 +882,8 @@ def _update_prefill_budget( self.rem_input_tokens -= extend_input_len if self.is_hybrid_swa: - # Unified-KV: SWA is a fixed per-request ring slot reserved once at - # first admission and already reflected in swa_available_size() on - # later rounds. Charging it again on a chunked continuation would - # double-count the slot and over-throttle admission, so skip it. + # The ring slot is reserved once at first admission; charging it + # again on a continuation would double-count and over-throttle. if not (self._unified_kv and is_chunked_continuation): self.rem_swa_token_offset += self._swa_budget_for_req( extend_input_len, max_new_tokens @@ -1032,10 +1025,8 @@ def add_chunked_req(self, req: Req): if self.is_hybrid_swa and not self._unified_kv: # alloc_extend needs extend_num_tokens + page_size per request, # so reserve one page here to avoid OOM. - # Unified-KV: rem_swa_tokens is ring capacity (free_slots * ring - # cost), not a linear per-chunk token budget, and this request's - # ring slot is already reserved -- mixing units here would wrongly - # truncate the chunk, so skip the SWA clamp. + # Unified-KV: rem_swa_tokens is ring capacity, not a per-chunk + # token budget; clamping against it would truncate the chunk. _rem_tokens = min( _rem_tokens, int(self.rem_swa_tokens) - self.page_size ) @@ -1284,9 +1275,8 @@ def add_one_req( self._swa_new_tokens(req), swa_host_hit_length=req.swa_host_hit_length, ) - # Unified-KV: rem_swa_tokens is an exact ring-slot capacity, so a - # request needing exactly what is left still fits. The legacy - # SWA-token path keeps its original conservative `>=`. + # Ring-slot capacity is exact, so needing exactly what is left still + # fits; the legacy SWA-token path keeps its conservative `>=`. if ( swa_needed > self.rem_swa_tokens if self._unified_kv diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index 5cfbc5949369..a1003474208b 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -241,9 +241,8 @@ def alloc_extend( return None if self._unified: - # Unified SWA ring is slot-addressed and not paged here: allocate only - # the full-attention KV and skip the vestigial SWA allocator / mapping - # (unused by the DSV4 kernels). + # The unified SWA ring is slot-addressed, not paged here, so the + # vestigial paged allocator and full->swa mapping are skipped. return self.full_attn_allocator.alloc_extend( prefix_lens, prefix_lens_cpu, @@ -305,10 +304,8 @@ def alloc_extend_swa_tail( return None if self._unified: - # Unified SWA ring is slot-addressed and not paged here, so the tail - # needs no SWA allocation and no full->swa mapping, as in - # alloc_extend. new_pages_available already ignored num_swa_pages, - # so the paged SWA allocator below has no capacity gate left. + # See alloc_extend. new_pages_available already ignored + # num_swa_pages, so the paged allocator has no capacity gate left. return self.full_attn_allocator.alloc_extend( prefix_lens, prefix_lens_cpu, @@ -366,8 +363,7 @@ def alloc_decode( ): assert self.page_size > 1 if self._unified: - # See alloc_extend: unified SWA ring is slot-addressed, allocate full - # only and skip the vestigial SWA allocator / mapping. + # See alloc_extend: slot-addressed ring, so full-attention KV only. return self.full_attn_allocator.alloc_decode( seq_lens, seq_lens_cpu, last_loc ) diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index ae4d00b16b32..30aec1900417 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -1057,11 +1057,8 @@ def get_online_c128_mtp_pending_seq_lens(self) -> torch.Tensor: return self.online_c128_mtp_pending_seq_lens def clear_c4_req_states(self, req_pool_indices: Sequence[int]) -> None: - """Reset newly allocated unified C4 attention and indexer state rings. - - Only the request-owned rows are touched. The extra sentinel/ring padding - allocated by :class:`CompressStatePool` remains intact. - """ + """Reset the request-owned C4 rows; the sentinel/ring padding that + CompressStatePool allocates is left intact.""" if not self._unified_kv or not req_pool_indices: return @@ -1114,12 +1111,9 @@ def clear_unaccepted_c128_draft_states( num_draft_tokens: int, ) -> None: """Clear offline C128 ring slots written for rejected speculative tokens. - - C4 needs no equivalent cleanup: draft states are written in position order, - and every rejected position is overwritten before it can become the prior - state of a later accepted token. C128 cleanup is required because its - compression boundary can consume a previously written draft slot directly. - """ + C4 needs none: its draft states are overwritten in position order before + they can be read, while a C128 compression boundary can consume a + previously written draft slot directly.""" if ONLINE_C128 or num_draft_tokens <= 1 or req_pool_indices.numel() == 0: return diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 1709db211e4d..71195a1c7043 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -338,8 +338,7 @@ def configure(self, *, pre_model_load_memory: int) -> KVCacheConfigResult: swa_max_total_num_tokens = sizes.swa_max_total_num_tokens # Unified-KV DSV4: swa_max_total_num_tokens was sized from the vestigial - # SWA pool; reconcile it to real ring capacity so the idle-leak invariant - # holds. Safe: swa_kv_pool is None here, so no buffer is resized. + # SWA pool; reconcile to real ring capacity. swa_kv_pool is None here. if ( self.is_hybrid_swa and not self.is_draft_worker @@ -2263,8 +2262,7 @@ def config_from_budget( if cap_tokens is not None: max_tokens = min(max_tokens, cap_tokens) # calculate_pool_sizes_from_max_tokens takes a token count, not a byte - # budget, so it cannot re-subtract the request-scoped fixed pools; it is - # only safe because every constraint above is a min(). Assert that here. + # budget; it cannot re-subtract the fixed pools, so capacity must not rise. assert max_tokens <= config.max_total_num_tokens, ( f"token constraints must not raise capacity: {max_tokens} > " f"{config.max_total_num_tokens}" diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 7665b176a8b6..af148fc36222 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -847,10 +847,8 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): Splits available memory across full / swa / c4 / c128 + c4_state / c128_state pools. coeff is bytes_per_full_token (inflated by (T+D)/T when speculative - decode reserves a draft worker, mirroring dflash's cell_size scaling). The - bias is the sum of request-scoped fixed pools that do not scale with - full_token: the c128 state pool and, on the unified_kv path, the fixed SWA - per-request ring (bf16, see _fixed_swa_bytes). + decode reserves a draft worker, mirroring dflash's cell_size scaling). bias + is the request-scoped fixed pools that do not scale with full_token. """ def __init__(self, kvc: KVCacheConfigurator): @@ -907,13 +905,6 @@ def __init__(self, kvc: KVCacheConfigurator): self.num_layers_ca4 = sum(1 for r in self.compression_ratios if r == 4) self.num_layers_ca128 = sum(1 for r in self.compression_ratios if r == 128) - # Unified-KV uses a different physical layout than the fp8 path: - # * KV is stored bf16 over the full latent (attn_head_dim * 2 bytes), - # not the fp8(nope) + bf16(rope) + scales 584-byte cell. - # * SWA is a fixed per-request ring (num_req_slots * ring_size), - # independent of full_token, so it is a fixed *bias* rather than a - # per-token term. Gate on the same switch the pool itself uses so the - # sizing and the allocation never drift apart. from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( is_unified_kv_triton, ) @@ -1034,13 +1025,8 @@ def _get_bytes_per_full_token(self) -> float: + c4_frac * kv_bytes * self.num_layers_ca4 + 1 / 128 * kv_bytes * self.num_layers_ca128 + 1 / 4 * self.indexer_bytes_per_token * self.num_layers_ca4 - # Unified_kv: the c4 (attn + indexer) compress-state is a ring buffer - # addressed off the SWA slot ((swa_loc // swa_page_size) * ring_size), - # and the unified SWA pool is a fixed per-request ring - # (swa_pages = num_req_slots * swa_ring_size), so the state ring is - # request-scoped, not full_token-scoped. It is therefore a fixed bias - # (see _fixed_c4_state_bytes), not a per-token term. On the non-unified - # path the SWA pool scales with full_token, so it stays per-token. + # Unified_kv: c4 state is addressed off the SWA slot, and that pool + # is a fixed per-request ring, so the state ring is request-scoped. + ( 0.0 if self._unified @@ -1069,9 +1055,7 @@ def _compute_dsv4_sizes(self, full_token: int, page_size: int) -> _DSV4PoolSizes swa_max_total_num_tokens=swa_tokens, c4_max_total_num_tokens=full_token // (4 * self.c4_shrink_factor), c128_max_total_num_tokens=full_token // 128, - # Unified_kv sizes the c4 state ring from the fixed SWA ring - # (request-scoped), finalized once max_running_requests is known -- so - # it must not scale with full_token here (mirrors c128_state below). + # Unified_kv: request-scoped, finalized once concurrency is known. c4_state_pool_size=( 0 if self._unified @@ -1108,26 +1092,17 @@ def _get_c128_state_fixed_bytes(self, max_running_requests: int) -> int: ) def _unified_c4_state_pool_size(self, max_running_requests: int) -> int: - """Exact request-scoped C4 ring size for the unified address contract. - - Unified C4 state locations are - ``req_pool_idx * c4_ring_size + position % c4_ring_size``. - """ + # Unified C4 state loc is req_pool_idx * c4_ring_size + pos % c4_ring_size. num_req_slots = self._get_num_req_slots(max_running_requests) return num_req_slots * self.c4_ring_size def _fixed_c4_state_bytes(self, max_running_requests: int) -> int: - """Unified_kv c4 (attn + indexer) compress-state is a fixed per-request - ring, sized by concurrency rather than full_token. Return its byte - footprint across all c4 layers. Returns 0 on the non-unified path (where - the c4 state pool scales with the SWA pool and is accounted per-token).""" if not self._unified or self.num_layers_ca4 == 0: return 0 c4_state_dtype_size, _ = _get_dsv4_compress_state_dtype_sizes() - # CompressStatePool allocates `size + ring_size + 1` rows, padded to the - # compress ratio (see CompressStatePool.__init__). Mirror that here so the - # reserved bias covers the real allocation. + # Mirror CompressStatePool.__init__: it allocates `size + ring_size + 1` + # rows, padded to the compress ratio. state_rows = self._unified_c4_state_pool_size(max_running_requests) state_rows = ceil_div(state_rows + self.c4_ring_size + 1, 4) * 4 # overlap c4: last_dim = 2 * (1 + overlap) * head_dim = 4 * head_dim. @@ -1136,9 +1111,8 @@ def _fixed_c4_state_bytes(self, max_running_requests: int) -> int: return state_rows * (core_bytes + indexer_bytes) * self.num_layers_ca4 def _resolve_max_running_requests_per_worker(self, available_bytes: int) -> int: - """Approximate ModelRunner._resolve_max_num_reqs closely enough to size - the request-scoped fixed pools (c128 state, unified SWA ring). Over- - estimating is safe: a larger fixed bias yields a smaller full_token.""" + # Approximates ModelRunner._resolve_max_num_reqs. Over-estimating is safe: + # a larger fixed bias yields a smaller full_token. if self.requested_max_running_requests_per_worker is not None: return self.requested_max_running_requests_per_worker @@ -1148,11 +1122,6 @@ def _resolve_max_running_requests_per_worker(self, available_bytes: int) -> int: return min(estimated, full_token // 2) def _fixed_swa_bytes(self, max_running_requests: int) -> int: - """Unified_kv SWA is a fixed per-request ring, sized by concurrency - (num_req_slots) rather than by full_token. Return its bf16 byte - footprint across all full layers, inflated for the draft worker the same - way as the per-token coeff. Returns 0 on the non-unified path (where SWA - is already accounted per-token).""" if not self._unified: return 0 num_req_slots = self._get_num_req_slots(max_running_requests) @@ -1194,9 +1163,8 @@ def finalize_with_max_running_requests( config.c128_state_pool_size = num_req_slots else: config.c128_state_pool_size = num_req_slots * self.c128_ring_size - # Unified_kv: the c4 state ring is request-scoped (fixed SWA pool), so - # finalize it here from the now-known concurrency. On the non-unified path - # it was already sized from full_token in _compute_dsv4_sizes. + # Unified_kv: request-scoped, so it is sized here from the now-known + # concurrency rather than from full_token in _compute_dsv4_sizes. if self._unified and self.num_layers_ca4 > 0: config.c4_state_pool_size = self._unified_c4_state_pool_size( config.max_running_requests @@ -1245,13 +1213,10 @@ def calculate_pool_sizes( def calculate_pool_sizes_from_max_tokens( self, max_total_num_tokens: int, page_size: int ) -> MemoryPoolConfig: - """Caller contract: max_total_num_tokens must not exceed the value + """Caller contract: max_total_num_tokens must not exceed what calculate_pool_sizes derived from the same budget (config_from_budget - asserts this). No fixed-pool bias is subtracted here -- the input is a - token count, not a byte budget, so subtracting bytes would double-count - what calculate_pool_sizes already removed. The request-scoped pools are - re-derived from the capped token count by - finalize_with_max_running_requests, so they shrink with it.""" + asserts it). Subtracting the fixed-pool bias again here would + double-count -- the input is a token count, not a byte budget.""" assert page_size % 128 == 0, ( "page_size must be multiple of 128 for compressed attention" ) diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index 69f951dee34e..818115187e4e 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -71,12 +71,8 @@ def create_token_allocator( allocator.swa_available_size.return_value = swa_available_size allocator.available_size.return_value = available_size allocator.size_swa = size_swa - # get_kvcache().[_unified_kv] gates the unified-KV SWA-ring accounting - # path in schedule_policy.add_chunked_req / rem_swa_tokens. A bare - # MagicMock auto-creates any attribute access as a truthy Mock, so - # without this the getattr(..., "_unified_kv", False) default never - # triggers and these tests silently exercise the unified-KV branch - # instead of the standard hybrid-SWA one they intend to cover. + # A bare MagicMock auto-creates a truthy `_unified_kv`, silently flipping + # these tests onto the unified-KV branch; pin it to the hybrid-SWA path. allocator.get_kvcache.return_value._unified_kv = False return allocator diff --git a/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py b/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py index e3de143f2e21..70208bfa24f4 100644 --- a/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py +++ b/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py @@ -36,10 +36,8 @@ def set_full_to_swa_mapping( return SimpleNamespace( page_size=page_size, - # alloc_extend branches on self._unified to skip the vestigial paged SWA - # allocator on the unified-KV path. This stub exercises the standard - # hybrid-SWA path, so pin it False rather than letting the attribute go - # missing (SimpleNamespace raises instead of defaulting). + # alloc_extend branches on _unified; pin it to the hybrid-SWA path + # (SimpleNamespace raises instead of defaulting if it goes missing). _unified=False, full_attn_allocator=SimpleNamespace( available_size=lambda: full_available, diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 0878fce0c2be..4183ce70a38a 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -1071,9 +1071,8 @@ def test_dsv4_accepts_pool_above_floor(self): self.assertEqual(sizes.c4_state_pool_size, 3072 // 128 * 8) def test_dsv4_token_cap_never_grows_total_footprint(self): - """Regression: the token-cap path subtracts no fixed-pool bias, which is - only safe while every constraint in config_from_budget is a min(). Cap - the budget-derived token count and check the total footprint shrinks.""" + """Regression: the token-cap path subtracts no fixed-pool bias, so + capping the budget-derived token count must still shrink the total.""" cfg = self._dsv4_configurator_for_budget() page_size = 128 budget = 256 * (1 << 30) @@ -1091,9 +1090,8 @@ def test_dsv4_token_cap_never_grows_total_footprint(self): with self.subTest(numerator=numerator): self.assertLessEqual(capped_bytes, base_bytes) + # White-box 671B-class shape: the byte arithmetic runs without a model fixture. def _dsv4_configurator_for_budget(self): - """DSV4 configurator with a 671B-class shape, built white-box so the - byte arithmetic runs without a real model fixture.""" from sglang.srt.model_executor.pool_configurator import DSV4PoolConfigurator cfg = object.__new__(DSV4PoolConfigurator) @@ -1120,9 +1118,9 @@ def _dsv4_configurator_for_budget(self): cfg._unified = True return cfg + # Token pool plus the three request-scoped fixed pools, sized from the + # concurrency resolve_max_num_reqs derives from this token count. def _dsv4_total_bytes(self, cfg, tokens): - """Token pool plus the three request-scoped fixed pools, sized from the - concurrency that resolve_max_num_reqs derives from this token count.""" estimated = max(min(int(tokens / cfg.context_len * 512), 4096), 2048) max_running_requests = min(estimated, tokens // 2) return int( From c062146f6ca8c2f692d37bec6ce96d370925a0ff Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 7 Sep 2026 02:48:14 -0700 Subject: [PATCH 10/15] typed swa_req_ring flag instead of getattr; c4 reset via req pool alloc hook --- python/sglang/srt/disaggregation/decode.py | 22 ++- python/sglang/srt/managers/schedule_batch.py | 46 +----- python/sglang/srt/managers/schedule_policy.py | 25 ++-- .../scheduler_components/invariant_checker.py | 15 +- .../pool_stats_observer.py | 8 +- python/sglang/srt/mem_cache/allocation.py | 25 +--- .../srt/mem_cache/allocator/hisparse.py | 4 - python/sglang/srt/mem_cache/allocator/swa.py | 48 ++++--- .../srt/mem_cache/base_swa_memory_pool.py | 6 +- .../srt/mem_cache/deepseek_v4_memory_pool.py | 1 + .../srt/mem_cache/kv_cache_configurator.py | 41 +++--- python/sglang/srt/mem_cache/memory_pool.py | 13 +- .../unit/managers/test_prefill_adder.py | 3 - .../unit/mem_cache/test_dllm_fdfo_kv_reuse.py | 3 - .../mem_cache/test_dsv4_c4_state_lifecycle.py | 135 ++++++++++-------- .../unit/mem_cache/test_hisparse_allocator.py | 1 - .../test_swa_alloc_extend_page_estimation.py | 4 +- 17 files changed, 180 insertions(+), 220 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index a9ae8d8661b0..24daaecdebbe 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/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 @@ -138,6 +138,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, @@ -203,6 +206,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: @@ -220,6 +225,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__( @@ -1768,22 +1777,11 @@ def _pre_alloc( if total_prefix_len is None: total_prefix_len = prefix_len - is_new_req_slot = req.kv.req_pool_idx is None req_pool_indices = self.req_to_token_pool.alloc([req]) assert req_pool_indices is not None, ( "req_pool_indices is full! There is a bug in memory estimation." ) - if ( - is_new_req_slot - and getattr(self.token_to_kv_pool, "_unified_kv", False) is True - ): - # Unified-KV DSV4 only: reset the stale C4 ring on a fresh slot. - clear_c4_req_states = getattr( - self.token_to_kv_pool, "clear_c4_req_states", None - ) - if clear_c4_req_states is not None: - clear_c4_req_states(req_pool_indices) fill_len = self._pre_alloc_fill_len(req) req.kv.kv_committed_len = fill_len diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 857eca34c1ef..db3d02c2921a 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -57,7 +57,6 @@ import logging import re import sys -import time from array import array from concurrent.futures import Future from enum import Enum, auto @@ -162,9 +161,6 @@ logger = logging.getLogger(__name__) -# Throttle for the unified-KV SWA bottleneck diagnostic (seconds). -_last_swa_bottleneck_log = 0.0 - ReturnHiddenStatesMode = Union[bool, Literal["last"]] @@ -3081,49 +3077,9 @@ def check_decode_mem(self, selected_indices: Optional[List[int]] = None): shortfalls retract gracefully instead of tripping fail-loud alloc errors.""" num_tokens = self.new_tokens_required_next_decode(selected_indices) - allocator = self.token_to_kv_pool_allocator - ok = allocator.check_decode_capacity( + return self.token_to_kv_pool_allocator.check_decode_capacity( num_tokens=num_tokens, tree_cache=self.tree_cache ) - if not ok and getattr(allocator.get_kvcache(), "_unified_kv", False) is True: - self._log_unified_swa_bottleneck(allocator, num_tokens, selected_indices) - return ok - - def _log_unified_swa_bottleneck(self, allocator, num_tokens, selected_indices): - """Compare the SWA token bookkeeping against real per-slot ring - utilization; throttled so retract storms cannot flood the log.""" - global _last_swa_bottleneck_log - now = time.monotonic() - if now - _last_swa_bottleneck_log < 1.0: - return - _last_swa_bottleneck_log = now - try: - full_avail = allocator.full_available_size() - swa_avail = allocator.swa_available_size() - reqs = ( - self.reqs - if selected_indices is None - else [self.reqs[i] for i in selected_indices] - ) - active_slots = len( - {int(r.kv.req_pool_idx) for r in reqs if r.kv.req_pool_idx is not None} - ) - unified = getattr(allocator.get_kvcache(), "unified_kv_pool", None) - if unified is not None: - num_slots = unified.num_slots - ring_util = ( - active_slots * unified.swa_ring_size / max(unified.swa_pages, 1) - ) - else: - num_slots, ring_util = -1, -1.0 - logger.warning( - "[SWA-BOTTLENECK] check_decode_mem short: " - f"need={num_tokens}, full_avail={full_avail}, swa_avail={swa_avail}, " - f"active_slots={active_slots}/{num_slots}, " - f"ring_util_upper={ring_util:.4f}" - ) - except Exception as e: # diagnostics must never break scheduling - logger.warning(f"[SWA-BOTTLENECK] logging failed: {e}") def retract_decode(self) -> Tuple[List[Req], float, List[Req]]: """Retract the decoding requests when there is not enough memory.""" diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index d1e3eb427ec4..c4a75f748ce0 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -52,6 +52,7 @@ from sglang.srt.mem_cache.allocator.swa import ( PureSWATokenToKVPoolAllocator, SWATokenToKVPoolAllocator, + is_swa_req_ring, ) from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( UnifiedMambaSWATokenToKVPoolAllocator, @@ -502,16 +503,8 @@ def __init__( self.prefill_tile_block_m = prefill_tile_block_m self.tree_cache = tree_cache self.token_to_kv_pool_allocator = token_to_kv_pool_allocator - # Unified-KV DSV4 sizes SWA as a fixed per-request ring slot, not a - # token budget. `is True`: a duck-typed stub must not select that path. - self._unified_kv = ( - getattr( - getattr(token_to_kv_pool_allocator, "get_kvcache", lambda: None)(), - "_unified_kv", - False, - ) - is True - ) + # Per-request SWA ring: one fixed slot per request, not a token budget. + self._swa_req_ring = is_swa_req_ring(token_to_kv_pool_allocator) self.running_batch = running_batch self.new_token_ratio = new_token_ratio self.rem_input_tokens = rem_input_tokens - num_mixed_decode_tokens @@ -674,7 +667,7 @@ def rem_total_tokens(self): @property def rem_swa_tokens(self): allocator = self.token_to_kv_pool_allocator - if self._unified_kv: + if self._swa_req_ring: # swa_available_size() already reports ring capacity; tree # swa_evictable is in linear token units and frees no ring space. return allocator.swa_available_size() - self.rem_swa_token_offset @@ -727,7 +720,7 @@ def _swa_budget_for_req( from double-counting extend, so budget <= extend + max_new_tokens + page. """ allocator = self.token_to_kv_pool_allocator - if self._unified_kv: + if self._swa_req_ring: # One fixed ring slot per request, independent of context or chunk # length; pairs with the ring-based swa_available_size. return allocator.swa_ring_cost_tokens @@ -884,7 +877,7 @@ def _update_prefill_budget( if self.is_hybrid_swa: # The ring slot is reserved once at first admission; charging it # again on a continuation would double-count and over-throttle. - if not (self._unified_kv and is_chunked_continuation): + if not (self._swa_req_ring and is_chunked_continuation): self.rem_swa_token_offset += self._swa_budget_for_req( extend_input_len, max_new_tokens ) @@ -1022,7 +1015,7 @@ def add_chunked_req(self, req: Req): _rem_tokens = self._get_dllm_remain_tokens() else: _rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens)) - if self.is_hybrid_swa and not self._unified_kv: + if self.is_hybrid_swa and not self._swa_req_ring: # alloc_extend needs extend_num_tokens + page_size per request, # so reserve one page here to avoid OOM. # Unified-KV: rem_swa_tokens is ring capacity, not a per-chunk @@ -1279,7 +1272,7 @@ def add_one_req( # fits; the legacy SWA-token path keeps its conservative `>=`. if ( swa_needed > self.rem_swa_tokens - if self._unified_kv + if self._swa_req_ring else swa_needed >= self.rem_swa_tokens ): if not self._swa_req_never_fits( @@ -1319,7 +1312,7 @@ def add_one_req( ) if ( swa_needed > self.rem_swa_tokens - if self._unified_kv + if self._swa_req_ring else swa_needed >= self.rem_swa_tokens ): if not self._swa_req_never_fits( diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index a640344661f5..a7340ca20346 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -21,6 +21,7 @@ SchedulerPoolStatsObserver, ) 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.allocator.unified_hybrid_swa import ( UnifiedMambaSWATokenToKVPoolAllocator, ) @@ -152,10 +153,9 @@ def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str] def _check_swa_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]: allocator = self.token_to_kv_pool_allocator - kv = allocator.get_kvcache() - if getattr(kv, "_unified_kv", False) is True: - # Unified-KV DSV4: a per-request SWA ring does not satisfy the token-pool - # invariant; ring-slot leaks are caught by the req_to_token check instead. + if is_swa_req_ring(allocator): + # Per-request SWA ring: there is no token pool to conserve; ring-slot + # leaks are caught by the req_to_token check instead. return False, ( "[swa] unified ring (leak-check skipped): " f"available={ps.swa_available_size}, " @@ -398,13 +398,6 @@ def _add_owner(req_or_slot, label, rpi, committed, allocated): # Sub-allocators to check: a flat allocator is its own single sub; a # hybrid-SWA wrapper exposes full_attn_allocator + swa_attn_allocator. alloc = self.token_to_kv_pool_allocator - # Unified-KV DSV4-HiSparse nests the real SWA allocator one level - # down; elsewhere the wrapper is the object this invariant asserts on. - if ( - getattr(getattr(alloc, "get_kvcache", lambda: None)(), "_unified_kv", False) - is True - ): - alloc = getattr(alloc, "logical_attn_allocator", alloc) sub_allocs = ( [alloc] if getattr(alloc, "free_pages", None) is not None diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index 532513aa1230..9e1b962bac21 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -11,6 +11,7 @@ Tuple, ) +from sglang.srt.mem_cache.allocator.swa import is_swa_req_ring from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( UnifiedMambaSWATokenToKVPoolAllocator, ) @@ -301,10 +302,9 @@ def _get_swa_token_info(self) -> PoolStats: swa_available_size = allocator.swa_available_size() full_evictable_size = self.tree_cache.full_evictable_size() swa_evictable_size = self.tree_cache.swa_evictable_size() - # Unified-KV DSV4: the SWA ring is released with the req slot, but cached - # radix prefixes still report swa_evictable; counting it drives usage negative. - _swa_kv = self.token_to_kv_pool_allocator.get_kvcache() - if getattr(_swa_kv, "_unified_kv", False) is True: + # Per-request SWA ring: released with the req slot, yet cached radix + # prefixes still report swa_evictable; counting it drives usage negative. + if is_swa_req_ring(self.token_to_kv_pool_allocator): swa_evictable_size = 0 full_num_used = self.full_tokens_per_layer - ( full_available_size + full_evictable_size diff --git a/python/sglang/srt/mem_cache/allocation.py b/python/sglang/srt/mem_cache/allocation.py index 78affd3f3654..5e109a4c897e 100644 --- a/python/sglang/srt/mem_cache/allocation.py +++ b/python/sglang/srt/mem_cache/allocation.py @@ -230,7 +230,6 @@ def alloc_req_slots( req_to_token_pool: ReqToTokenPool, reqs: list[Req], tree_cache: BasePrefixCache | None, - token_to_kv_pool=None, ) -> list[int]: """Allocate request slots from the pool. @@ -261,18 +260,6 @@ def alloc_req_slots( tree_cache.evict_for_alloc( EvictParams(num_tokens=0, mamba_num=mamba_num) ) - # Unified-KV DSV4 only: a freshly allocated req slot carries stale C4 ring - # state. Resolve the hook before alloc() overwrites req_pool_idx. - clear_c4_req_states = ( - getattr(token_to_kv_pool, "clear_c4_req_states", None) - if getattr(token_to_kv_pool, "_unified_kv", False) is True - else None - ) - newly_allocated = ( - [req.kv.req_pool_idx is None for req in reqs] - if clear_c4_req_states is not None - else None - ) req_pool_indices = req_to_token_pool.alloc(reqs) if req_pool_indices is None: raise RuntimeError( @@ -280,13 +267,6 @@ def alloc_req_slots( "Please set a smaller number for `--max-running-requests`. " f"{req_to_token_pool.available_size()=}, {num_reqs=}, " ) - - if clear_c4_req_states is not None: - new_req_pool_indices = [ - idx for idx, is_new in zip(req_pool_indices, newly_allocated) if is_new - ] - if new_req_pool_indices: - clear_c4_req_states(new_req_pool_indices) return req_pool_indices @@ -331,10 +311,7 @@ def alloc_for_extend( # Allocate req slots (raises RuntimeError if the pool is exhausted) req_pool_indices = alloc_req_slots( - batch.req_to_token_pool, - batch.reqs, - batch.tree_cache, - token_to_kv_pool=batch.token_to_kv_pool_allocator.get_kvcache(), + batch.req_to_token_pool, batch.reqs, batch.tree_cache ) req_pool_indices_cpu = torch.tensor( req_pool_indices, dtype=torch.int64, pin_memory=pin_memory diff --git a/python/sglang/srt/mem_cache/allocator/hisparse.py b/python/sglang/srt/mem_cache/allocator/hisparse.py index 44047c9b5c02..5647154f70b1 100644 --- a/python/sglang/srt/mem_cache/allocator/hisparse.py +++ b/python/sglang/srt/mem_cache/allocator/hisparse.py @@ -343,10 +343,6 @@ def debug_print(self) -> str: def get_kvcache(self): return self._kvcache - @property - def swa_ring_cost_tokens(self) -> int: - return self.logical_attn_allocator.swa_ring_cost_tokens - def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor): return self.logical_attn_allocator.translate_loc_from_full_to_swa(kv_indices) diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index a1003474208b..bb7368682014 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -32,6 +32,10 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): """Allocator for SWA hybrid KV cache.""" + # Per-request SWA ring (BaseSWAKVPool.swa_req_ring_size). Class default so + # subclasses that bypass this __init__ read False. + _swa_req_ring = False + def __init__( self, size: int, @@ -110,22 +114,19 @@ def __init__( self._kvcache = kvcache - # Unified-KV DSV4: the DSV4 kernels address SWA as a per-request ring, so - # the paged indices built here are unused and the bound is num_req_slots. - self._unified = getattr(kvcache, "_unified_kv", False) is True + # Per-request SWA ring: the paged SWA indices built here are unused and + # SWA capacity is bounded by req slots, not tokens. + ring_size = kvcache.swa_req_ring_size + self._swa_req_ring = ring_size is not None self._req_to_token_pool = req_to_token_pool - if self._unified: - ring_size = getattr(kvcache, "unified_swa_ring_size", self.page_size) + if self._swa_req_ring: self._swa_ring_cost = ( (ring_size + self.page_size - 1) // self.page_size ) * self.page_size logger.info( - "[SWA-BOOKKEEPING] unified ring accounting enabled: " - f"num_slots={getattr(kvcache, 'num_req_slots', '?')}, " - f"swa_ring_size={ring_size}, " - f"ring_cost_tokens={self._swa_ring_cost}, " - f"unified_swa_pages={getattr(kvcache, 'unified_swa_pages', '?')} | " - f"legacy paged size_swa={self._size_swa} (bypassed)" + "SWA per-request ring accounting enabled: " + f"ring_size={ring_size}, ring_cost_tokens={self._swa_ring_cost}, " + f"paged size_swa={self._size_swa} (bypassed)" ) else: self._swa_ring_cost = 0 @@ -133,13 +134,18 @@ def __init__( self.clear() self._kvcache.register_mapping(self.full_to_swa_index_mapping) + @property + def swa_req_ring(self) -> bool: + """SWA is a per-request ring; no per-token SWA budget applies.""" + return self._swa_req_ring + @property def swa_ring_cost_tokens(self) -> int: - """Unified: paged SWA cost of one request's ring slot (0 otherwise).""" + """Ring mode: paged SWA cost of one request's ring slot (0 otherwise).""" return self._swa_ring_cost def available_size(self): - if self._unified: + if self._swa_req_ring: # The SWA ring is pre-allocated per slot and reused by decode, so it # never constrains token growth; full attention is the real limiter. return self.full_attn_allocator.available_size() @@ -152,7 +158,7 @@ def full_available_size(self): return self.full_attn_allocator.available_size() def swa_available_size(self): - if self._unified: + if self._swa_req_ring: # Ring-based availability: free request slots * per-slot ring cost. # Fall back to non-binding if the req pool wasn't wired in. if self._req_to_token_pool is None: @@ -216,7 +222,7 @@ def new_pages_available(self, num_full_pages: int, num_swa_pages: int) -> bool: num_full_pages <= self.full_attn_allocator.available_size() // self.page_size ) - if self._unified: + if self._swa_req_ring: # SWA ring rows are pre-allocated per slot; no per-token SWA paging. return full_ok return full_ok and ( @@ -240,7 +246,7 @@ def alloc_extend( if not self.new_pages_available(num_new_pages, num_new_pages): return None - if self._unified: + if self._swa_req_ring: # The unified SWA ring is slot-addressed, not paged here, so the # vestigial paged allocator and full->swa mapping are skipped. return self.full_attn_allocator.alloc_extend( @@ -303,7 +309,7 @@ def alloc_extend_swa_tail( if not self.new_pages_available(num_full_pages, num_swa_pages): return None - if self._unified: + if self._swa_req_ring: # See alloc_extend. new_pages_available already ignored # num_swa_pages, so the paged allocator has no capacity gate left. return self.full_attn_allocator.alloc_extend( @@ -362,7 +368,7 @@ def alloc_decode( last_loc: torch.Tensor, # last_loc for full layers ): assert self.page_size > 1 - if self._unified: + if self._swa_req_ring: # See alloc_extend: slot-addressed ring, so full-attention KV only. return self.full_attn_allocator.alloc_decode( seq_lens, seq_lens_cpu, last_loc @@ -693,3 +699,9 @@ def free_group_end(self): def clear(self): self.swa_attn_allocator.clear() self.free_group = None + + +def is_swa_req_ring(allocator) -> bool: + """True when the allocator's SWA side is a per-request ring (see + BaseSWAKVPool.swa_req_ring_size), so SWA carries no per-token budget.""" + return isinstance(allocator, SWATokenToKVPoolAllocator) and allocator.swa_req_ring diff --git a/python/sglang/srt/mem_cache/base_swa_memory_pool.py b/python/sglang/srt/mem_cache/base_swa_memory_pool.py index 06ff43984c2c..24a61da749f8 100644 --- a/python/sglang/srt/mem_cache/base_swa_memory_pool.py +++ b/python/sglang/srt/mem_cache/base_swa_memory_pool.py @@ -1,5 +1,5 @@ import abc -from typing import List, Tuple +from typing import List, Optional, Tuple import torch @@ -15,6 +15,10 @@ class BaseSWAKVPool(KVCache): """ swa_kv_pool: KVCache + # Set when SWA KV is a fixed per-request ring of this many tokens, addressed + # by req_pool_idx, instead of a paged token pool. The paged SWA allocator is + # then vestigial and SWA must not be budgeted per token. + swa_req_ring_size: Optional[int] = None @abc.abstractmethod def register_mapping(self, full_to_swa_index_mapping: torch.Tensor) -> None: diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index 30aec1900417..804d2d5cae6d 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -656,6 +656,7 @@ def __init__( self.unified_swa_window = self.sliding_window self.unified_swa_ring_size = self.sliding_window + spec_extra self.unified_swa_pages = self.unified_kv_pool.swa_pages + self.swa_req_ring_size = self.unified_swa_ring_size else: self.unified_kv_pool = None self.swa_kv_pool = self._make_kv_pool( diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 71195a1c7043..6dda4eb9091a 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -46,6 +46,7 @@ from sglang.srt.mem_cache.allocator.swa import ( PureSWATokenToKVPoolAllocator, SWATokenToKVPoolAllocator, + is_swa_req_ring, ) from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( UnifiedSWATokenToKVPoolAllocator, @@ -337,25 +338,20 @@ def configure(self, *, pre_model_load_memory: int) -> KVCacheConfigResult: ) swa_max_total_num_tokens = sizes.swa_max_total_num_tokens - # Unified-KV DSV4: swa_max_total_num_tokens was sized from the vestigial - # SWA pool; reconcile to real ring capacity. swa_kv_pool is None here. - if ( - self.is_hybrid_swa - and not self.is_draft_worker - and getattr(pools.token_to_kv_pool, "_unified_kv", False) is True - ): - alloc = pools.token_to_kv_pool_allocator - if hasattr(alloc, "swa_available_size"): - ring_capacity = int(alloc.swa_available_size()) - # Only reconcile downward: a value >= the current total means - # swa_available_size() hit its non-binding fallback. - if 0 < ring_capacity < swa_max_total_num_tokens: - logger.info( - "Unified-KV: reconciling swa_max_total_num_tokens " - f"{swa_max_total_num_tokens} -> {ring_capacity} " - "(fixed per-request SWA ring capacity)." - ) - swa_max_total_num_tokens = ring_capacity + # Per-request SWA ring: swa_max_total_num_tokens was sized from the + # vestigial paged SWA pool; reconcile to the real ring capacity. + alloc = pools.token_to_kv_pool_allocator + if not self.is_draft_worker and is_swa_req_ring(alloc): + ring_capacity = int(alloc.swa_available_size()) + # Only reconcile downward: a value >= the current total means + # swa_available_size() hit its non-binding fallback. + if 0 < ring_capacity < swa_max_total_num_tokens: + logger.info( + "SWA ring: reconciling swa_max_total_num_tokens " + f"{swa_max_total_num_tokens} -> {ring_capacity} " + "(fixed per-request SWA ring capacity)." + ) + swa_max_total_num_tokens = ring_capacity logger.info( f"Memory pool end. " @@ -1368,6 +1364,13 @@ def _build_dsv4_kv_pool( enable_hisparse=get_memory().enable_hisparse, online_mtp_max_draft_tokens=(max_speculative_num_draft_tokens() or 0), ) + if not self.is_draft_worker and token_to_kv_pool._unified_kv: + # Unified-KV C4 state is a per-request ring: reset a row's ring + # whenever its req slot is handed out again. The draft pool has no + # C4 layers and shares this req pool, so only the target registers. + req_to_token_pool.register_on_alloc_rows( + token_to_kv_pool.clear_c4_req_states + ) return token_to_kv_pool def _build_oot_dsa_kv_pool(self, *, max_total_num_tokens: int) -> KVCache: diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 59cec0fb538b..51f8ced084c4 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -31,7 +31,7 @@ from contextlib import contextmanager, nullcontext from dataclasses import dataclass, fields from functools import cached_property -from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple, Union import numpy as np import torch @@ -258,6 +258,9 @@ class ReqToTokenPool: """A memory pool that maps a request to its token locations.""" enable_mamba_extra_buffer_lazy: bool = False + # Class default: some decode pools borrow another __init__ (see + # DecodeReqToTokenPool) but inherit alloc_rows. + _on_alloc_rows: Optional[Callable[[List[int]], None]] = None def __init__( self, @@ -321,6 +324,8 @@ def alloc_rows(self, need_size: int) -> Optional[List[int]]: select_index = self.free_slots[-need_size:] del self.free_slots[-need_size:] self.req_generation[select_index] += 1 + if self._on_alloc_rows is not None: + self._on_alloc_rows(select_index) return select_index def free_rows(self, indices: List[int]) -> None: @@ -346,6 +351,12 @@ def attach_aux_cache(self, aux_cache: Any) -> None: assert self._aux_cache is None self._aux_cache = aux_cache + def register_on_alloc_rows(self, hook: Callable[[List[int]], None]) -> None: + """Receive the row indices of every fresh allocation. Per-request + state pools use it to reset a row's state when the slot is reused.""" + assert self._on_alloc_rows is None + self._on_alloc_rows = hook + def reset_aux_cache_allocator(self) -> None: if self._aux_cache is not None: self._aux_cache.reset_allocator() diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index 818115187e4e..d74db97d6cbb 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -71,9 +71,6 @@ def create_token_allocator( allocator.swa_available_size.return_value = swa_available_size allocator.available_size.return_value = available_size allocator.size_swa = size_swa - # A bare MagicMock auto-creates a truthy `_unified_kv`, silently flipping - # these tests onto the unified-KV branch; pin it to the hybrid-SWA path. - allocator.get_kvcache.return_value._unified_kv = False return allocator def create_running_batch(self, reqs=None) -> MagicMock: diff --git a/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py b/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py index 54440519b62c..6152e59cc4d4 100644 --- a/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py +++ b/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py @@ -23,9 +23,6 @@ def __init__(self, base=1000, page_size=1): self.alloc_calls = [] self.extend_calls = [] - def get_kvcache(self): - return None - def available_size(self): return 1 << 30 diff --git a/test/registered/unit/mem_cache/test_dsv4_c4_state_lifecycle.py b/test/registered/unit/mem_cache/test_dsv4_c4_state_lifecycle.py index 11102b4e5378..c0b2b7a1a8e1 100644 --- a/test/registered/unit/mem_cache/test_dsv4_c4_state_lifecycle.py +++ b/test/registered/unit/mem_cache/test_dsv4_c4_state_lifecycle.py @@ -6,6 +6,7 @@ import torch +from sglang.srt.disaggregation.decode import DecodeReqToTokenPool from sglang.srt.mem_cache.allocation import alloc_req_slots from sglang.srt.mem_cache.deepseek_v4_compress_state import KVAndScore from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool @@ -28,6 +29,13 @@ def _request(req_pool_idx=None, *, reused=False): ) +def _mark_reused(req): + req.kv.kv_committed_len = 1 + req.kv.kv_allocated_len = 1 + req.kv.holds_kv = True + req.inflight_middle_chunks = 1 + + def _c4_pool(rows: int, width: int, ring_size: int): return SimpleNamespace( ratio=4, @@ -36,6 +44,24 @@ def _c4_pool(rows: int, width: int, ring_size: int): ) +def _token_pool(unified: bool, ring_size: int = 8): + logical_rows = 4 * ring_size + physical_rows = logical_rows + ring_size + 4 + attn = _c4_pool(physical_rows, width=12, ring_size=ring_size) + indexer = _c4_pool(physical_rows, width=8, ring_size=ring_size) + c128 = SimpleNamespace( + ratio=128, + ring_size=128, + kv_score_buffer=KVAndScore(torch.full((physical_rows, 8), 9.0)), + ) + token_pool = object.__new__(DeepSeekV4TokenToKVPool) + token_pool._unified_kv = unified + token_pool.compress_state_pools = [attn, c128] + token_pool.indexer_compress_state_pools = [indexer, None] + token_pool.get_ring_size = MagicMock(return_value=ring_size) + return token_pool, attn, indexer, c128, logical_rows + + class TestUnifiedC4StateLifecycle(unittest.TestCase): def test_pool_size_is_exact_request_ring_product(self): configurator = object.__new__(DSV4PoolConfigurator) @@ -47,22 +73,10 @@ def test_pool_size_is_exact_request_ring_product(self): def test_clear_resets_only_selected_request_rings(self): ring_size = 8 - logical_rows = 4 * ring_size - physical_rows = logical_rows + ring_size + 4 - attn = _c4_pool(physical_rows, width=12, ring_size=ring_size) - indexer = _c4_pool(physical_rows, width=8, ring_size=ring_size) - c128 = SimpleNamespace( - ratio=128, - ring_size=128, - kv_score_buffer=KVAndScore(torch.full((physical_rows, 8), 9.0)), + token_pool, attn, indexer, c128, logical_rows = _token_pool( + unified=True, ring_size=ring_size ) - token_pool = object.__new__(DeepSeekV4TokenToKVPool) - token_pool._unified_kv = True - token_pool.compress_state_pools = [attn, c128] - token_pool.indexer_compress_state_pools = [indexer, None] - token_pool.get_ring_size = MagicMock(return_value=ring_size) - token_pool.clear_c4_req_states([1, 3]) selected = torch.tensor(list(range(8, 16)) + list(range(24, 32))) @@ -80,59 +94,68 @@ def test_clear_resets_only_selected_request_rings(self): self.assertTrue((state[logical_rows:] == 7).all()) self.assertTrue((c128.kv_score_buffer.kv_score == 9).all()) - def test_alloc_clears_new_slots_but_not_reused_slots(self): + def test_clear_is_noop_off_the_unified_path(self): + """The non-unified (fp8) pool addresses C4 state by SWA page, so a + req-slot reset must not touch it.""" + token_pool, attn, indexer, _, _ = _token_pool(unified=False) + + token_pool.clear_c4_req_states([1, 3]) + + for pool in (attn, indexer): + self.assertTrue((pool.kv_score_buffer.kv_score == 7).all()) + + def test_req_pool_hook_fires_for_new_slots_only(self): req_pool = ReqToTokenPool(3, 16, "cpu", enable_memory_saver=False) - token_pool = MagicMock() - # A bare MagicMock auto-creates a truthy `_unified_kv`, so the stub - # must declare it; the gate identity-compares against True. - token_pool._unified_kv = True + hook = MagicMock() + req_pool.register_on_alloc_rows(hook) reused = _request() # First admission: a brand-new slot, so its C4 ring must be cleared. - (reused_idx,) = alloc_req_slots( - req_pool, [reused], None, token_to_kv_pool=token_pool - ) - token_pool.clear_c4_req_states.assert_called_once_with([reused_idx]) + (reused_idx,) = alloc_req_slots(req_pool, [reused], None) + hook.assert_called_once_with([reused_idx]) # Chunked continuation reuses the same slot -- clearing it here would # wipe the state captured by the previous chunk. - token_pool.clear_c4_req_states.reset_mock() - reused.kv.req_pool_idx = reused_idx - reused.kv.kv_committed_len = 1 - reused.kv.kv_allocated_len = 1 - reused.kv.holds_kv = True - reused.inflight_middle_chunks = 1 - self.assertEqual( - alloc_req_slots(req_pool, [reused], None, token_to_kv_pool=token_pool), - [reused_idx], - ) - token_pool.clear_c4_req_states.assert_not_called() + hook.reset_mock() + _mark_reused(reused) + self.assertEqual(alloc_req_slots(req_pool, [reused], None), [reused_idx]) + hook.assert_not_called() - # Mixed batch: only the newly allocated slot is cleared. + # Mixed batch: only the newly allocated slot is reported. fresh = _request() - indices = alloc_req_slots( - req_pool, [reused, fresh], None, token_to_kv_pool=token_pool - ) + indices = alloc_req_slots(req_pool, [reused, fresh], None) self.assertEqual(indices[0], reused_idx) self.assertNotEqual(indices[1], reused_idx) - token_pool.clear_c4_req_states.assert_called_once_with([indices[1]]) - - def test_alloc_does_not_clear_c4_state_off_the_unified_path(self): - """The reset must not reach the non-unified (fp8) path. - - A duck-typed `hasattr(pool, "clear_c4_req_states")` check is not enough: - the attribute exists on every DeepSeekV4TokenToKVPool, unified or not. - """ - for unified in (False, None, 1, "yes"): - with self.subTest(unified_kv=unified): - # Fresh pool per subtest: each iteration consumes a req slot. - req_pool = ReqToTokenPool(1, 16, "cpu", enable_memory_saver=False) - token_pool = MagicMock() - token_pool._unified_kv = unified - alloc_req_slots( - req_pool, [_request()], None, token_to_kv_pool=token_pool - ) - token_pool.clear_c4_req_states.assert_not_called() + hook.assert_called_once_with([indices[1]]) + + def test_decode_req_pool_hook_fires_for_new_slots_only(self): + """PD decode pre-allocates through DecodeReqToTokenPool, which has its + own alloc; it must report fresh rows the same way.""" + req_pool = DecodeReqToTokenPool( + 2, 16, "cpu", enable_memory_saver=False, pre_alloc_size=2 + ) + hook = MagicMock() + req_pool.register_on_alloc_rows(hook) + + first = _request() + (first_idx,) = req_pool.alloc([first]) + hook.assert_called_once_with([first_idx]) + + hook.reset_mock() + _mark_reused(first) + second = _request() + indices = req_pool.alloc([first, second]) + self.assertEqual(indices[0], first_idx) + hook.assert_called_once_with([indices[1]]) + + hook.reset_mock() + self.assertEqual(req_pool.alloc([first]), [first_idx]) + hook.assert_not_called() + + def test_req_pool_without_hook_is_unchanged(self): + req_pool = ReqToTokenPool(2, 16, "cpu", enable_memory_saver=False) + (idx,) = alloc_req_slots(req_pool, [_request()], None) + self.assertGreater(idx, 0) if __name__ == "__main__": diff --git a/test/registered/unit/mem_cache/test_hisparse_allocator.py b/test/registered/unit/mem_cache/test_hisparse_allocator.py index ebc173a4ce59..7ea10c1f4f0b 100644 --- a/test/registered/unit/mem_cache/test_hisparse_allocator.py +++ b/test/registered/unit/mem_cache/test_hisparse_allocator.py @@ -131,7 +131,6 @@ def write(self, indices, values): queue = DecodePreallocQueue.__new__(DecodePreallocQueue) queue.req_to_token_pool = req_to_token_pool queue.token_to_kv_pool_allocator = allocator - queue.token_to_kv_pool = None queue.tree_cache = SimpleNamespace( evictable_size=MagicMock(return_value=0), protected_size=MagicMock(return_value=0), diff --git a/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py b/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py index 70208bfa24f4..2dcac0647635 100644 --- a/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py +++ b/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py @@ -36,9 +36,9 @@ def set_full_to_swa_mapping( return SimpleNamespace( page_size=page_size, - # alloc_extend branches on _unified; pin it to the hybrid-SWA path + # alloc_extend branches on _swa_req_ring; pin it to the paged-SWA path # (SimpleNamespace raises instead of defaulting if it goes missing). - _unified=False, + _swa_req_ring=False, full_attn_allocator=SimpleNamespace( available_size=lambda: full_available, alloc_extend=MagicMock(return_value=full_indices), From fbbb2bd2ec1f6afed9e371f268fa30f8897bbb60 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 7 Sep 2026 03:21:57 -0700 Subject: [PATCH 11/15] shared get_swa_ring_size helper --- .../srt/mem_cache/deepseek_v4_memory_pool.py | 17 +++++++++++------ .../srt/model_executor/pool_configurator.py | 9 +++------ 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index 804d2d5cae6d..eef4a6475dd1 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -63,6 +63,13 @@ def get_compress_state_write_pad(compress_ratio: int, ring_size: int) -> int: return ring_size - window_size + 2 if ring_size > window_size else 0 +def get_swa_ring_size(sliding_window: int, is_speculative: bool = False) -> int: + """Rows per request in the per-request SWA ring: the window plus room for + the draft tokens a verify batch writes ahead of the committed position.""" + spec_extra = (get_spec().speculative_num_draft_tokens - 1) if is_speculative else 0 + return sliding_window + spec_extra + + class DeepSeekV4SingleKVPool(KVCache): def __init__( self, @@ -635,10 +642,8 @@ def __init__( self.swa_kv_pool = None self.c4_kv_pool = None self.c128_kv_pool = None - spec_extra = ( - (get_spec().speculative_num_draft_tokens - 1) - if get_spec().speculative_algorithm is not None - else 0 + swa_ring_size = get_swa_ring_size( + self.sliding_window, get_spec().speculative_algorithm is not None ) self.unified_kv_pool = DeepSeekV4UnifiedKVPool( stage_ratios=stage_ratios, @@ -650,11 +655,11 @@ def __init__( device=device, memory_saver_adapter=self.memory_saver_adapter, custom_mem_pool=self.custom_mem_pool, - swa_ring_size=self.sliding_window + spec_extra, + swa_ring_size=swa_ring_size, ) self.unified_swa_window = self.sliding_window - self.unified_swa_ring_size = self.sliding_window + spec_extra + self.unified_swa_ring_size = swa_ring_size self.unified_swa_pages = self.unified_kv_pool.swa_pages self.swa_req_ring_size = self.unified_swa_ring_size else: diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index af148fc36222..3f7589beb13f 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -38,6 +38,7 @@ get_compress_state_ring_size, get_compress_state_write_pad, get_dsv4_indexer_bytes_per_token, + get_swa_ring_size, ) from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool from sglang.srt.runtime_context import ( @@ -911,12 +912,8 @@ def __init__(self, kvc: KVCacheConfigurator): self._unified = is_unified_kv_triton() self.attn_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim - # Mirror DeepSeekV4TokenToKVPool: swa_ring_size = sliding_window + - # (speculative_num_draft_tokens - 1). - spec_num_draft = get_spec().speculative_num_draft_tokens or 1 - self._swa_ring_size = self.swa_page_size + ( - (spec_num_draft - 1) if self.is_speculative else 0 - ) + # swa_page_size is the model's sliding window (cfg.window_size). + self._swa_ring_size = get_swa_ring_size(self.swa_page_size, self.is_speculative) self._spec_infl = 1.0 if self.is_speculative: From 0fa933449bbc8f8edc0da03d403e4560a3508784 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 7 Sep 2026 03:30:49 -0700 Subject: [PATCH 12/15] revert comment-only kernel hunks; ring mode asserts req pool instead of fallback --- .../kernels/jit/csrc/deepseek_v4/c_plan.cuh | 31 ++++++++++--------- python/sglang/srt/mem_cache/allocator/swa.py | 6 ++-- .../srt/mem_cache/kv_cache_configurator.py | 20 +++++------- 3 files changed, 28 insertions(+), 29 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh index c87f23e89f76..b157eebea28b 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh @@ -47,8 +47,8 @@ struct Prefill0Params { uint32_t num_q_tokens; int32_t compress_ratio; int32_t swa_page_size; - /// \brief Trailing tokens the write plan keeps resident in the compress state - /// ring; the bound is derived in `plan_compress_prefill`. + /// \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; }; @@ -213,8 +213,8 @@ __global__ __launch_bounds__(1024, 1) // } } } else { - // Path 2: general prefill (long extend_len). Iterate batches in an outer - // loop; the whole block sweeps each batch's tokens in parallel. + // Path 2: general prefill (long extend_len). Iterate batches in an outer loop; + // the whole block sweeps each batch's tokens in parallel. uint32_t base_e = 0; for (uint32_t batch_id = 0; batch_id < params.batch_size; ++batch_id) { const int32_t pl = s_prefix_len[batch_id]; @@ -511,20 +511,24 @@ inline PrefillPlan plan_compress_prefill( 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); - // Write pad: trailing tokens kept resident so a verify batch's committed tail - // survives any accept length. A write at `w` aliases onto `w - ring_size`, and - // the earliest position a future compression still needs is - // `prefix_len - window_size + 2` -- the next batch commits >= 1 token, and - // `run_prefill` launches compress before write, so a batch's own compressions - // read the pre-write ring. + // Write pad: trailing tokens kept resident so a verify batch's committed tail survives + // any accept length. Zero without speculation -- nothing rolls back, and the ring is + // then exactly one window wide. Otherwise the ring bounds it: a write at `w` aliases + // onto `w - ring_size`, and the earliest position a future compression still needs is + // `prefix_len - window_size + 2` (the next batch commits >= 1 token, and `run_prefill` + // launches the compress kernel before the write kernel, so a batch's own compressions + // read the pre-write ring). Padding past the extend range is harmless: the loops only + // span `[prefix_len, seq_len)`. const auto mtp_pad = ring_size > window_size ? ring_size - window_size + 2 : 0; const auto device = device_.unwrap(); const auto stream = LaunchKernel::resolve_device(device); if (cpu_or_gpu.unwrap().device_type == kDLGPU) { - // GPU input path for MTP / cuda-graph capture, where a host sync would be - // expensive: kernel0 builds the plan metadata, kernel_1 translates SWA locs. + // GPU input path: kernel0 builds the (CPU-loop-equivalent) plan metadata directly + // on device, padding to num_q_tokens with invalid; kernel_1 then finalizes the + // SWA-translated read/write locations. Used for MTP / cuda-graph capture where + // a host sync would be expensive. RuntimeCheck(batch_size <= kMaxPrefillBatchSize, "GPU plan only support batch size up to ", kMaxPrefillBatchSize); auto C = ffi::empty({num_q_tokens, sizeof(PlanC)}, kDLUInt8, device); auto W = ffi::empty({num_q_tokens, sizeof(PlanW)}, kDLUInt8, device); @@ -541,8 +545,7 @@ inline PrefillPlan plan_compress_prefill( .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. + // kernel_1 sees the already-padded buffers, so num_c == num_w == num_padded == num_q_tokens. const auto params1 = Prefill1Params{ .plan_c = static_cast(C.data_ptr()), .plan_w = static_cast(W.data_ptr()), diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index bb7368682014..e5b3690ddb6e 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -120,6 +120,9 @@ def __init__( self._swa_req_ring = ring_size is not None self._req_to_token_pool = req_to_token_pool if self._swa_req_ring: + assert req_to_token_pool is not None, ( + "per-request SWA ring: capacity is counted in req slots" + ) self._swa_ring_cost = ( (ring_size + self.page_size - 1) // self.page_size ) * self.page_size @@ -160,9 +163,6 @@ def full_available_size(self): def swa_available_size(self): if self._swa_req_ring: # Ring-based availability: free request slots * per-slot ring cost. - # Fall back to non-binding if the req pool wasn't wired in. - if self._req_to_token_pool is None: - return self.full_attn_allocator.available_size() return self._req_to_token_pool.available_size() * self._swa_ring_cost return self.swa_attn_allocator.available_size() diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 6dda4eb9091a..679bdfbc3f9c 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -338,20 +338,16 @@ def configure(self, *, pre_model_load_memory: int) -> KVCacheConfigResult: ) swa_max_total_num_tokens = sizes.swa_max_total_num_tokens - # Per-request SWA ring: swa_max_total_num_tokens was sized from the - # vestigial paged SWA pool; reconcile to the real ring capacity. alloc = pools.token_to_kv_pool_allocator if not self.is_draft_worker and is_swa_req_ring(alloc): - ring_capacity = int(alloc.swa_available_size()) - # Only reconcile downward: a value >= the current total means - # swa_available_size() hit its non-binding fallback. - if 0 < ring_capacity < swa_max_total_num_tokens: - logger.info( - "SWA ring: reconciling swa_max_total_num_tokens " - f"{swa_max_total_num_tokens} -> {ring_capacity} " - "(fixed per-request SWA ring capacity)." - ) - swa_max_total_num_tokens = ring_capacity + # Per-request SWA ring: the sizer's swa token count describes the + # vestigial paged pool; the real capacity is every req slot's ring. + swa_max_total_num_tokens = int(alloc.swa_available_size()) + logger.info( + "SWA ring: swa_max_total_num_tokens " + f"{sizes.swa_max_total_num_tokens} -> {swa_max_total_num_tokens} " + "(fixed per-request SWA ring capacity)." + ) logger.info( f"Memory pool end. " From 0e595f46ada6d96b38667b005a82d73efefbdb72 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 7 Sep 2026 03:47:33 -0700 Subject: [PATCH 13/15] hoist req_ring_state in compress_hip; drop is True on bool --- .../srt/layers/attention/dsv4/compress_hip.py | 26 ++++++++++--------- .../srt/layers/attention/dsv4/compressor.py | 2 +- .../layers/attention/dsv4/compressor_v2.py | 2 +- 3 files changed, 16 insertions(+), 14 deletions(-) diff --git a/python/sglang/srt/layers/attention/dsv4/compress_hip.py b/python/sglang/srt/layers/attention/dsv4/compress_hip.py index 1c05e92b4f0d..57847d9d70c3 100644 --- a/python/sglang/srt/layers/attention/dsv4/compress_hip.py +++ b/python/sglang/srt/layers/attention/dsv4/compress_hip.py @@ -118,6 +118,11 @@ def compress_extend_paged( assert isinstance(backend, DeepseekV4HipRadixBackend) token_to_kv_pool = backend.token_to_kv_pool assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool) + # C128 state is always request-addressed; C4 state is too under the + # unified-KV request ring. + 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,9 +149,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 or ( - self.ratio == 4 and token_to_kv_pool._unified_kv is True - ): + if req_ring_state: state_loc = state_pool.translate_from_req_position_to_state_loc( req_pool_indices[i], pre_state_indices ) @@ -168,9 +171,7 @@ def compress_extend_paged( post_state_len = post_state_indices.size(0) assert post_state_len <= valid_kv_len - if self.ratio == 128 or ( - self.ratio == 4 and token_to_kv_pool._unified_kv is True - ): + if req_ring_state: post_state_loc = state_pool.translate_from_req_position_to_state_loc( req_pool_indices[i], post_state_indices ) @@ -264,6 +265,11 @@ 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) + # C128 state is always request-addressed; C4 state is too under the + # unified-KV request ring. + 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 @@ -275,9 +281,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 or ( - self.ratio == 4 and token_to_kv_pool._unified_kv is True - ): + if req_ring_state: state_locs = state_pool.translate_from_req_position_to_state_loc( req_pool_indices, seq_lens - 1 ) @@ -292,9 +296,7 @@ def compress_decode_paged( -compress_bulk_len, 0, device=seq_lens.device ) compress_indices.clamp_(min=-1) - if self.ratio == 128 or ( - self.ratio == 4 and token_to_kv_pool._unified_kv is True - ): + if req_ring_state: compress_indices_state = ( state_pool.translate_from_req_position_to_state_loc( req_pool_indices[:, None], compress_indices diff --git a/python/sglang/srt/layers/attention/dsv4/compressor.py b/python/sglang/srt/layers/attention/dsv4/compressor.py index 684fe1131341..2e420f65367a 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor.py @@ -264,7 +264,7 @@ def create_paged_compressor_data( ) -> FusedCompressMetadata: swa_page_size = token_to_kv_pool.swa_page_size ring_size = token_to_kv_pool.get_ring_size(compress_ratio=compress_ratio) - use_req_ring = compress_ratio == 4 and token_to_kv_pool._unified_kv is True + use_req_ring = compress_ratio == 4 and token_to_kv_pool._unified_kv # assert ring_size % compress_ratio == 0 def clip_down(positions: torch.Tensor) -> torch.Tensor: diff --git a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py index 04f33ba45c47..c3515413119b 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py @@ -441,7 +441,7 @@ def create_paged_compressor_data( swa_page_size = token_to_kv_pool.swa_page_size ring_size = token_to_kv_pool.get_ring_size(compress_ratio=compress_ratio) - use_req_ring = compress_ratio == 4 and token_to_kv_pool._unified_kv is True + use_req_ring = compress_ratio == 4 and token_to_kv_pool._unified_kv # NOTE: This is actually a proxy, which encounter some bug with tvm-ffi. # As a workaround, we use `.detach()` to get the real tensor. full_to_swa = token_to_kv_pool.full_to_swa_index_mapping.detach() From c2e9b8dfb71dacef2933faaf5ba734fed7e87bdc Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 7 Sep 2026 03:56:04 -0700 Subject: [PATCH 14/15] trim comments --- .../sglang/kernels/ops/attention/dsv4/compress.py | 3 --- .../srt/layers/attention/dsv4/compress_hip.py | 4 ---- python/sglang/srt/managers/schedule_policy.py | 6 ++---- python/sglang/srt/mem_cache/allocator/swa.py | 10 ++-------- .../sglang/srt/mem_cache/base_swa_memory_pool.py | 5 ++--- .../srt/mem_cache/deepseek_v4_memory_pool.py | 10 +++------- .../sglang/srt/mem_cache/kv_cache_configurator.py | 5 ++--- python/sglang/srt/mem_cache/memory_pool.py | 2 -- .../srt/model_executor/pool_configurator.py | 15 +++++---------- .../test_swa_alloc_extend_page_estimation.py | 3 +-- 10 files changed, 17 insertions(+), 46 deletions(-) diff --git a/python/sglang/kernels/ops/attention/dsv4/compress.py b/python/sglang/kernels/ops/attention/dsv4/compress.py index 8dadcadc75bd..9650bb5228c8 100644 --- a/python/sglang/kernels/ops/attention/dsv4/compress.py +++ b/python/sglang/kernels/ops/attention/dsv4/compress.py @@ -179,8 +179,6 @@ def generate( int(swa_page_size), int(ring_size), ) - # The XPU plan builder has no use_req_ring parameter, so unified-KV - # request-ring addressing cannot be expressed there. assert not (_is_xpu and use_req_ring), ( "use_req_ring is not supported by the XPU compress plan builder" ) @@ -294,7 +292,6 @@ def generate( int(swa_page_size), int(ring_size), ) - # See plan_decode: XPU cannot express the unified-KV request ring. assert not (_is_xpu and use_req_ring), ( "use_req_ring is not supported by the XPU compress plan builder" ) diff --git a/python/sglang/srt/layers/attention/dsv4/compress_hip.py b/python/sglang/srt/layers/attention/dsv4/compress_hip.py index 57847d9d70c3..5c0addf33211 100644 --- a/python/sglang/srt/layers/attention/dsv4/compress_hip.py +++ b/python/sglang/srt/layers/attention/dsv4/compress_hip.py @@ -118,8 +118,6 @@ def compress_extend_paged( assert isinstance(backend, DeepseekV4HipRadixBackend) token_to_kv_pool = backend.token_to_kv_pool assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool) - # C128 state is always request-addressed; C4 state is too under the - # unified-KV request ring. req_ring_state = self.ratio == 128 or ( self.ratio == 4 and token_to_kv_pool._unified_kv ) @@ -265,8 +263,6 @@ 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) - # C128 state is always request-addressed; C4 state is too under the - # unified-KV request ring. req_ring_state = self.ratio == 128 or ( self.ratio == 4 and token_to_kv_pool._unified_kv ) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index c4a75f748ce0..e2864ca1fe2e 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -721,8 +721,7 @@ def _swa_budget_for_req( """ allocator = self.token_to_kv_pool_allocator if self._swa_req_ring: - # One fixed ring slot per request, independent of context or chunk - # length; pairs with the ring-based swa_available_size. + # One ring slot per request, in the same unit as swa_available_size. return allocator.swa_ring_cost_tokens if self.rem_chunk_tokens is not None: alloc = min(extend_input_len, self.rem_chunk_tokens) @@ -1018,8 +1017,7 @@ def add_chunked_req(self, req: Req): if self.is_hybrid_swa and not self._swa_req_ring: # alloc_extend needs extend_num_tokens + page_size per request, # so reserve one page here to avoid OOM. - # Unified-KV: rem_swa_tokens is ring capacity, not a per-chunk - # token budget; clamping against it would truncate the chunk. + # Ring mode skips it: rem_swa_tokens counts slots, not chunk tokens. _rem_tokens = min( _rem_tokens, int(self.rem_swa_tokens) - self.page_size ) diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index e5b3690ddb6e..e21b153a0b61 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -139,12 +139,10 @@ def __init__( @property def swa_req_ring(self) -> bool: - """SWA is a per-request ring; no per-token SWA budget applies.""" return self._swa_req_ring @property def swa_ring_cost_tokens(self) -> int: - """Ring mode: paged SWA cost of one request's ring slot (0 otherwise).""" return self._swa_ring_cost def available_size(self): @@ -247,8 +245,7 @@ def alloc_extend( return None if self._swa_req_ring: - # The unified SWA ring is slot-addressed, not paged here, so the - # vestigial paged allocator and full->swa mapping are skipped. + # Ring mode pages full KV only; full_to_swa_index_mapping stays unwritten. return self.full_attn_allocator.alloc_extend( prefix_lens, prefix_lens_cpu, @@ -310,8 +307,7 @@ def alloc_extend_swa_tail( return None if self._swa_req_ring: - # See alloc_extend. new_pages_available already ignored - # num_swa_pages, so the paged allocator has no capacity gate left. + # See alloc_extend: full KV only. return self.full_attn_allocator.alloc_extend( prefix_lens, prefix_lens_cpu, @@ -702,6 +698,4 @@ def clear(self): def is_swa_req_ring(allocator) -> bool: - """True when the allocator's SWA side is a per-request ring (see - BaseSWAKVPool.swa_req_ring_size), so SWA carries no per-token budget.""" return isinstance(allocator, SWATokenToKVPoolAllocator) and allocator.swa_req_ring diff --git a/python/sglang/srt/mem_cache/base_swa_memory_pool.py b/python/sglang/srt/mem_cache/base_swa_memory_pool.py index 24a61da749f8..2b44bf942e0e 100644 --- a/python/sglang/srt/mem_cache/base_swa_memory_pool.py +++ b/python/sglang/srt/mem_cache/base_swa_memory_pool.py @@ -15,9 +15,8 @@ class BaseSWAKVPool(KVCache): """ swa_kv_pool: KVCache - # Set when SWA KV is a fixed per-request ring of this many tokens, addressed - # by req_pool_idx, instead of a paged token pool. The paged SWA allocator is - # then vestigial and SWA must not be budgeted per token. + # Set when SWA KV is a per-request ring of this many tokens (addressed by + # req_pool_idx) rather than a paged token pool; SWA is then not budgeted per token. swa_req_ring_size: Optional[int] = None @abc.abstractmethod diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index eef4a6475dd1..35e3930cadac 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -64,8 +64,7 @@ def get_compress_state_write_pad(compress_ratio: int, ring_size: int) -> int: def get_swa_ring_size(sliding_window: int, is_speculative: bool = False) -> int: - """Rows per request in the per-request SWA ring: the window plus room for - the draft tokens a verify batch writes ahead of the committed position.""" + # A verify batch writes its draft tokens ahead of the committed position. spec_extra = (get_spec().speculative_num_draft_tokens - 1) if is_speculative else 0 return sliding_window + spec_extra @@ -1063,8 +1062,6 @@ def get_online_c128_mtp_pending_seq_lens(self) -> torch.Tensor: return self.online_c128_mtp_pending_seq_lens def clear_c4_req_states(self, req_pool_indices: Sequence[int]) -> None: - """Reset the request-owned C4 rows; the sentinel/ring padding that - CompressStatePool allocates is left intact.""" if not self._unified_kv or not req_pool_indices: return @@ -1117,9 +1114,8 @@ def clear_unaccepted_c128_draft_states( num_draft_tokens: int, ) -> None: """Clear offline C128 ring slots written for rejected speculative tokens. - C4 needs none: its draft states are overwritten in position order before - they can be read, while a C128 compression boundary can consume a - previously written draft slot directly.""" + C4 needs no counterpart: its draft states are overwritten in position order + before any read; a C128 compression boundary can read a stale draft slot.""" if ONLINE_C128 or num_draft_tokens <= 1 or req_pool_indices.numel() == 0: return diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 679bdfbc3f9c..9e61e544c980 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -1361,9 +1361,8 @@ def _build_dsv4_kv_pool( online_mtp_max_draft_tokens=(max_speculative_num_draft_tokens() or 0), ) if not self.is_draft_worker and token_to_kv_pool._unified_kv: - # Unified-KV C4 state is a per-request ring: reset a row's ring - # whenever its req slot is handed out again. The draft pool has no - # C4 layers and shares this req pool, so only the target registers. + # The draft pool has no C4 layers and shares this req pool, so only + # the target registers the per-slot C4 reset. req_to_token_pool.register_on_alloc_rows( token_to_kv_pool.clear_c4_req_states ) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 51f8ced084c4..14db7141d5c8 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -352,8 +352,6 @@ def attach_aux_cache(self, aux_cache: Any) -> None: self._aux_cache = aux_cache def register_on_alloc_rows(self, hook: Callable[[List[int]], None]) -> None: - """Receive the row indices of every fresh allocation. Per-request - state pools use it to reset a row's state when the slot is reused.""" assert self._on_alloc_rows is None self._on_alloc_rows = hook diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 3f7589beb13f..b8bb4facc81f 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -1012,8 +1012,7 @@ def _get_bytes_per_full_token(self) -> float: c4_frac = 1 / (4 * self.c4_shrink_factor) return ( - # Unified_kv: SWA is a fixed per-request ring (see _fixed_swa_bytes), - # not a per-token pool, so it is excluded from the per-token coeff. + # Ring mode: SWA is a fixed per-request pool (see _fixed_swa_bytes). ( 0.0 if self._unified @@ -1022,8 +1021,7 @@ def _get_bytes_per_full_token(self) -> float: + c4_frac * kv_bytes * self.num_layers_ca4 + 1 / 128 * kv_bytes * self.num_layers_ca128 + 1 / 4 * self.indexer_bytes_per_token * self.num_layers_ca4 - # Unified_kv: c4 state is addressed off the SWA slot, and that pool - # is a fixed per-request ring, so the state ring is request-scoped. + # Ring mode: C4 state is per-request too (see _fixed_c4_state_bytes). + ( 0.0 if self._unified @@ -1160,8 +1158,7 @@ def finalize_with_max_running_requests( config.c128_state_pool_size = num_req_slots else: config.c128_state_pool_size = num_req_slots * self.c128_ring_size - # Unified_kv: request-scoped, so it is sized here from the now-known - # concurrency rather than from full_token in _compute_dsv4_sizes. + # Ring mode: C4 state is request-scoped, so size it from the known concurrency. if self._unified and self.num_layers_ca4 > 0: config.c4_state_pool_size = self._unified_c4_state_pool_size( config.max_running_requests @@ -1210,10 +1207,8 @@ def calculate_pool_sizes( def calculate_pool_sizes_from_max_tokens( self, max_total_num_tokens: int, page_size: int ) -> MemoryPoolConfig: - """Caller contract: max_total_num_tokens must not exceed what - calculate_pool_sizes derived from the same budget (config_from_budget - asserts it). Subtracting the fixed-pool bias again here would - double-count -- the input is a token count, not a byte budget.""" + # Token count, not a byte budget: the fixed pools are not re-subtracted, so + # the input must not exceed what calculate_pool_sizes derived for it. assert page_size % 128 == 0, ( "page_size must be multiple of 128 for compressed attention" ) diff --git a/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py b/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py index 2dcac0647635..712d26f5e667 100644 --- a/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py +++ b/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py @@ -36,8 +36,7 @@ def set_full_to_swa_mapping( return SimpleNamespace( page_size=page_size, - # alloc_extend branches on _swa_req_ring; pin it to the paged-SWA path - # (SimpleNamespace raises instead of defaulting if it goes missing). + # alloc_extend reads _swa_req_ring; pin the paged-SWA path. _swa_req_ring=False, full_attn_allocator=SimpleNamespace( available_size=lambda: full_available, From 8e005cf4106f162790bd10712df2ca2fd83fb7e9 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 7 Sep 2026 04:15:23 -0700 Subject: [PATCH 15/15] ring mode: size_swa is the ring total; PD swa budget ignores phantom evictable --- python/sglang/srt/disaggregation/decode.py | 9 ++++++++- python/sglang/srt/mem_cache/allocator/swa.py | 8 ++++++-- python/sglang/srt/mem_cache/kv_cache_configurator.py | 4 ++-- python/sglang/srt/model_executor/pool_configurator.py | 4 +++- 4 files changed, 19 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 24daaecdebbe..1e18e91451aa 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -73,6 +73,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, @@ -1715,7 +1716,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) diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index e21b153a0b61..13f57eb118ea 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -126,10 +126,12 @@ def __init__( self._swa_ring_cost = ( (ring_size + self.page_size - 1) // self.page_size ) * self.page_size + # Total SWA capacity is every req slot's ring; all slots are free here. + self._size_swa = req_to_token_pool.available_size() * self._swa_ring_cost logger.info( "SWA per-request ring accounting enabled: " f"ring_size={ring_size}, ring_cost_tokens={self._swa_ring_cost}, " - f"paged size_swa={self._size_swa} (bypassed)" + f"size_swa={self._size_swa} (paged size_swa={size_swa} bypassed)" ) else: self._swa_ring_cost = 0 @@ -532,7 +534,9 @@ def resize(self, config) -> None: size_full = int(config.full_max_total_num_tokens) size_swa = int(config.swa_max_total_num_tokens) self._size_full = size_full - self._size_swa = size_swa + if not self._swa_req_ring: + # Ring capacity follows the req slot count, not the token config. + self._size_swa = size_swa for alloc, sz in ( (self.full_attn_allocator, size_full), (self.swa_attn_allocator, size_swa), diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 9e61e544c980..9aeefec84dfe 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -341,8 +341,8 @@ def configure(self, *, pre_model_load_memory: int) -> KVCacheConfigResult: alloc = pools.token_to_kv_pool_allocator if not self.is_draft_worker and is_swa_req_ring(alloc): # Per-request SWA ring: the sizer's swa token count describes the - # vestigial paged pool; the real capacity is every req slot's ring. - swa_max_total_num_tokens = int(alloc.swa_available_size()) + # vestigial paged pool; the allocator knows the real ring total. + swa_max_total_num_tokens = alloc.size_swa logger.info( "SWA ring: swa_max_total_num_tokens " f"{sizes.swa_max_total_num_tokens} -> {swa_max_total_num_tokens} " diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index b8bb4facc81f..513b801f0c86 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -1044,7 +1044,9 @@ def _get_bytes_per_full_token(self) -> float: def _compute_dsv4_sizes(self, full_token: int, page_size: int) -> _DSV4PoolSizes: full_token = full_token // page_size * page_size swa_tokens = int(full_token * self.swa_ratio) // page_size * page_size - self.validate_swa_pool_size(swa_tokens, self.sliding_window_size, page_size) + if not self._unified: + # Ring mode: the paged SWA pool is vestigial, so its floor does not apply. + self.validate_swa_pool_size(swa_tokens, self.sliding_window_size, page_size) return _DSV4PoolSizes( full_max_total_num_tokens=full_token, swa_max_total_num_tokens=swa_tokens,