Optimize C128 state pool allocation using request state pool (#28612)

python/sglang/jit_kernel/csrc/deepseek_v4/c128_online_v2.cuh

@@ -533,9 +533,7 @@ struct OnlineDecodePlanParams {
   const int64_t* __restrict__ seq_lens;
   const int64_t* __restrict__ req_pool_indices;
   const int32_t* __restrict__ req_to_token;
-  const int64_t* __restrict__ full_to_swa;  // (full_cache_size,) int64
   int64_t stride_r2t;
-  int32_t swa_page_size;
   int32_t state_slot_offset;
   uint32_t batch_size;
 };
@@ -545,10 +543,7 @@ __global__ void plan_c128_online_decode_kernel(const OnlineDecodePlanParams para
   if (idx >= params.batch_size) return;
   const auto seq_len = static_cast<uint32_t>(params.seq_lens[idx]);
   const auto rid = params.req_pool_indices[idx];
-  const int32_t chunk_start = static_cast<int32_t>((seq_len - 1u) / 128u * 128u);
-  const int32_t full_loc = params.req_to_token[rid * params.stride_r2t + chunk_start];
-  const int32_t swa_loc = static_cast<int32_t>(params.full_to_swa[full_loc]);
-  const int32_t slot = swa_loc / params.swa_page_size + params.state_slot_offset;
+  const int32_t slot = static_cast<int32_t>(rid) + params.state_slot_offset;
   params.plan_d[idx] = DecodePlan{
       .seq_len = seq_len,
       .write_loc = slot,
@@ -566,9 +561,7 @@ inline void plan_online_decode(
     const tvm::ffi::TensorView seq_lens,
     const tvm::ffi::TensorView req_pool_indices,
     const tvm::ffi::TensorView req_to_token,
-    const tvm::ffi::TensorView full_to_swa,
     const tvm::ffi::TensorView plan_d_dev_,
-    const int32_t swa_page_size,
     const int32_t state_slot_offset) {
   auto B = SymbolicSize{"batch_size"};
   auto device_ = SymbolicDevice{};
@@ -587,15 +580,10 @@ inline void plan_online_decode(
       .with_dtype<int32_t>()
       .with_device(device_)
       .verify(req_to_token);
-  TensorMatcher({-1})  //
-      .with_dtype<int64_t>()
-      .with_device(device_)
-      .verify(full_to_swa);
   TensorMatcher({B, sizeof(DecodePlan)})  //
       .with_dtype<uint8_t>()
       .with_device(device_)
       .verify(plan_d_dev_);
-  RuntimeCheck(swa_page_size > 0);
   RuntimeCheck(state_slot_offset >= 0);
 
   const auto batch_size = static_cast<uint32_t>(B.unwrap());
@@ -610,9 +598,7 @@ inline void plan_online_decode(
       .seq_lens = static_cast<const int64_t*>(seq_lens.data_ptr()),
       .req_pool_indices = static_cast<const int64_t*>(req_pool_indices.data_ptr()),
       .req_to_token = static_cast<const int32_t*>(req_to_token.data_ptr()),
-      .full_to_swa = static_cast<const int64_t*>(full_to_swa.data_ptr()),
       .stride_r2t = stride_r2t,
-      .swa_page_size = swa_page_size,
       .state_slot_offset = state_slot_offset,
       .batch_size = batch_size,
   };
@@ -685,9 +671,7 @@ struct OnlinePrefillStage1Params {
   CompressPlan* __restrict__ plan_w;
   const int64_t* __restrict__ req_pool_indices;  // (batch_size,)
   const int32_t* __restrict__ req_to_token;      // (num_reqs, max_tokens)
-  const int64_t* __restrict__ full_to_swa;       // (full_cache_size,)
   int64_t stride_r2t;
-  int32_t swa_page_size;
   int32_t state_slot_offset;
   uint32_t num_c;
   uint32_t num_w;
@@ -704,11 +688,7 @@ __global__ void plan_c128_online_prefill_kernel(const OnlinePrefillStage1Params
   if (plan.is_invalid()) return;
   const auto batch_id = plan.read_page_0;
   const auto rid = params.req_pool_indices[batch_id];
-  const int32_t position = static_cast<int32_t>(plan.seq_len - 1u);
-  const int32_t chunk_start = (position / 128) * 128;
-  const int32_t full_loc = params.req_to_token[rid * params.stride_r2t + chunk_start];
-  const int32_t swa_loc = static_cast<int32_t>(params.full_to_swa[full_loc]);
-  const int32_t main_slot = swa_loc / params.swa_page_size;
+  const int32_t main_slot = static_cast<int32_t>(rid);
   plan.read_page_0 = main_slot + params.state_slot_offset;
   plan.read_page_1 = main_slot;
   *plan_ptr = plan;
@@ -721,12 +701,10 @@ inline OnlinePrefillPlan plan_online_prefill(
     const tvm::ffi::TensorView extend_lens,
     const tvm::ffi::TensorView req_pool_indices,
     const tvm::ffi::TensorView req_to_token,
-    const tvm::ffi::TensorView full_to_swa,
     const tvm::ffi::TensorView plan_c_pin,
     const tvm::ffi::TensorView plan_w_pin,
     const tvm::ffi::TensorView plan_c_dev_,
     const tvm::ffi::TensorView plan_w_dev_,
-    const int32_t swa_page_size,
     const int32_t state_slot_offset,
     const bool use_cuda_graph) {
   auto B = SymbolicSize{"batch_size"};
@@ -749,10 +727,6 @@ inline OnlinePrefillPlan plan_online_prefill(
       .with_dtype<int32_t>()
       .with_device(device_)
       .verify(req_to_token);
-  TensorMatcher({-1})  //
-      .with_dtype<int64_t>()
-      .with_device(device_)
-      .verify(full_to_swa);
   TensorMatcher({N, sizeof(CompressPlan)})  //
       .with_dtype<uint8_t>()
       .with_device(cpu)
@@ -878,9 +852,7 @@ inline OnlinePrefillPlan plan_online_prefill(
         .plan_w = plan_w_dev_ptr,
         .req_pool_indices = static_cast<const int64_t*>(req_pool_indices.data_ptr()),
         .req_to_token = static_cast<const int32_t*>(req_to_token.data_ptr()),
-        .full_to_swa = static_cast<const int64_t*>(full_to_swa.data_ptr()),
         .stride_r2t = req_to_token.stride(0),
-        .swa_page_size = swa_page_size,
         .state_slot_offset = state_slot_offset,
         .num_c = num_c_padded,
         .num_w = num_w_padded,

python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh

@@ -53,7 +53,7 @@ struct Prefill1Params {
   PlanW* plan_w;
   const RID_T* rid_ptr;  // [batch_size]
   const R2T_T* r2t_ptr;  // [num_reqs, stride_r2t]
-  const F2S_T* f2s_ptr;  // [num_swa_slots]
+  const F2S_T* f2s_ptr;  // [num_full_slots], full_loc -> swa_loc
   int64_t stride_r2t;
   uint32_t num_c;
   uint32_t num_w;
@@ -69,7 +69,7 @@ struct DecodeParams {
   PlanD* plan_d;
   const RID_T* rid_ptr;  // [batch_size]
   const R2T_T* r2t_ptr;  // [num_reqs, stride_r2t]
-  const F2S_T* f2s_ptr;  // [num_swa_slots]
+  const F2S_T* f2s_ptr;  // [num_full_slots], full_loc -> swa_loc
   const IDX_T* seq_ptr;  // [batch_size]
   int64_t stride_r2t;
   uint32_t batch_size;
@@ -297,6 +297,9 @@ __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) {
+    return static_cast<int32_t>(rid * params.ring_size + position % params.ring_size);
+  };
 
   if (!plan_c.is_invalid()) {  // 1. in bound. 2. not masked
     if (plan_c.buffer_len > 0) {
@@ -307,12 +310,17 @@ __global__ void plan_compress_prefill_kernel_1(const Prefill1Params params) {
       const auto position_1 = static_cast<int32_t>(plan_c.seq_len - 1);
       // only used for c4, harmless for c128
       const auto position_0 = max(position_1 - params.compress_ratio, 0);
-      const auto raw_loc_0 = mapping[position_0];
-      const auto raw_loc_1 = mapping[position_1];
-      const auto swa_loc_0 = params.f2s_ptr[raw_loc_0];
-      const auto swa_loc_1 = params.f2s_ptr[raw_loc_1];
-      plan_c.read_page_0 = compute_loc(swa_loc_0) / params.compress_ratio;
-      plan_c.read_page_1 = compute_loc(swa_loc_1) / params.compress_ratio;
+      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;
+      } else {
+        const auto raw_loc_0 = mapping[position_0];
+        const auto raw_loc_1 = mapping[position_1];
+        const auto state_loc_0 = params.f2s_ptr[raw_loc_0];
+        const auto state_loc_1 = params.f2s_ptr[raw_loc_1];
+        plan_c.read_page_0 = compute_loc(state_loc_0) / params.compress_ratio;
+        plan_c.read_page_1 = compute_loc(state_loc_1) / params.compress_ratio;
+      }
       params.plan_c[idx] = plan_c;
     }
   } else if (idx < params.num_c_padded) {
@@ -325,10 +333,13 @@ __global__ void plan_compress_prefill_kernel_1(const Prefill1Params params) {
     const auto mapping = params.r2t_ptr + rid * params.stride_r2t;
     // `seq_len` (`write_loc`) may not be aligned here
     const auto position = static_cast<int32_t>(plan_w.write_loc - 1);
-    const auto raw_loc = mapping[position];
-    const auto swa_loc = params.f2s_ptr[raw_loc];
     plan_w.ragged_id = ragged_id;
-    plan_w.write_loc = compute_loc(swa_loc);
+    if (params.compress_ratio == 128) {
+      plan_w.write_loc = compute_c128_loc(rid, position);
+    } else {
+      const auto raw_loc = mapping[position];
+      plan_w.write_loc = compute_loc(params.f2s_ptr[raw_loc]);
+    }
     params.plan_w[idx] = plan_w;
   } else if (idx < params.num_w_padded) {
     params.plan_w[idx] = PlanW::invalid();
@@ -345,16 +356,28 @@ __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) {
+    return static_cast<int32_t>(rid * params.ring_size + position % params.ring_size);
+  };
   const auto seq_len = static_cast<int32_t>(params.seq_ptr[idx]);
   const auto position_1 = static_cast<int32_t>(seq_len - 1);
   const auto position_0 = max(position_1 - params.compress_ratio, 0);
-  const auto raw_loc_0 = mapping[position_0];
-  const auto raw_loc_1 = mapping[position_1];
-  const auto swa_loc_0 = params.f2s_ptr[raw_loc_0];
-  const auto swa_loc_1 = params.f2s_ptr[raw_loc_1];
-  const auto write_loc = compute_loc(swa_loc_1);
-  const auto read_page_0 = compute_loc(swa_loc_0) / params.compress_ratio;
-  const auto read_page_1 = write_loc / params.compress_ratio;
+  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;
+  } else {
+    const auto raw_loc_0 = mapping[position_0];
+    const auto raw_loc_1 = mapping[position_1];
+    const auto state_loc_0 = params.f2s_ptr[raw_loc_0];
+    const auto state_loc_1 = params.f2s_ptr[raw_loc_1];
+    write_loc = static_cast<int32_t>(compute_loc(state_loc_1));
+    read_page_0 = static_cast<int32_t>(compute_loc(state_loc_0) / params.compress_ratio);
+    read_page_1 = static_cast<int32_t>(write_loc / params.compress_ratio);
+  }
   params.plan_d[idx] = {
       .seq_len = static_cast<uint32_t>(seq_len),
       .write_loc = write_loc,
@@ -425,9 +448,9 @@ __global__ void plan_compress_decode_legacy_kernel(const DecodeParamsLegacy para
   const auto seq_len = static_cast<int32_t>(params.seq_ptr[idx]);
   const auto position_1 = seq_len - 1;
   const auto position_0 = max(position_1 - params.compress_ratio, 0);
-  const auto write_loc = legacy_compute_loc(rid, position_1);
-  const auto read_page_0 = legacy_compute_page(rid, position_0);
-  const auto read_page_1 = legacy_compute_page(rid, position_1);
+  const int32_t write_loc = legacy_compute_loc(rid, position_1);
+  const int32_t read_page_0 = legacy_compute_page(rid, position_0);
+  const int32_t read_page_1 = legacy_compute_page(rid, position_1);
   params.plan_d[idx] = {
       .seq_len = static_cast<uint32_t>(seq_len),
       .write_loc = write_loc,
@@ -443,7 +466,9 @@ using PrefillPlan = tvm::ffi::Tuple<tvm::ffi::Tensor, tvm::ffi::Tensor>;
  * Inputs (all CPU-resident):
  * @param req_pool_indices  `[batch_size]` int64_t
  * @param req_to_token      `[num_reqs, max_tokens_per_req]` int64_t
- * @param full_to_swa       `[num_swa_slots]` int64_t
+ * @param full_to_state     `[full_cache_size]` int64_t. For c4 this maps
+ *                          full loc -> SWA loc; ignored for c128, whose
+ *                          state slot is request-scoped.
  * @param seq_lens          `[batch_size]` int64
  * @param extend_lens       `[batch_size]` int64
  * @param compress_plan     `[num_q_tokens, 16]` uint8 (output)
@@ -455,7 +480,7 @@ using PrefillPlan = tvm::ffi::Tuple<tvm::ffi::Tensor, tvm::ffi::Tensor>;
 inline PrefillPlan plan_compress_prefill(
     const tvm::ffi::TensorView req_pool_indices,  // GPU
     const tvm::ffi::TensorView req_to_token,      // GPU
-    const tvm::ffi::TensorView full_to_swa,       // GPU
+    const tvm::ffi::TensorView full_to_state,     // GPU
     const tvm::ffi::TensorView seq_lens,          // CPU/GPU
     const tvm::ffi::TensorView extend_lens,       // CPU/GPU
     const tvm::ffi::TensorView pin_buffer,        // CPU
@@ -482,7 +507,7 @@ inline PrefillPlan plan_compress_prefill(
   TensorMatcher({-1})  //
       .with_dtype<F2S_T>()
       .with_device(device_)
-      .verify(full_to_swa);
+      .verify(full_to_state);
   TensorMatcher({B})  //
       .with_dtype<IDX_T>()
       .with_device(cpu_or_gpu)
@@ -500,7 +525,7 @@ inline PrefillPlan plan_compress_prefill(
   const auto ext_ptr = static_cast<const IDX_T*>(extend_lens.data_ptr());
   const auto rid_ptr = static_cast<const RID_T*>(req_pool_indices.data_ptr());
   const auto r2t_ptr = static_cast<const R2T_T*>(req_to_token.data_ptr());
-  const auto f2s_ptr = static_cast<const F2S_T*>(full_to_swa.data_ptr());
+  const auto f2s_ptr = static_cast<const F2S_T*>(full_to_state.data_ptr());
 
   const auto batch_size = static_cast<uint32_t>(B.unwrap());
   constexpr auto kMaxTokens = static_cast<uint32_t>(std::numeric_limits<uint16_t>::max());
@@ -637,7 +662,7 @@ inline PrefillPlan plan_compress_prefill(
 inline tvm::ffi::Tensor plan_compress_decode(
     const tvm::ffi::TensorView req_pool_indices,  // GPU
     const tvm::ffi::TensorView req_to_token,      // GPU
-    const tvm::ffi::TensorView full_to_swa,       // GPU
+    const tvm::ffi::TensorView full_to_state,     // GPU
     const tvm::ffi::TensorView seq_lens,          // CPU/GPU
     const int32_t compress_ratio,
     const int32_t swa_page_size,
@@ -657,7 +682,7 @@ inline tvm::ffi::Tensor plan_compress_decode(
   TensorMatcher({-1})  //
       .with_dtype<F2S_T>()
       .with_device(device_)
-      .verify(full_to_swa);
+      .verify(full_to_state);
   TensorMatcher({B})  //
       .with_dtype<IDX_T>()
       .with_device(device_)
@@ -670,7 +695,7 @@ inline tvm::ffi::Tensor plan_compress_decode(
       .plan_d = static_cast<PlanD*>(D.data_ptr()),
       .rid_ptr = static_cast<const RID_T*>(req_pool_indices.data_ptr()),
       .r2t_ptr = static_cast<const R2T_T*>(req_to_token.data_ptr()),
-      .f2s_ptr = static_cast<const F2S_T*>(full_to_swa.data_ptr()),
+      .f2s_ptr = static_cast<const F2S_T*>(full_to_state.data_ptr()),
       .seq_ptr = static_cast<const IDX_T*>(seq_lens.data_ptr()),
       .stride_r2t = req_to_token.size(1),
       .batch_size = batch_size,

python/sglang/jit_kernel/csrc/deepseek_v4/online_c128_mtp.cuh

@@ -23,7 +23,6 @@ struct OnlineC128MTPWritePrefixParams {
   const TSeq* __restrict__ seq_lens;
   const TReq* __restrict__ req_pool_indices;
   const int32_t* __restrict__ req_to_token;
-  const int64_t* __restrict__ full_to_swa;
   const float* __restrict__ ape;
   float* __restrict__ state;
   int64_t kv_score_stride_b;
@@ -31,7 +30,6 @@ struct OnlineC128MTPWritePrefixParams {
   int64_t ape_stride_r;
   int64_t state_stride_b;
   int64_t layer_bs;
-  int64_t swa_page_size;
   int64_t num_verify_tokens;
   int64_t state_slot_stride;
 };
@@ -50,13 +48,11 @@ struct OnlineC128MTPCommitPendingParams {
   const TSeq* __restrict__ cur_seq_lens;
   const TReq* __restrict__ cur_req_pool_indices;
   const int32_t* __restrict__ req_to_token;
-  const int64_t* __restrict__ full_to_swa;
   const int64_t* __restrict__ pending_seq_lens;
   float* __restrict__ state;
   int64_t cur_bs;
   int64_t req_to_token_stride_b;
   int64_t state_stride_b;
-  int64_t swa_page_size;
   int64_t num_verify_tokens;
   int64_t state_slot_stride;
   int64_t max_num_reqs;
@@ -94,10 +90,7 @@ __global__ void online_c128_mtp_commit_pending_kernel(const OnlineC128MTPCommitP
   const int64_t final_seq = old_seq + accept;
   if ((final_seq & 127) == 0) return;
 
-  const int64_t chunk_start = ((final_seq - 1) / 128) * 128;
-  const int64_t full_loc = static_cast<int64_t>(params.req_to_token[req * params.req_to_token_stride_b + chunk_start]);
-  const int64_t swa_loc = params.full_to_swa[full_loc];
-  const int64_t slot = swa_loc / params.swa_page_size;
+  const int64_t slot = req;
   const float* const src = params.state + (slot + accept * params.state_slot_stride) * params.state_stride_b;
   float* const dst = params.state + slot * params.state_stride_b;
 
@@ -118,11 +111,7 @@ __global__ void online_c128_mtp_write_prefix_kernel(const OnlineC128MTPWritePref
 
   int64_t init_slot = 0;
   if (has_partial) {
-    const int64_t chunk_start = ((seq_before - 1) / 128) * 128;
-    const int64_t full_loc =
-        static_cast<int64_t>(params.req_to_token[req_idx * params.req_to_token_stride_b + chunk_start]);
-    const int64_t swa_loc = params.full_to_swa[full_loc];
-    init_slot = swa_loc / params.swa_page_size;
+    init_slot = req_idx;
   }
 
   const int64_t d = static_cast<int64_t>(threadIdx.x);
@@ -173,11 +162,7 @@ __global__ void online_c128_mtp_write_prefix_kernel(const OnlineC128MTPWritePref
 
     const int64_t final_seq = seq_before + step + 1;
     if ((final_seq & 127) != 0) {
-      const int64_t chunk_start = ((final_seq - 1) / 128) * 128;
-      const int64_t full_loc =
-          static_cast<int64_t>(params.req_to_token[req_idx * params.req_to_token_stride_b + chunk_start]);
-      const int64_t swa_loc = params.full_to_swa[full_loc];
-      const int64_t slot = swa_loc / params.swa_page_size + (step + 1) * params.state_slot_stride;
+      const int64_t slot = req_idx + (step + 1) * params.state_slot_stride;
       float* const out = params.state + slot * params.state_stride_b;
       out[d] = run_max;
       out[kHeadDim + d] = run_sum;
@@ -192,19 +177,16 @@ __global__ void online_c128_mtp_write_prefix_kernel(const OnlineC128MTPWritePref
   }
 }
 
-template <int64_t kHeadDim>
+template <int64_t kHeadDim, typename TSeq, typename TReq>
 struct OnlineC128MTPWritePrefixKernel {
-  template <typename TSeq, typename TReq>
   static void launch(
       tvm::ffi::TensorView kv_score_input,
       tvm::ffi::TensorView seq_lens,
       tvm::ffi::TensorView req_pool_indices,
       tvm::ffi::TensorView req_to_token,
-      tvm::ffi::TensorView full_to_swa,
       tvm::ffi::TensorView ape,
       tvm::ffi::TensorView state,
       int64_t layer_bs,
-      int64_t swa_page_size,
       int64_t num_verify_tokens,
       int64_t state_slot_stride,
       DLDevice device) {
@@ -215,7 +197,6 @@ struct OnlineC128MTPWritePrefixKernel {
         .seq_lens = static_cast<const TSeq*>(seq_lens.data_ptr()),
         .req_pool_indices = static_cast<const TReq*>(req_pool_indices.data_ptr()),
         .req_to_token = static_cast<const int32_t*>(req_to_token.data_ptr()),
-        .full_to_swa = static_cast<const int64_t*>(full_to_swa.data_ptr()),
         .ape = static_cast<const float*>(ape.data_ptr()),
         .state = static_cast<float*>(state.data_ptr()),
         .kv_score_stride_b = kv_score_input.stride(0),
@@ -223,7 +204,6 @@ struct OnlineC128MTPWritePrefixKernel {
         .ape_stride_r = ape.stride(0),
         .state_stride_b = state.stride(0),
         .layer_bs = layer_bs,
-        .swa_page_size = swa_page_size,
         .num_verify_tokens = num_verify_tokens,
         .state_slot_stride = state_slot_stride,
     };
@@ -239,25 +219,20 @@ struct OnlineC128MTPWritePrefixKernel {
       tvm::ffi::TensorView seq_lens,
       tvm::ffi::TensorView req_pool_indices,
       tvm::ffi::TensorView req_to_token,
-      tvm::ffi::TensorView full_to_swa,
       tvm::ffi::TensorView ape,
       tvm::ffi::TensorView state,
       int64_t layer_bs,
-      int64_t swa_page_size,
       int64_t num_verify_tokens,
       int64_t state_slot_stride) {
     using namespace host;
 
-    auto seq_dtype = SymbolicDType{};
-    auto req_dtype = SymbolicDType{};
     auto device = SymbolicDevice{};
     device.set_options<kDLCUDA>();
 
     TensorMatcher({-1, kHeadDim * 2}).with_dtype<float>().with_device(device).verify(kv_score_input);
-    TensorMatcher({-1}).with_dtype<int32_t, int64_t>(seq_dtype).with_device(device).verify(seq_lens);
-    TensorMatcher({-1}).with_dtype<int32_t, int64_t>(req_dtype).with_device(device).verify(req_pool_indices);
+    TensorMatcher({-1}).with_dtype<TSeq>().with_device(device).verify(seq_lens);
+    TensorMatcher({-1}).with_dtype<TReq>().with_device(device).verify(req_pool_indices);
     TensorMatcher({-1, -1}).with_dtype<int32_t>().with_device(device).verify(req_to_token);
-    TensorMatcher({-1}).with_dtype<int64_t>().with_device(device).verify(full_to_swa);
     TensorMatcher({128, kHeadDim}).with_dtype<float>().with_device(device).verify(ape);
     TensorMatcher({-1, kHeadDim * 3}).with_dtype<float>().with_device(device).verify(state);
 
@@ -268,73 +243,22 @@ struct OnlineC128MTPWritePrefixKernel {
     RuntimeCheck(layer_bs <= req_pool_indices.shape()[0], "layer_bs exceeds req_pool_indices rows");
     RuntimeCheck(layer_bs * num_verify_tokens <= kv_score_input.shape()[0], "kv_score_input is too small");
 
-    if (seq_dtype.is_type<int32_t>()) {
-      if (req_dtype.is_type<int32_t>()) {
-        launch<int32_t, int32_t>(
-            kv_score_input,
-            seq_lens,
-            req_pool_indices,
-            req_to_token,
-            full_to_swa,
-            ape,
-            state,
-            layer_bs,
-            swa_page_size,
-            num_verify_tokens,
-            state_slot_stride,
-            device.unwrap());
-      } else {
-        launch<int32_t, int64_t>(
-            kv_score_input,
-            seq_lens,
-            req_pool_indices,
-            req_to_token,
-            full_to_swa,
-            ape,
-            state,
-            layer_bs,
-            swa_page_size,
-            num_verify_tokens,
-            state_slot_stride,
-            device.unwrap());
-      }
-    } else {
-      if (req_dtype.is_type<int32_t>()) {
-        launch<int64_t, int32_t>(
-            kv_score_input,
-            seq_lens,
-            req_pool_indices,
-            req_to_token,
-            full_to_swa,
-            ape,
-            state,
-            layer_bs,
-            swa_page_size,
-            num_verify_tokens,
-            state_slot_stride,
-            device.unwrap());
-      } else {
-        launch<int64_t, int64_t>(
-            kv_score_input,
-            seq_lens,
-            req_pool_indices,
-            req_to_token,
-            full_to_swa,
-            ape,
-            state,
-            layer_bs,
-            swa_page_size,
-            num_verify_tokens,
-            state_slot_stride,
-            device.unwrap());
-      }
-    }
+    launch(
+        kv_score_input,
+        seq_lens,
+        req_pool_indices,
+        req_to_token,
+        ape,
+        state,
+        layer_bs,
+        num_verify_tokens,
+        state_slot_stride,
+        device.unwrap());
   }
 };
 
-template <int64_t kHeadDim>
+template <int64_t kHeadDim, typename TSeq, typename TReq>
 struct OnlineC128MTPMarkPendingKernel {
-  template <typename TSeq, typename TReq>
   static void launch(
       tvm::ffi::TensorView seq_lens,
       tvm::ffi::TensorView req_pool_indices,
@@ -368,13 +292,11 @@ struct OnlineC128MTPMarkPendingKernel {
       int64_t max_num_reqs) {
     using namespace host;
 
-    auto seq_dtype = SymbolicDType{};
-    auto req_dtype = SymbolicDType{};
     auto device = SymbolicDevice{};
     device.set_options<kDLCUDA>();
 
-    TensorMatcher({-1}).with_dtype<int32_t, int64_t>(seq_dtype).with_device(device).verify(seq_lens);
-    TensorMatcher({-1}).with_dtype<int32_t, int64_t>(req_dtype).with_device(device).verify(req_pool_indices);
+    TensorMatcher({-1}).with_dtype<TSeq>().with_device(device).verify(seq_lens);
+    TensorMatcher({-1}).with_dtype<TReq>().with_device(device).verify(req_pool_indices);
     TensorMatcher({-1}).with_dtype<int64_t>().with_device(device).verify(pending_seq_lens);
 
     if (bs <= 0) return;
@@ -382,34 +304,19 @@ struct OnlineC128MTPMarkPendingKernel {
     RuntimeCheck(bs <= req_pool_indices.shape()[0], "bs exceeds req_pool_indices rows");
     RuntimeCheck(max_num_reqs <= pending_seq_lens.shape()[0], "max_num_reqs exceeds pending rows");
 
-    if (seq_dtype.is_type<int32_t>()) {
-      if (req_dtype.is_type<int32_t>()) {
-        launch<int32_t, int32_t>(seq_lens, req_pool_indices, pending_seq_lens, bs, max_num_reqs, device.unwrap());
-      } else {
-        launch<int32_t, int64_t>(seq_lens, req_pool_indices, pending_seq_lens, bs, max_num_reqs, device.unwrap());
-      }
-    } else {
-      if (req_dtype.is_type<int32_t>()) {
-        launch<int64_t, int32_t>(seq_lens, req_pool_indices, pending_seq_lens, bs, max_num_reqs, device.unwrap());
-      } else {
-        launch<int64_t, int64_t>(seq_lens, req_pool_indices, pending_seq_lens, bs, max_num_reqs, device.unwrap());
-      }
-    }
+    launch(seq_lens, req_pool_indices, pending_seq_lens, bs, max_num_reqs, device.unwrap());
   }
 };
 
-template <int64_t kHeadDim>
+template <int64_t kHeadDim, typename TSeq, typename TReq>
 struct OnlineC128MTPCommitPendingKernel {
-  template <typename TSeq, typename TReq>
   static void launch(
       tvm::ffi::TensorView cur_seq_lens,
       tvm::ffi::TensorView cur_req_pool_indices,
       tvm::ffi::TensorView req_to_token,
-      tvm::ffi::TensorView full_to_swa,
       tvm::ffi::TensorView pending_seq_lens,
       tvm::ffi::TensorView state,
       int64_t cur_bs,
-      int64_t swa_page_size,
       int64_t num_verify_tokens,
       int64_t state_slot_stride,
       int64_t max_num_reqs,
@@ -420,13 +327,11 @@ struct OnlineC128MTPCommitPendingKernel {
         .cur_seq_lens = static_cast<const TSeq*>(cur_seq_lens.data_ptr()),
         .cur_req_pool_indices = static_cast<const TReq*>(cur_req_pool_indices.data_ptr()),
         .req_to_token = static_cast<const int32_t*>(req_to_token.data_ptr()),
-        .full_to_swa = static_cast<const int64_t*>(full_to_swa.data_ptr()),
         .pending_seq_lens = static_cast<const int64_t*>(pending_seq_lens.data_ptr()),
         .state = static_cast<float*>(state.data_ptr()),
         .cur_bs = cur_bs,
         .req_to_token_stride_b = req_to_token.stride(0),
         .state_stride_b = state.stride(0),
-        .swa_page_size = swa_page_size,
         .num_verify_tokens = num_verify_tokens,
         .state_slot_stride = state_slot_stride,
         .max_num_reqs = max_num_reqs,
@@ -441,25 +346,20 @@ struct OnlineC128MTPCommitPendingKernel {
   run(tvm::ffi::TensorView cur_seq_lens,
       tvm::ffi::TensorView cur_req_pool_indices,
       tvm::ffi::TensorView req_to_token,
-      tvm::ffi::TensorView full_to_swa,
       tvm::ffi::TensorView pending_seq_lens,
       tvm::ffi::TensorView state,
       int64_t cur_bs,
-      int64_t swa_page_size,
       int64_t num_verify_tokens,
       int64_t state_slot_stride,
       int64_t max_num_reqs) {
     using namespace host;
 
-    auto seq_dtype = SymbolicDType{};
-    auto req_dtype = SymbolicDType{};
     auto device = SymbolicDevice{};
     device.set_options<kDLCUDA>();
 
-    TensorMatcher({-1}).with_dtype<int32_t, int64_t>(seq_dtype).with_device(device).verify(cur_seq_lens);
-    TensorMatcher({-1}).with_dtype<int32_t, int64_t>(req_dtype).with_device(device).verify(cur_req_pool_indices);
+    TensorMatcher({-1}).with_dtype<TSeq>().with_device(device).verify(cur_seq_lens);
+    TensorMatcher({-1}).with_dtype<TReq>().with_device(device).verify(cur_req_pool_indices);
     TensorMatcher({-1, -1}).with_dtype<int32_t>().with_device(device).verify(req_to_token);
-    TensorMatcher({-1}).with_dtype<int64_t>().with_device(device).verify(full_to_swa);
     TensorMatcher({-1}).with_dtype<int64_t>().with_device(device).verify(pending_seq_lens);
     TensorMatcher({-1, kHeadDim * 3}).with_dtype<float>().with_device(device).verify(state);
 
@@ -470,67 +370,17 @@ struct OnlineC128MTPCommitPendingKernel {
     RuntimeCheck(cur_bs <= cur_req_pool_indices.shape()[0], "cur_bs exceeds req rows");
     RuntimeCheck(max_num_reqs <= pending_seq_lens.shape()[0], "max_num_reqs exceeds pending rows");
 
-    if (seq_dtype.is_type<int32_t>()) {
-      if (req_dtype.is_type<int32_t>()) {
-        launch<int32_t, int32_t>(
-            cur_seq_lens,
-            cur_req_pool_indices,
-            req_to_token,
-            full_to_swa,
-            pending_seq_lens,
-            state,
-            cur_bs,
-            swa_page_size,
-            num_verify_tokens,
-            state_slot_stride,
-            max_num_reqs,
-            device.unwrap());
-      } else {
-        launch<int32_t, int64_t>(
-            cur_seq_lens,
-            cur_req_pool_indices,
-            req_to_token,
-            full_to_swa,
-            pending_seq_lens,
-            state,
-            cur_bs,
-            swa_page_size,
-            num_verify_tokens,
-            state_slot_stride,
-            max_num_reqs,
-            device.unwrap());
-      }
-    } else {
-      if (req_dtype.is_type<int32_t>()) {
-        launch<int64_t, int32_t>(
-            cur_seq_lens,
-            cur_req_pool_indices,
-            req_to_token,
-            full_to_swa,
-            pending_seq_lens,
-            state,
-            cur_bs,
-            swa_page_size,
-            num_verify_tokens,
-            state_slot_stride,
-            max_num_reqs,
-            device.unwrap());
-      } else {
-        launch<int64_t, int64_t>(
-            cur_seq_lens,
-            cur_req_pool_indices,
-            req_to_token,
-            full_to_swa,
-            pending_seq_lens,
-            state,
-            cur_bs,
-            swa_page_size,
-            num_verify_tokens,
-            state_slot_stride,
-            max_num_reqs,
-            device.unwrap());
-      }
-    }
+    launch(
+        cur_seq_lens,
+        cur_req_pool_indices,
+        req_to_token,
+        pending_seq_lens,
+        state,
+        cur_bs,
+        num_verify_tokens,
+        state_slot_stride,
+        max_num_reqs,
+        device.unwrap());
   }
 };
 

python/sglang/jit_kernel/dsv4/__init__.py

@@ -3,6 +3,7 @@ from .attn import (
     get_paged_mqa_logits_metadata,
     triton_create_paged_compress_data,
 )
+from .c128_cleanup import clear_unaccepted_c128_draft_states
 from .compress import (
     CompressorDecodePlan,
     CompressorPrefillPlan,
@@ -35,6 +36,7 @@ __all__ = [
     "CompressorPrefillPlan",
     "compress_forward",
     "compress_norm_rope_store",
+    "clear_unaccepted_c128_draft_states",
     "fused_norm_rope_inplace",
     "fused_store_cache",
     "fused_rope_inplace",

python/sglang/jit_kernel/dsv4/attn.py

@@ -122,18 +122,21 @@ def create_paged_compress_data_kernel(
         else:
             pos = write_overlap_pos
         pos = tl.maximum(pos, 0)
-        loc = tl.load(
-            req_to_token_ptr
-            + rid.to(tl.int64) * stride_req_to_token_0
-            + pos.to(tl.int64) * stride_req_to_token_1,
-            mask=mask,
-            other=0,
-        ).to(tl.int32)
-        swa_loc = tl.load(full_to_swa_index_mapping_ptr + loc, mask=mask, other=0).to(
-            tl.int32
-        )
-        swa_page = swa_loc // swa_page_size
-        state_loc = swa_page * ring_size + (swa_loc % ring_size)
+        if compress_ratio == 128:
+            state_loc = rid * ring_size + (pos % ring_size)
+        else:
+            loc = tl.load(
+                req_to_token_ptr
+                + rid.to(tl.int64) * stride_req_to_token_0
+                + pos.to(tl.int64) * stride_req_to_token_1,
+                mask=mask,
+                other=0,
+            ).to(tl.int32)
+            swa_loc = tl.load(
+                full_to_swa_index_mapping_ptr + loc, mask=mask, other=0
+            ).to(tl.int32)
+            swa_page = swa_loc // swa_page_size
+            state_loc = swa_page * ring_size + (swa_loc % ring_size)
         state_loc = state_loc // cr
         if i == 0:
             v0 = state_loc

python/sglang/jit_kernel/dsv4/c128_cleanup.py

@@ -0,0 +1,58 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _clear_unaccepted_c128_draft_states_kernel(
+    state,
+    req_pool_indices,
+    seq_lens,
+    accept_lens,
+    ring_size: tl.constexpr,
+    half: tl.constexpr,
+    num_draft_tokens: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    bid = tl.program_id(0)
+    draft_offset = tl.program_id(1)
+    block_id = tl.program_id(2)
+
+    accept_len = tl.load(accept_lens + bid)
+    if draft_offset < accept_len:
+        return
+
+    req_pool_idx = tl.load(req_pool_indices + bid).to(tl.int64)
+    seq_len = tl.load(seq_lens + bid).to(tl.int64)
+    slot = (seq_len + draft_offset) % ring_size
+    row = req_pool_idx * ring_size + slot
+
+    offsets = block_id * BLOCK_D + tl.arange(0, BLOCK_D)
+    mask = offsets < half
+    row_base = row * (half * 2)
+    tl.store(state + row_base + offsets, 0.0, mask=mask)
+    tl.store(state + row_base + half + offsets, float("-inf"), mask=mask)
+
+
+def clear_unaccepted_c128_draft_states(
+    state: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    seq_lens: torch.Tensor,
+    accept_lens: torch.Tensor,
+    *,
+    ring_size: int,
+    num_draft_tokens: int,
+) -> None:
+    half = state.shape[-1] // 2
+    _clear_unaccepted_c128_draft_states_kernel[
+        (req_pool_indices.numel(), num_draft_tokens, triton.cdiv(half, 256))
+    ](
+        state,
+        req_pool_indices,
+        seq_lens,
+        accept_lens,
+        ring_size,
+        half,
+        num_draft_tokens,
+        BLOCK_D=256,
+    )

python/sglang/jit_kernel/dsv4/compress.py

@@ -125,7 +125,7 @@ class CompressorDecodePlan(NamedTuple):
         compress_ratio: Literal[4, 128],
         req_pool_indices: torch.Tensor,
         req_to_token: torch.Tensor,
-        full_to_swa: torch.Tensor,
+        full_to_state: torch.Tensor,
         seq_lens: torch.Tensor,
         swa_page_size: int,
         ring_size: int,
@@ -134,7 +134,7 @@ class CompressorDecodePlan(NamedTuple):
         plan_d = module.plan_decode(
             req_pool_indices,
             req_to_token,
-            full_to_swa,
+            full_to_state,
             seq_lens,
             int(compress_ratio),
             int(swa_page_size),
@@ -157,8 +157,6 @@ class CompressorDecodePlan(NamedTuple):
         seq_lens: torch.Tensor,
         req_pool_indices: torch.Tensor,
         req_to_token: torch.Tensor,
-        full_to_swa: torch.Tensor,
-        swa_page_size: int,
         state_slot_offset: int = 0,
     ) -> CompressorDecodePlan:
         batch_size = int(seq_lens.shape[0])
@@ -172,9 +170,7 @@ class CompressorDecodePlan(NamedTuple):
             seq_lens,
             req_pool_indices,
             req_to_token,
-            full_to_swa,
             plan_d,
-            swa_page_size,
             int(state_slot_offset),
         )
         return CompressorDecodePlan(128, plan_d)
@@ -203,7 +199,7 @@ class CompressorPrefillPlan(NamedTuple):
         seq_lens: torch.Tensor,
         extend_lens: torch.Tensor,
         req_to_token: torch.Tensor,
-        full_to_swa: torch.Tensor,
+        full_to_state: torch.Tensor,
         swa_page_size: int,
         ring_size: int,
         num_q_tokens: int,
@@ -219,7 +215,7 @@ class CompressorPrefillPlan(NamedTuple):
         plan_c, plan_w = module.plan_prefill(
             req_pool_indices,
             req_to_token,
-            full_to_swa,
+            full_to_state,
             seq_lens,
             extend_lens,
             pin_buffer,
@@ -274,9 +270,7 @@ class CompressorPrefillPlan(NamedTuple):
         extend_lens: torch.Tensor,
         req_pool_indices: torch.Tensor,
         req_to_token: torch.Tensor,
-        full_to_swa: torch.Tensor,
         num_q_tokens: int,
-        swa_page_size: int,
         use_cuda_graph: bool = False,
         state_slot_offset: int = 0,
     ) -> CompressorPrefillPlan:
@@ -284,7 +278,6 @@ class CompressorPrefillPlan(NamedTuple):
         extend_lens_cpu = extend_lens.detach().to(torch.int64).cpu()
         rid_i64 = req_pool_indices.to(torch.int64)
         r2t_i32 = req_to_token.to(torch.int32)
-        f2s_i64 = full_to_swa.to(torch.int64)
         pin_buffer = torch.empty(
             (2, num_q_tokens, 16), dtype=torch.uint8, pin_memory=True
         )
@@ -298,12 +291,10 @@ class CompressorPrefillPlan(NamedTuple):
             extend_lens_cpu,
             rid_i64,
             r2t_i32,
-            f2s_i64,
             plan_c_pin,
             plan_w_pin,
             plan_c_dev,
             plan_w_dev,
-            int(swa_page_size),
             int(state_slot_offset),
             bool(use_cuda_graph),
         )

python/sglang/jit_kernel/dsv4/online_c128_mtp.py

@@ -14,8 +14,10 @@ if TYPE_CHECKING:
 
 
 @cache_once
-def _jit_online_c128_mtp_module(head_dim: int) -> Module:
-    args = make_cpp_args(head_dim)
+def _jit_online_c128_mtp_module(
+    head_dim: int, seq_dtype: torch.dtype, req_dtype: torch.dtype
+) -> Module:
+    args = make_cpp_args(head_dim, seq_dtype, req_dtype)
     return load_jit(
         make_name(f"online_c128_mtp_{head_dim}"),
         *args,
@@ -77,12 +79,14 @@ class OnlineC128MTPController:
         if head_dim is None or self._num_verify_tokens() == 0:
             return
         token_to_kv_pool = self.backend.token_to_kv_pool
-        _jit_online_c128_mtp_module(head_dim).mark_pending(
+        _jit_online_c128_mtp_module(
+            head_dim, seq_lens.dtype, req_pool_indices.dtype
+        ).mark_pending(
             seq_lens,
             req_pool_indices,
             token_to_kv_pool.get_online_c128_mtp_pending_seq_lens(),
             min(seq_lens.shape[0], req_pool_indices.shape[0]),
-            token_to_kv_pool.max_num_reqs,
+            token_to_kv_pool.get_online_c128_state_num_req_slots(),
         )
 
     def clear(self) -> None:
@@ -157,16 +161,16 @@ class OnlineC128MTPController:
         if layer_bs <= 0:
             return
 
-        _jit_online_c128_mtp_module(head_dim).write_prefix_states(
+        _jit_online_c128_mtp_module(
+            head_dim, ctx.seq_lens.dtype, ctx.req_pool_indices.dtype
+        ).write_prefix_states(
             kv_score_input,
             ctx.seq_lens,
             ctx.req_pool_indices,
             self.backend.req_to_token,
-            token_to_kv_pool.full_to_swa_index_mapping,
             compressor.ape.reshape(128, head_dim),
             state_pool.kv_score_buffer.kv_score,
             layer_bs,
-            token_to_kv_pool.swa_page_size,
             num_verify_tokens,
             state_pool.online_mtp_state_slot_offset,
         )
@@ -195,18 +199,18 @@ class OnlineC128MTPController:
         cur_bs = min(seq_lens.shape[0], req_pool_indices.shape[0])
 
         for runtime in self._iter_layer_runtimes():
-            _jit_online_c128_mtp_module(runtime.head_dim).commit_pending(
+            _jit_online_c128_mtp_module(
+                runtime.head_dim, seq_lens.dtype, req_pool_indices.dtype
+            ).commit_pending(
                 seq_lens,
                 req_pool_indices,
                 backend.req_to_token,
-                token_to_kv_pool.full_to_swa_index_mapping,
                 pending_seq_lens,
                 runtime.main_state,
                 cur_bs,
-                token_to_kv_pool.swa_page_size,
                 num_verify_tokens,
                 runtime.state_slot_offset,
-                token_to_kv_pool.max_num_reqs,
+                token_to_kv_pool.get_online_c128_state_num_req_slots(),
             )
 
         self.clear()

python/sglang/jit_kernel/tests/deepseek_v4/common.py

@@ -103,7 +103,7 @@ class PagedContext:
             seq_lens=seq_lens_cpu,
             extend_lens=extend_lens_cpu,
             req_to_token=self.req_to_token,
-            full_to_swa=self.full_to_swa,
+            full_to_state=self.full_to_swa,
             swa_page_size=self.swa_page_size,
             ring_size=self.ring_size,
             num_q_tokens=num_q_tokens,
@@ -114,7 +114,7 @@ class PagedContext:
             compress_ratio=self.compress_ratio,  # type: ignore
             req_pool_indices=self.req_pool_indices,
             req_to_token=self.req_to_token,
-            full_to_swa=self.full_to_swa,
+            full_to_state=self.full_to_swa,
             seq_lens=seq_lens_gpu,
             swa_page_size=self.swa_page_size,
             ring_size=self.ring_size,

python/sglang/srt/disaggregation/base/conn.py

@@ -22,6 +22,8 @@ class StateType(str, enum.Enum):
     # DeepSeek-V4 unified_kv SWA ring: addressed per-row by ring slot
     # (req_pool_idx * ring_stride + pos % ring_stride), needs its own component.
     SWA_RING = "swa_ring"
+    # DeepSeek-V4 online C128 request-scoped state.
+    C128_STATE = "c128_state"
 
 
 @dataclasses.dataclass

python/sglang/srt/disaggregation/common/conn.py

@@ -23,6 +23,7 @@ from sglang.srt.disaggregation.base.conn import (
     KVArgs,
     KVPoll,
     KVTransferMetric,
+    StateType,
 )
 from sglang.srt.disaggregation.utils import (
     DisaggregationMode,
@@ -522,7 +523,10 @@ class CommonKVManager(BaseKVManager):
         return src_k_ptrs, src_v_ptrs, dst_k_ptrs, dst_v_ptrs, layers_current_pp_stage
 
     def get_mla_kv_ptrs_with_pp(
-        self, src_kv_ptrs: List[int], dst_kv_ptrs: List[int]
+        self,
+        src_kv_ptrs: List[int],
+        dst_kv_ptrs: List[int],
+        state_type: Optional[StateType] = None,
     ) -> Tuple[List[int], List[int], int]:
         # Fast path: both sides use exactly the same PP layout
         if len(src_kv_ptrs) == len(dst_kv_ptrs):
@@ -535,7 +539,7 @@ class CommonKVManager(BaseKVManager):
             # layer, so we locate the sub-range for this PP stage inside each
             # section of the dst flat list.
             sliced_src_kv_ptrs, sliced_dst_kv_ptrs = self._mla_slice_ptrs_for_pp(
-                src_kv_ptrs, dst_kv_ptrs, mla_ratios
+                src_kv_ptrs, dst_kv_ptrs, mla_ratios, state_type
             )
             return (
                 sliced_src_kv_ptrs,
@@ -555,6 +559,7 @@ class CommonKVManager(BaseKVManager):
         src_kv_ptrs: List[int],
         dst_kv_ptrs: List[int],
         mla_ratios: List[int],
+        state_type: Optional[StateType] = None,
     ) -> Tuple[List[int], List[int]]:
         """Produce aligned (src, dst) pointer lists for compressed-MLA
         pools (e.g. DeepSeek V4) under PP.
@@ -569,16 +574,19 @@ class CommonKVManager(BaseKVManager):
           Each section is indexed by compressed-layer id within that
           compression bucket.
 
-        - state_data layout, length = swa_L + 2 * c4_L + c128_L:
+        - SWA state_data layout, length = swa_L + 2 * c4_L:
             [swa_layer_{0..swa_L-1},
-             compress_state_{non-None, c4_L + c128_L},
-             indexer_compress_state_{non-None, c4_L}]
+             c4_compress_state_{0..c4_L-1},
+             c4_indexer_compress_state_{0..c4_L-1}]
           ``swa_L`` is the SWA pool's actual buffer count
           (``num_effective_layers``), which can be smaller than
           ``len(mla_ratios)`` when the HF config's ``compress_ratios``
           list contains entries for layers not materialized into the SWA
           pool (e.g. an MTP/nextn slot at the tail).
 
+        - C128_STATE layout, length = c128_L:
+            [c128_compress_state_{0..c128_L-1}]
+
         src is already PP-filtered on the prefill side. dst is the
         decode-side full-model list (when decode is PP=1). We slice dst to
         match src's PP stage. If src itself is also full-model, it is
@@ -600,7 +608,18 @@ class CommonKVManager(BaseKVManager):
         c128_off_s = sum(1 for r in mla_ratios[:start_layer] if r == 128)
         c128_off_e = sum(1 for r in mla_ratios[:end_layer] if r == 128)
 
-        if len(dst_kv_ptrs) == kv_layout_len:
+        if state_type == StateType.C128_STATE:
+            return src_kv_ptrs, list(dst_kv_ptrs[c128_off_s:c128_off_e])
+
+        if state_type == StateType.SWA_RING:
+            swa_s = min(start_layer, len(dst_kv_ptrs))
+            swa_e = min(end_layer, len(dst_kv_ptrs))
+            return src_kv_ptrs, list(dst_kv_ptrs[swa_s:swa_e])
+
+        if (
+            state_type not in (StateType.SWA, StateType.SWA_RING, StateType.C128_STATE)
+            and len(dst_kv_ptrs) == kv_layout_len
+        ):
             sliced_dst = (
                 list(dst_kv_ptrs[c4_off_s:c4_off_e])
                 + list(dst_kv_ptrs[c4_full + c4_off_s : c4_full + c4_off_e])
@@ -608,39 +627,33 @@ class CommonKVManager(BaseKVManager):
             )
             return src_kv_ptrs, sliced_dst
 
-        # State-data layout. ``swa_L`` is derived from the actual dst
-        # length so we tolerate cases where the SWA pool has fewer
-        # buffers than ``len(mla_ratios)`` (e.g. nextn padding).
-        swa_L = len(dst_kv_ptrs) - 2 * c4_full - c128_full
+        # SWA state-data layout. ``swa_L`` is derived from the actual dst
+        # length so we tolerate cases where the SWA pool has fewer buffers
+        # than ``len(mla_ratios)`` (e.g. nextn padding). C128 state ships as
+        # a separate StateType.C128_STATE component and must not be counted
+        # here.
+        swa_L = len(dst_kv_ptrs) - 2 * c4_full
         if swa_L < 0 or swa_L > len(mla_ratios):
             raise ValueError(
                 f"Unexpected compressed-MLA dst_kv_ptrs length "
                 f"{len(dst_kv_ptrs)}; expected either {kv_layout_len} "
-                f"(kv_data) or swa_L + {2 * c4_full + c128_full} "
+                f"(kv_data) or swa_L + {2 * c4_full} "
                 f"(state_data) given compression_ratios "
                 f"(c4={c4_full}, c128={c128_full}, "
                 f"total={len(mla_ratios)})."
             )
-        # Guard against asking the prefill side to read past the SWA
-        # pool boundary.
-        assert end_layer <= swa_L, (
-            f"prefill_end_layer ({end_layer}) exceeds dst SWA pool "
-            f"buffer count ({swa_L}); compression_ratios may include "
-            f"layers (e.g. nextn) that the SWA pool does not cover."
-        )
 
-        # compress_state non-None count up to L = count(r != 0).
-        c_non_zero_s = sum(1 for r in mla_ratios[:start_layer] if r != 0)
-        c_non_zero_e = sum(1 for r in mla_ratios[:end_layer] if r != 0)
+        swa_s = min(start_layer, swa_L)
+        swa_e = min(end_layer, swa_L)
         compress_section_start = swa_L
-        indexer_section_start = swa_L + (c4_full + c128_full)
+        indexer_section_start = swa_L + c4_full
         sliced_dst = (
-            list(dst_kv_ptrs[start_layer:end_layer])
+            list(dst_kv_ptrs[swa_s:swa_e])
             + list(
                 dst_kv_ptrs[
                     compress_section_start
-                    + c_non_zero_s : compress_section_start
-                    + c_non_zero_e
+                    + c4_off_s : compress_section_start
+                    + c4_off_e
                 ]
             )
             + list(
@@ -1127,6 +1140,7 @@ class CommonKVReceiver(BaseKVReceiver):
         kv_indices: npt.NDArray[np.int32],
         aux_index: Optional[int] = None,
         state_indices: Optional[List[int]] = None,
+        decode_prefix_len: Optional[int] = None,
     ):
         raise NotImplementedError
 

python/sglang/srt/disaggregation/decode.py

@@ -50,7 +50,9 @@ from sglang.srt.disaggregation.utils import (
     ReqToMetadataIdxAllocator,
     TransferBackend,
     _is_fake_transfer,
+    get_dsv4_c128_state_indices,
     get_kv_class,
+    is_dsv4_c128_online_enabled,
     is_mla_backend,
     poll_and_all_reduce,
     poll_and_all_reduce_with_staging,
@@ -1049,8 +1051,24 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
                 ring_rows = state_slot * ring_stride + (positions % ring_stride)
                 return ring_rows.astype(np.int32)
 
+            def _c128_state_payload():
+                online = is_dsv4_c128_online_enabled()
+                ring_size = 1 if online else self.token_to_kv_pool.get_ring_size(128)
+                return get_dsv4_c128_state_indices(
+                    int(decode_req.req.req_pool_idx),
+                    seq_len,
+                    online=online,
+                    ring_size=ring_size,
+                )
+
             state_types = self.kv_manager.kv_args.state_types
             state_indices: Optional[List] = []
+            if StateType.C128_STATE in state_types:
+                clear_c128_state = getattr(
+                    self.token_to_kv_pool, "clear_c128_req_state", None
+                )
+                if clear_c128_state is not None:
+                    clear_c128_state(int(decode_req.req.req_pool_idx))
             for st in state_types:
                 if st == StateType.MAMBA:
                     state_indices.append(_mamba_payload())
@@ -1064,6 +1082,8 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
                     state_indices.append(_dsa_payload())
                 elif st == StateType.SWA_RING:
                     state_indices.append(_swa_ring_payload())
+                elif st == StateType.C128_STATE:
+                    state_indices.append(_c128_state_payload())
                 else:
                     state_indices.append(None)
 

python/sglang/srt/disaggregation/mooncake/conn.py

@@ -586,6 +586,7 @@ class MooncakeKVManager(CommonKVManager):
         prefill_data_indices: npt.NDArray[np.int32],
         dst_data_indices: npt.NDArray[np.int32],
         executor: concurrent.futures.ThreadPoolExecutor,
+        state_type: Optional[StateType] = None,
         force_flat: bool = False,
     ) -> int:
         """
@@ -606,7 +607,7 @@ class MooncakeKVManager(CommonKVManager):
         # Decode pp size should be equal to prefill pp size or 1
         if self.is_mla_backend or force_flat:
             src_kv_ptrs, dst_kv_ptrs, layers_current_pp_stage = (
-                self.get_mla_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs)
+                self.get_mla_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs, state_type)
             )
             layers_params = [
                 (
@@ -1001,7 +1002,12 @@ class MooncakeKVManager(CommonKVManager):
                         )
                         or rc
                     )
-            elif st in (StateType.SWA, StateType.DSA, StateType.SWA_RING):
+            elif st in (
+                StateType.SWA,
+                StateType.DSA,
+                StateType.SWA_RING,
+                StateType.C128_STATE,
+            ):
                 if (
                     target_rank_registration_info is not None
                     and not self.is_mla_backend
@@ -1013,13 +1019,20 @@ class MooncakeKVManager(CommonKVManager):
                     )
                 src_indices = list(indices)
                 dst_indices_local = list(dst_indices)
+                if (
+                    st == StateType.C128_STATE
+                    and len(src_indices) == 0
+                    and len(dst_indices_local) == 0
+                ):
+                    continue
                 if len(src_indices) != len(dst_indices_local):
-                    # SWA_RING is positional: truncating silently misaligns rows
-                    # and corrupts KV, so fail loud. Paged SWA/DSA tolerate a
-                    # 1-page drift -> keep the lenient truncation below.
-                    if st == StateType.SWA_RING:
+                    # These components are position- or request-indexed:
+                    # truncating silently misaligns rows and corrupts KV.
+                    # Paged SWA/DSA tolerate a 1-page drift -> keep the
+                    # lenient truncation below.
+                    if st in (StateType.SWA_RING, StateType.C128_STATE):
                         raise RuntimeError(
-                            "SWA_RING state index length mismatch: "
+                            f"{st.upper()} state index length mismatch: "
                             f"prefill={len(src_indices)}, dst={len(dst_indices_local)}"
                         )
                     logger.warning(
@@ -1038,6 +1051,7 @@ class MooncakeKVManager(CommonKVManager):
                         prefill_data_indices=np.array(src_indices, dtype=np.int32),
                         dst_data_indices=np.array(dst_indices_local, dtype=np.int32),
                         executor=executor,
+                        state_type=st,
                     )
                     or rc
                 )

python/sglang/srt/disaggregation/mori/conn.py

@@ -1092,7 +1092,7 @@ class MoriKVManager(CommonKVManager):
                         dst_dims,
                     )
                 )
-            elif st in ("swa", "dsa", "swa_ring"):
+            elif st in ("swa", "dsa", "swa_ring", "c128_state"):
                 statuses.extend(
                     self._send_swa_dsa_state(
                         peer_info,
@@ -1221,17 +1221,24 @@ class MoriKVManager(CommonKVManager):
             )
 
         common_len = min(src_state_indices.size, dst_state_indices.size)
+        if (
+            state_type == "c128_state"
+            and common_len == 0
+            and src_state_indices.size == 0
+            and dst_state_indices.size == 0
+        ):
+            return []
         if common_len == 0 and max(src_state_indices.size, dst_state_indices.size) > 0:
             raise RuntimeError(
                 f"No overlapping state indices for state_type={state_type}"
             )
         if src_state_indices.size != dst_state_indices.size:
-            # SWA_RING is positional: truncating silently misaligns rows and
-            # corrupts KV, so fail loud. Paged swa/dsa tolerate a 1-page drift
-            # -> keep truncation.
-            if state_type == "swa_ring":
+            # These components are position- or request-indexed: truncating
+            # silently misaligns rows and corrupts KV. Paged swa/dsa tolerate
+            # a 1-page drift -> keep truncation.
+            if state_type in ("swa_ring", "c128_state"):
                 raise RuntimeError(
-                    "SWA_RING state index length mismatch: "
+                    f"{state_type.upper()} state index length mismatch: "
                     f"src={src_state_indices.size}, dst={dst_state_indices.size}"
                 )
             logger.warning(

python/sglang/srt/disaggregation/nixl/conn.py

@@ -1277,6 +1277,7 @@ class NixlKVManager(CommonKVManager):
         dst_data_indices: npt.NDArray[np.int32],
         dst_gpu_id: int,
         notif: str,
+        state_type: Optional[StateType] = None,
         src_mem_kind: str = "VRAM",
         dst_mem_kind: str = "VRAM",
         force_flat: bool = False,
@@ -1341,7 +1342,7 @@ class NixlKVManager(CommonKVManager):
         # Make descs
         if self.is_mla_backend or force_flat:
             src_kv_ptrs, dst_kv_ptrs, layers_current_pp_stage = (
-                self.get_mla_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs)
+                self.get_mla_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs, state_type)
             )
             layers_params = [
                 (
@@ -1994,11 +1995,22 @@ class NixlKVManager(CommonKVManager):
                         dst_gpu_id,
                         comp_notif,
                     )
-            elif st in (StateType.SWA, StateType.DSA, StateType.SWA_RING):
+            elif st in (
+                StateType.SWA,
+                StateType.DSA,
+                StateType.SWA_RING,
+                StateType.C128_STATE,
+            ):
                 if not self.is_mla_backend and self.attn_tp_size != decode_tp_size:
                     raise RuntimeError(
                         f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {st.upper()} hybrid models yet."
                     )
+                if (
+                    st == StateType.C128_STATE
+                    and len(src_indices) == 0
+                    and len(dst_indices) == 0
+                ):
+                    continue
                 if len(src_indices) != len(dst_indices):
                     raise RuntimeError(
                         f"State index length mismatch at component {i}: "
@@ -2013,6 +2025,7 @@ class NixlKVManager(CommonKVManager):
                     dst_data_indices=np.array(dst_indices, dtype=np.int32),
                     dst_gpu_id=dst_gpu_id,
                     notif=comp_notif,
+                    state_type=st,
                 )
             elif st == StateType.MINIMAX_INDEX_K:
                 # Equal-TP / PP=1 only. Sub-pools are compacted sparse-layer

python/sglang/srt/disaggregation/prefill.py

@@ -39,8 +39,10 @@ from sglang.srt.disaggregation.utils import (
     MetadataBuffers,
     ReqToMetadataIdxAllocator,
     TransferBackend,
+    get_dsv4_c128_state_indices,
     get_kv_class,
     is_aborted,
+    is_dsv4_c128_online_enabled,
     is_mla_backend,
     poll_and_all_reduce_attn_cp_tp_group,
     prepare_abort,
@@ -269,9 +271,8 @@ class PrefillBootstrapQueue:
             return False
 
         req.time_stats.set_bootstrap_done_time()
-        num_kv_indices = len(req.origin_input_ids)
-
         decode_prefix_len = req.disagg_kv_sender.pop_decode_prefix_len()
+        num_kv_indices = len(req.origin_input_ids)
         req.start_send_idx = decode_prefix_len
         num_kv_indices_to_send = num_kv_indices - decode_prefix_len
         num_pages = kv_to_page_num(
@@ -979,10 +980,11 @@ class SchedulerDisaggregationPrefillMixin:
         """
         page_size = self.token_to_kv_pool_allocator.page_size
         start_idx = req.start_send_idx
+        transfer_input_len = len(req.origin_input_ids)
         end_idx = (
             end_idx
             if end_idx is not None
-            else min(req.extend_range.end, len(req.origin_input_ids))
+            else min(req.extend_range.end, transfer_input_len)
         )
 
         if not last_chunk:
@@ -1007,13 +1009,12 @@ class SchedulerDisaggregationPrefillMixin:
         if last_chunk:
             self.disagg_metadata_buffers.set_buf(req)
 
-            # fill_ids includes the token sampled during prefill, but decode
-            # registers state pages over origin_input_ids (DecodePreallocQueue)
-            # and the main pool send is clamped to end_idx above. Matching that
-            # length here avoids emitting an extra state page when the sampled
-            # token crosses a page boundary, which mismatched src/dst lengths in
-            # group_concurrent_contiguous.
-            seq_len = min(req.extend_range.end, len(req.origin_input_ids))
+            # Most state payloads read token-pool rows and should match the KV
+            # range actually materialized on prefill. C128 state is request
+            # scoped, so its transfer index must use the logical input length
+            # that decode used to register the destination row.
+            seq_len = min(req.extend_range.end, transfer_input_len)
+            c128_seq_len = transfer_input_len
 
             def _mamba_payload():
                 return [
@@ -1059,6 +1060,22 @@ class SchedulerDisaggregationPrefillMixin:
                 ring_rows = state_slot * ring_stride + (positions % ring_stride)
                 return ring_rows.astype(np.int32)
 
+            def _c128_state_payload():
+                online = is_dsv4_c128_online_enabled()
+                ring_size = (
+                    1
+                    if online
+                    else self.token_to_kv_pool_allocator.get_kvcache().get_ring_size(
+                        128
+                    )
+                )
+                return get_dsv4_c128_state_indices(
+                    int(req.req_pool_idx),
+                    c128_seq_len,
+                    online=online,
+                    ring_size=ring_size,
+                )
+
             state_types = (
                 self.disagg_prefill_bootstrap_queue.kv_manager.kv_args.state_types
             )
@@ -1076,6 +1093,8 @@ class SchedulerDisaggregationPrefillMixin:
                     state_indices.append(_dsa_payload())
                 elif st == StateType.SWA_RING:
                     state_indices.append(_swa_ring_payload())
+                elif st == StateType.C128_STATE:
+                    state_indices.append(_c128_state_payload())
                 else:
                     state_indices.append(None)
 

python/sglang/srt/disaggregation/utils.py

@@ -13,7 +13,7 @@ import torch.distributed as dist
 
 from sglang.srt.disaggregation.base import KVPoll
 from sglang.srt.environ import envs
-from sglang.srt.utils import is_npu
+from sglang.srt.utils import is_hip, is_npu
 
 if TYPE_CHECKING:
     from sglang.srt.disaggregation.base.conn import KVArgs, StateType
@@ -30,6 +30,31 @@ if TYPE_CHECKING:
 # Constants & Enums
 #########################
 FAKE_BOOTSTRAP_HOST = "2.2.2.2"
+_IS_HIP = is_hip()
+
+
+def is_dsv4_c128_online_enabled() -> bool:
+    """Return whether DSV4 C128 uses request-scoped online state."""
+    return not _IS_HIP and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get()
+
+
+def get_dsv4_c128_state_indices(
+    req_pool_idx: int,
+    seq_len: int,
+    *,
+    online: bool,
+    ring_size: int,
+) -> np.ndarray:
+    """Return the PD transfer row/page indices for DSV4 C128 state."""
+    if seq_len == 0 or seq_len % 128 == 0:
+        return np.empty((0,), dtype=np.int32)
+    if online:
+        return np.array([int(req_pool_idx)], dtype=np.int32)
+
+    assert ring_size % 128 == 0, f"C128 ring_size must be 128-aligned, got {ring_size}"
+    pages_per_req = ring_size // 128
+    page = int(req_pool_idx) * pages_per_req + ((seq_len - 1) % ring_size) // 128
+    return np.array([page], dtype=np.int32)
 
 
 class DisaggregationMode(Enum):
@@ -689,6 +714,18 @@ def setup_state_kv_args(
                         ring_lens,
                         ring_item_lens,
                     )
+            if hasattr(token_to_kv_pool, "get_c128_state_buf_infos"):
+                c128_ptrs, c128_lens, c128_item_lens = (
+                    token_to_kv_pool.get_c128_state_buf_infos()
+                )
+                if c128_ptrs:
+                    append_state_component(
+                        kv_args,
+                        StateType.C128_STATE,
+                        c128_ptrs,
+                        c128_lens,
+                        c128_item_lens,
+                    )
         elif isinstance(token_to_kv_pool, HybridLinearKVPool):
             dim = (
                 token_to_kv_pool.get_state_dim_per_tensor()

python/sglang/srt/environ.py

@@ -855,6 +855,7 @@ class Envs:
     SGLANG_OPT_USE_ONLINE_COMPRESS = EnvBool(False)
     SGLANG_EXPERIMENTAL_ONLINE_C128_MTP = EnvBool(False)
     SGLANG_DSV4_COMPRESS_STATE_DTYPE = EnvStr("float32")
+    # Deprecated: DSV4 compressor V2 is always used.
     SGLANG_OPT_USE_COMPRESSOR_V2 = EnvBool(True)
     SGLANG_FP8_PAGED_MQA_LOGITS_TORCH = EnvBool(False)
     SGLANG_TOPK_TRANSFORM_512_TORCH = EnvBool(False)

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

@@ -18,25 +18,14 @@ from typing import (
 import torch
 import torch.nn.functional as F
 
+from sglang.jit_kernel.dsv4.online_c128_mtp import OnlineC128MTPController
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
-from sglang.srt.runtime_context import get_parallel
-
-if envs.SGLANG_OPT_USE_COMPRESSOR_V2.get():
-    # NOTE: should eventually be the only compressor backend
-    from sglang.srt.layers.attention.dsv4.compressor_v2 import (
-        CompressorBackendMixin,
-        FusedCompressMetadata,
-        create_paged_compressor_data,
-    )
-else:
-    from sglang.srt.layers.attention.dsv4.compressor import (
-        CompressorBackendMixin,
-        FusedCompressMetadata,
-        create_paged_compressor_data,
-    )
-
-from sglang.jit_kernel.dsv4.online_c128_mtp import OnlineC128MTPController
+from sglang.srt.layers.attention.dsv4.compressor_v2 import (
+    CompressorBackendMixin,
+    FusedCompressMetadata,
+    create_paged_compressor_data,
+)
 from sglang.srt.layers.attention.dsv4.dequant_k_cache import (
     dequantize_k_cache_paged,
 )
@@ -58,6 +47,7 @@ from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import (
 )
 from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.speculative.eagle_utils import per_step_draft_out_cache_loc
 from sglang.srt.utils import ceil_align
 from sglang.srt.utils.common import is_sm120_supported

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

@@ -20,21 +20,11 @@ import torch.nn.functional as F
 
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
-from sglang.srt.runtime_context import get_parallel
-
-if envs.SGLANG_OPT_USE_COMPRESSOR_V2.get():
-    from sglang.srt.layers.attention.dsv4.compressor_v2 import (
-        CompressorBackendMixin,
-        FusedCompressMetadata,
-        create_paged_compressor_data,
-    )
-else:
-    from sglang.srt.layers.attention.dsv4.compressor import (
-        CompressorBackendMixin,
-        FusedCompressMetadata,
-        create_paged_compressor_data,
-    )
-
+from sglang.srt.layers.attention.dsv4.compressor_v2 import (
+    CompressorBackendMixin,
+    FusedCompressMetadata,
+    create_paged_compressor_data,
+)
 from sglang.srt.layers.attention.dsv4.indexer import C4IndexerBackendMixin
 from sglang.srt.layers.attention.dsv4.metadata import (
     PagedIndexerMetadata,
@@ -49,6 +39,7 @@ from sglang.srt.layers.attention.dsv4.quant_k_cache import (
 )
 from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.speculative.eagle_utils import per_step_draft_out_cache_loc
 from sglang.srt.utils import ceil_align
 

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

@@ -210,13 +210,18 @@ class CompressorHip(_CompressorBase):
             pre_state_indices = self.compute_state_len_indices(
                 seq_len=prefix_lens[i], ratio=self.ratio
             ).to(device)
-            raw_loc = torch.where(
-                pre_state_indices < 0,
-                -1,
-                req_to_token[req_pool_indices[i], pre_state_indices],
-            )
-            swa_loc = token_to_kv_pool.translate_loc_from_full_to_swa(raw_loc)
-            state_loc = state_pool.translate_from_swa_loc_to_state_loc(swa_loc)
+            if self.ratio == 128:
+                state_loc = state_pool.translate_from_req_position_to_state_loc(
+                    req_pool_indices[i], pre_state_indices
+                )
+            else:
+                raw_loc = torch.where(
+                    pre_state_indices < 0,
+                    -1,
+                    req_to_token[req_pool_indices[i], pre_state_indices],
+                )
+                swa_loc = token_to_kv_pool.translate_loc_from_full_to_swa(raw_loc)
+                state_loc = state_pool.translate_from_swa_loc_to_state_loc(swa_loc)
             pre_kv_state = state_pool.get_state_by_state_loc(state_loc)
             kv_and_score_buffer = KVAndScore.cat([pre_kv_state, kv_and_score], dim=0)
             valid_kv_len = kv_and_score_buffer.kv.size(0)
@@ -227,15 +232,22 @@ class CompressorHip(_CompressorBase):
             post_state_len = post_state_indices.size(0)
 
             assert post_state_len <= valid_kv_len
-            post_raw_loc = torch.where(
-                post_state_indices < 0,
-                -1,
-                req_to_token[req_pool_indices[i], post_state_indices],
-            )
-            post_swa_loc = token_to_kv_pool.translate_loc_from_full_to_swa(post_raw_loc)
-            post_state_loc = state_pool.translate_from_swa_loc_to_state_loc(
-                post_swa_loc
-            )
+            if self.ratio == 128:
+                post_state_loc = state_pool.translate_from_req_position_to_state_loc(
+                    req_pool_indices[i], post_state_indices
+                )
+            else:
+                post_raw_loc = torch.where(
+                    post_state_indices < 0,
+                    -1,
+                    req_to_token[req_pool_indices[i], post_state_indices],
+                )
+                post_swa_loc = token_to_kv_pool.translate_loc_from_full_to_swa(
+                    post_raw_loc
+                )
+                post_state_loc = state_pool.translate_from_swa_loc_to_state_loc(
+                    post_swa_loc
+                )
             post_state_to_set = kv_and_score_buffer[valid_kv_len - post_state_len :]
             state_pool.set_state_by_state_loc(post_state_loc, post_state_to_set)
 
@@ -337,10 +349,14 @@ class CompressorHip(_CompressorBase):
             seq_lens = seq_lens_2d.view(-1)
             req_pool_indices = req_pool_indices.repeat_interleave(draft_tokens)
 
-        raw_locs = req_to_token[req_pool_indices, seq_lens - 1]
-
-        swa_locs = token_to_kv_pool.translate_loc_from_full_to_swa(raw_locs)
-        state_locs = state_pool.translate_from_swa_loc_to_state_loc(swa_locs)
+        if self.ratio == 128:
+            state_locs = state_pool.translate_from_req_position_to_state_loc(
+                req_pool_indices, seq_lens - 1
+            )
+        else:
+            raw_locs = req_to_token[req_pool_indices, seq_lens - 1]
+            swa_locs = token_to_kv_pool.translate_loc_from_full_to_swa(raw_locs)
+            state_locs = state_pool.translate_from_swa_loc_to_state_loc(swa_locs)
         state_pool.set_state_by_state_loc(state_locs, kv_and_scores)
 
         compress_bulk_len = self.ratio * self.coff
@@ -348,17 +364,24 @@ class CompressorHip(_CompressorBase):
             -compress_bulk_len, 0, device=seq_lens.device
         )
         compress_indices.clamp_(min=-1)
-        compress_indices_raw = torch.where(
-            compress_indices < 0,
-            -1,
-            req_to_token[req_pool_indices[:, None], compress_indices],
-        )
-        compress_indices_swa = token_to_kv_pool.translate_loc_from_full_to_swa(
-            compress_indices_raw
-        )
-        compress_indices_state = state_pool.translate_from_swa_loc_to_state_loc(
-            compress_indices_swa
-        )
+        if self.ratio == 128:
+            compress_indices_state = (
+                state_pool.translate_from_req_position_to_state_loc(
+                    req_pool_indices[:, None], compress_indices
+                )
+            )
+        else:
+            compress_indices_raw = torch.where(
+                compress_indices < 0,
+                -1,
+                req_to_token[req_pool_indices[:, None], compress_indices],
+            )
+            compress_indices_swa = token_to_kv_pool.translate_loc_from_full_to_swa(
+                compress_indices_raw
+            )
+            compress_indices_state = state_pool.translate_from_swa_loc_to_state_loc(
+                compress_indices_swa
+            )
         kv_and_score_to_compress = state_pool.get_state_by_state_loc(
             compress_indices_state.view(-1)
         ).view(-1, self.ratio, self.coff * self.head_dim)

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

@@ -31,15 +31,9 @@ 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
 from sglang.srt.runtime_context import get_parallel
-from sglang.srt.utils import add_prefix, get_bool_env_var, is_npu, set_weight_attrs
+from sglang.srt.utils import add_prefix, is_npu, set_weight_attrs
 
 _is_npu = is_npu()
-_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
-_tgemm = None
-if _use_aiter:
-    from aiter.tuned_gemm import tgemm
-
-    _tgemm = tgemm
 
 if TYPE_CHECKING:
     from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
@@ -287,10 +281,13 @@ def create_paged_compressor_data(
 
     def get_raw_loc(positions: torch.Tensor) -> torch.Tensor:
         positions = positions.masked_fill(positions < 0, 0)
-        loc = req_to_token[req_pool_indices, positions]
-        swa_loc = token_to_kv_pool.translate_loc_from_full_to_swa(loc)
-        swa_pages = swa_loc // swa_page_size
-        state_loc = swa_pages * ring_size + swa_loc % ring_size
+        if compress_ratio == 128:
+            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)
+            swa_pages = swa_loc // swa_page_size
+            state_loc = swa_pages * ring_size + swa_loc % ring_size
         return (state_loc // compress_ratio).to(torch.int32)
 
     is_overlap = is_overlap_compress(compress_ratio)
@@ -423,12 +420,7 @@ class Compressor(MultiPlatformOp):
         return ret
 
     def compute_kv_score(self, x: torch.Tensor, forward_batch: ForwardBatch):
-        if _tgemm is not None and not envs.SGLANG_OPT_USE_COMPRESSOR_V2.get():
-            # v1 compress goes through fused_compress_triton, which promotes
-            # bf16->fp32 internally, so skip the .float() cast.
-            kv_score = _tgemm.mm(x, self.wkv_gate.weight, otype=x.dtype)
-        else:
-            kv_score = linear_bf16_fp32(x, self.wkv_gate.weight)
+        kv_score = linear_bf16_fp32(x, self.wkv_gate.weight)
 
         # CUDA path: delegate to backend
         if dsa_use_prefill_cp(forward_batch):
@@ -487,9 +479,3 @@ class Compressor(MultiPlatformOp):
             )
 
         return get_attn_backend().forward_compress(self, x, forward_batch)
-
-
-if _is_hip and not envs.SGLANG_OPT_USE_COMPRESSOR_V2.get():
-    from sglang.srt.layers.attention.dsv4.compress_hip import (  # noqa: F811
-        CompressorHip as Compressor,
-    )

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

@@ -731,7 +731,7 @@ def create_paged_compressor_data(
             seq_lens=seq_lens_planner,
             extend_lens=extend_lens_planner,
             req_to_token=req_to_token,
-            full_to_swa=full_to_swa,
+            full_to_state=full_to_swa,
             swa_page_size=swa_page_size,
             ring_size=ring_size,
             num_q_tokens=num_q_tokens,
@@ -742,7 +742,7 @@ def create_paged_compressor_data(
             compress_ratio=compress_ratio,
             req_pool_indices=req_pool_indices_i64,
             req_to_token=req_to_token,
-            full_to_swa=full_to_swa,
+            full_to_state=full_to_swa,
             seq_lens=seq_lens.to(torch.int64),
             swa_page_size=swa_page_size,
             ring_size=ring_size,
@@ -763,8 +763,6 @@ def _create_online_paged_compressor_data(
     num_q_tokens: Optional[int],
     online_state_slot_offset: int = 0,
 ) -> CompressMetadata:
-    swa_page_size = int(token_to_kv_pool.swa_page_size)
-    full_to_swa = token_to_kv_pool.full_to_swa_index_mapping.detach()
     req_pool_indices = req_pool_indices.to(torch.int64)
 
     if is_prefill:
@@ -787,9 +785,7 @@ def _create_online_paged_compressor_data(
             extend_lens=extend_lens_planner,
             req_pool_indices=req_pool_indices,
             req_to_token=req_to_token,
-            full_to_swa=full_to_swa,
             num_q_tokens=int(num_q_tokens_planner),
-            swa_page_size=swa_page_size,
             use_cuda_graph=use_prefill_cuda_graph,
             state_slot_offset=online_state_slot_offset,
         )
@@ -798,7 +794,5 @@ def _create_online_paged_compressor_data(
             seq_lens=seq_lens.to(torch.int64),
             req_pool_indices=req_pool_indices,
             req_to_token=req_to_token,
-            full_to_swa=full_to_swa,
-            swa_page_size=swa_page_size,
             state_slot_offset=online_state_slot_offset,
         )

python/sglang/srt/managers/schedule_policy.py

@@ -185,7 +185,7 @@ class SchedulePolicy:
             and get_global_server_args().disaggregation_mode != "decode"
         ):
             for r in waiting_queue:
-                match_prefix_for_req(self.tree_cache, r)
+                match_prefix_for_req(self.tree_cache, r, include_req=True)
 
         if self.policy == CacheAgnosticPolicy.FCFS:
             if self.enable_priority_scheduling:
@@ -260,7 +260,9 @@ class SchedulePolicy:
         for r in waiting_queue:
             prefix_ids = r.origin_input_ids + r.output_ids
             extra_key = r.extra_key
-            match_result = match_prefix_for_req(self.tree_cache, r, prefix_ids)
+            match_result = match_prefix_for_req(
+                self.tree_cache, r, prefix_ids, include_req=True
+            )
 
             # NOTE(sang): This logic is for in-batch prefix caching;
             # If there are more than 1 request that have small matching prefix from

python/sglang/srt/mem_cache/deepseek_v4_compress_state.py

@@ -197,6 +197,13 @@ class CompressStatePool:
         state_loc = torch.where(swa_loc < 0, -1, state_loc)
         return state_loc
 
+    def translate_from_req_position_to_state_loc(
+        self, req_pool_indices: torch.Tensor, positions: torch.Tensor
+    ) -> torch.Tensor:
+        state_loc = req_pool_indices * self.ring_size + positions % self.ring_size
+        state_loc = torch.where(positions < 0, -1, state_loc)
+        return state_loc
+
     def get_state_by_state_loc(self, state_loc: torch.Tensor) -> KVAndScore:
         return self.kv_score_buffer[state_loc]
 

python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py

@@ -6,7 +6,11 @@ from typing import List, Literal, NamedTuple, Optional, Tuple
 
 import torch
 
-from sglang.jit_kernel.dsv4 import fused_k_norm_rope_flashmla, fused_store_cache
+from sglang.jit_kernel.dsv4 import (
+    clear_unaccepted_c128_draft_states,
+    fused_k_norm_rope_flashmla,
+    fused_store_cache,
+)
 from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.dsa import index_buf_accessor
@@ -494,15 +498,27 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
         self.c4_logical_size = c4_logical_size
         self.c128_size = c128_size
         self.c4_state_pool_size = c4_state_pool_size
+        c128_ring_size = self.get_ring_size(128)
+        if ONLINE_C128:
+            # Request-scoped online C128 state is indexed by req_pool_idx.
+            # PD decode can allocate pre-transfer slots beyond
+            # max_num_reqs, so size to the actual req_to_token row count.
+            c128_state_pool_size = max(c128_state_pool_size, self.num_req_slots)
+        else:
+            # Offline C128 keeps a per-request raw state ring.
+            c128_state_pool_size = max(
+                c128_state_pool_size, self.num_req_slots * c128_ring_size
+            )
         self.c128_state_pool_size = c128_state_pool_size
         self.c4_state_dtype = c4_state_dtype
         self.c128_state_dtype = c128_state_dtype
         self.compression_ratios = compression_ratios
         self.online_mtp_max_draft_tokens = online_mtp_max_draft_tokens
+        self.online_c128_state_num_req_slots = c128_state_pool_size
         self.online_c128_mtp_pending_seq_lens: Optional[torch.Tensor] = None
         if ONLINE_C128 and envs.SGLANG_EXPERIMENTAL_ONLINE_C128_MTP.get():
             self.online_c128_mtp_pending_seq_lens = torch.empty(
-                max_num_reqs, dtype=torch.int64, device=device
+                self.online_c128_state_num_req_slots, dtype=torch.int64, device=device
             )
 
         # Determine this PP stage's absolute layer range
@@ -603,11 +619,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
                 global_page_size=page_size,
             )
 
-        indexer_size = (
-            self.c4_logical_size
-            if (not _is_hip or envs.SGLANG_OPT_USE_COMPRESSOR_V2.get())
-            else c4_size
-        )
+        indexer_size = self.c4_logical_size
         self.c4_indexer_kv_pool = self._make_indexer_pool(
             indexer_size,
             c4_page_size,
@@ -732,6 +744,8 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
             for pool in pools:
                 if pool is None:
                     continue
+                if pool.ratio == 128:
+                    continue
                 t = pool.kv_score_buffer.kv_score
                 assert t.ndim == 2, f"expected 2D buffer, got {t.ndim}D"
                 data_ptrs.append(t.data_ptr())
@@ -740,6 +754,22 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
 
         return data_ptrs, data_lens, item_lens
 
+    def get_c128_state_buf_infos(
+        self,
+    ) -> Tuple[List[int], List[int], List[int]]:
+        data_ptrs: List[int] = []
+        data_lens: List[int] = []
+        item_lens: List[int] = []
+        for pool in self.compress_state_pools:
+            if pool is None or pool.ratio != 128:
+                continue
+            t = pool.kv_score_buffer.kv_score
+            assert t.ndim == 2, f"expected 2D buffer, got {t.ndim}D"
+            data_ptrs.append(t.data_ptr())
+            data_lens.append(t.nbytes)
+            item_lens.append(t[0].nbytes if ONLINE_C128 else t[0].nbytes * 128)
+        return data_ptrs, data_lens, item_lens
+
     def _make_kv_pool(
         self,
         *,
@@ -910,10 +940,57 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
                 return int(pool.online_mtp_max_draft_tokens)
         return 0
 
+    def get_online_c128_state_num_req_slots(self) -> int:
+        return self.online_c128_state_num_req_slots
+
     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_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:
+            if pool is None or pool.ratio != 128:
+                continue
+
+            state = pool.kv_score_buffer.kv_score
+            if ONLINE_C128:
+                row = state[req_pool_idx]
+                head_dim = row.shape[-1] // 3
+                row[:head_dim].fill_(float("-inf"))
+                row[head_dim:].zero_()
+            else:
+                start = req_pool_idx * pool.ring_size
+                rows = state[start : start + pool.ring_size]
+                half = rows.shape[-1] // 2
+                rows[:, :half].zero_()
+                rows[:, half:].fill_(float("-inf"))
+
+    def clear_unaccepted_c128_draft_states(
+        self,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        accept_lens: torch.Tensor,
+        num_draft_tokens: int,
+    ) -> None:
+        """Clear offline C128 ring slots written for rejected speculative tokens."""
+        if ONLINE_C128 or num_draft_tokens <= 1 or req_pool_indices.numel() == 0:
+            return
+
+        bs = req_pool_indices.numel()
+        for pool in self.compress_state_pools:
+            if pool is None or pool.ratio != 128:
+                continue
+
+            clear_unaccepted_c128_draft_states(
+                pool.kv_score_buffer.kv_score,
+                req_pool_indices,
+                seq_lens,
+                accept_lens,
+                ring_size=pool.ring_size,
+                num_draft_tokens=num_draft_tokens,
+            )
+
     def get_indexer_compress_states(self, layer_id: int) -> CompressStatePool:
         self.wait_layer_transfer(layer_id)
         indexer_compress_state_pool = self.indexer_compress_state_pools[layer_id]

python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py

@@ -988,7 +988,7 @@ class ModelRunnerKVCacheMixin:
                     self.token_to_kv_pool_allocator,
                 )
                 assert isinstance(swa_allocator, SWATokenToKVPoolAllocator)
-                self.token_to_kv_pool.full_to_swa_index_mapping = (
+                self.token_to_kv_pool.register_mapping(
                     swa_allocator.full_to_swa_index_mapping
                 )
 
@@ -1133,6 +1133,7 @@ class ModelRunnerKVCacheMixin:
         config.max_running_requests = self._resolve_max_num_reqs(
             config.max_total_num_tokens
         )
+        config = configurator.finalize_with_max_running_requests(config)
         config.mem_fraction_static = self.server_args.mem_fraction_static
         return config
 

python/sglang/srt/model_executor/pool_configurator.py

@@ -35,6 +35,7 @@ from sglang.srt.mem_cache.deepseek_v4_memory_pool import get_compress_state_ring
 from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
 from sglang.srt.utils.common import (
     ceil_align,
+    ceil_div,
     is_float4_e2m1fn_x2,
     spec_decode_alloc_len_per_request,
 )
@@ -108,6 +109,11 @@ class MemoryPoolConfigurator:
         """Constraint path: recalculate pool sizes from a constrained max_tokens."""
         raise NotImplementedError
 
+    def finalize_with_max_running_requests(
+        self, config: MemoryPoolConfig
+    ) -> MemoryPoolConfig:
+        return config
+
 
 class DefaultPoolConfigurator(MemoryPoolConfigurator):
     """Configurator for standard models: MHA, MLA, DSA, FP4.
@@ -517,6 +523,7 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
         self.qk_nope_head_dim = cfg.qk_nope_head_dim
         self.qk_rope_head_dim = cfg.qk_rope_head_dim
         self.indexer_head_dim = cfg.index_head_dim
+        self.context_len = mr.model_config.context_len
         # PP-local slice; matches DeepSeekV4TokenToKVPool's stage_ratios.
         self.compression_ratios = cfg.compress_ratios[mr.start_layer : mr.end_layer]
         if mr.pp_size > 1:
@@ -531,6 +538,15 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
         self.online_c128_mtp_max_draft_tokens = (
             mr.server_args.max_speculative_num_draft_tokens or 0
         )
+        self.requested_max_running_requests_per_worker = (
+            mr.server_args.max_running_requests // mr.dp_size
+            if mr.server_args.max_running_requests is not None
+            else None
+        )
+        self.disaggregation_mode = mr.server_args.disaggregation_mode
+        self.disaggregation_decode_extra_slots = (
+            mr.server_args.disaggregation_decode_extra_slots or 0
+        )
         if mr.enable_hisparse:
             from sglang.srt.mem_cache.sparsity import parse_hisparse_config
 
@@ -613,9 +629,10 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
         c4_indexer_state_bytes = 2 * 2 * self.indexer_head_dim * c4_state_dtype_size
 
         c4_state_ratio = self.c4_ring_size / self.swa_page_size
-        c128_state_ratio = self.c128_ring_size / self.swa_page_size
-        if c128_online and envs.SGLANG_EXPERIMENTAL_ONLINE_C128_MTP.get():
-            c128_state_ratio *= 1 + self.online_c128_mtp_max_draft_tokens
+        # C128 state is request-scoped and is finalized after
+        # max_running_requests is known, so it should not scale with
+        # full-token capacity here.
+        c128_state_ratio = 0
 
         c4_frac = 1 / (4 * self.c4_shrink_factor)
         return (
@@ -624,10 +641,7 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
             + 1 / 128 * kv_bytes * self.num_layers_ca128
             + 1 / 4 * indexer_bytes * self.num_layers_ca4
             + self.swa_ratio * c4_state_ratio * c4_state_bytes * self.num_layers_ca4
-            + self.swa_ratio
-            * c128_state_ratio
-            * c128_state_bytes
-            * self.num_layers_ca128
+            + c128_state_ratio * c128_state_bytes * self.num_layers_ca128
             + self.swa_ratio
             * c4_state_ratio
             * c4_indexer_state_bytes
@@ -643,9 +657,49 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
             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,
-            c128_state_pool_size=swa_tokens // self.swa_page_size * self.c128_ring_size,
+            c128_state_pool_size=0,
         )
 
+    def _get_num_req_slots(self, max_running_requests: int) -> int:
+        if self.disaggregation_mode == "decode":
+            return max_running_requests + self.disaggregation_decode_extra_slots + 1
+        return max_running_requests + 1
+
+    def _get_c128_state_fixed_bytes(self, max_running_requests: int) -> int:
+        if self.num_layers_ca128 == 0:
+            return 0
+
+        _, c128_state_dtype_size = _get_dsv4_compress_state_dtype_sizes()
+        attn_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+        num_req_slots = self._get_num_req_slots(max_running_requests)
+
+        if envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get():
+            state_rows = num_req_slots + self.c128_ring_size + 1
+            state_rows *= 1 + self.online_c128_mtp_max_draft_tokens
+            state_last_dim = 3 * attn_head_dim
+        else:
+            state_pool_size = num_req_slots * self.c128_ring_size
+            state_rows = state_pool_size + self.c128_ring_size + 1
+            state_rows = ceil_div(state_rows, 128) * 128
+            state_last_dim = 2 * attn_head_dim
+
+        return (
+            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:
+        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
+            )
+
+        estimated = int(token_capacity / 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)
+
     def _to_config(self, sizes: _DSV4PoolSizes) -> MemoryPoolConfig:
         full = sizes.full_max_total_num_tokens
         swa = sizes.swa_max_total_num_tokens
@@ -666,6 +720,17 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
             c128_state_pool_size=sizes.c128_state_pool_size,
         )
 
+    def finalize_with_max_running_requests(
+        self, config: MemoryPoolConfig
+    ) -> MemoryPoolConfig:
+        assert config.max_running_requests is not None
+        num_req_slots = self._get_num_req_slots(config.max_running_requests)
+        if envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get():
+            config.c128_state_pool_size = num_req_slots
+        else:
+            config.c128_state_pool_size = num_req_slots * self.c128_ring_size
+        return config
+
     def calculate_pool_sizes(
         self, available_bytes: int, page_size: int
     ) -> MemoryPoolConfig:
@@ -673,12 +738,25 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
             page_size % 128 == 0
         ), "page_size must be multiple of 128 for compressed attention"
 
-        full_token = int(available_bytes / self.bytes_per_full_token)
+        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)
+            )
+
+        available_bytes_for_tokens = max(available_bytes - c128_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"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"full_token={sizes.full_max_total_num_tokens}"
         )
         return self._to_config(sizes)

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -1544,6 +1544,18 @@ class EAGLEWorkerV2(BaseSpecWorker):
             accept_index,
         ) = eagle_sample(verify_input, batch, logits_output, vocab_mask)
         new_seq_lens = batch.seq_lens + accept_lens
+        clear_unaccepted_c128 = getattr(
+            self.token_to_kv_pool_allocator.get_kvcache(),
+            "clear_unaccepted_c128_draft_states",
+            None,
+        )
+        if clear_unaccepted_c128 is not None and not batch.forward_mode.is_idle():
+            clear_unaccepted_c128(
+                batch.req_pool_indices,
+                batch.seq_lens,
+                accept_lens,
+                self.speculative_num_draft_tokens,
+            )
 
         # Update mamba state for hybrid GDN models after verification
         commit_mamba_states_after_verify(

test/registered/jit/benchmark/bench_online_c128_mtp.py

@@ -43,7 +43,6 @@ class BenchmarkCase:
     seq_lens: torch.Tensor
     req_pool_indices: torch.Tensor
     req_to_token: torch.Tensor
-    full_to_swa: torch.Tensor
     ape: torch.Tensor
     state: torch.Tensor
     layer_bs: int
@@ -86,10 +85,6 @@ def make_case(batch_size: int, num_verify_tokens: int) -> BenchmarkCase:
     req_to_token = make_req_to_token(batch_size, max_seq_len, num_chunks)
 
     num_full_locs = batch_size * num_chunks
-    full_to_swa = (
-        torch.arange(num_full_locs, dtype=torch.int64, device=DEFAULT_DEVICE)
-        * SWA_PAGE_SIZE
-    )
 
     state_slot_stride = num_full_locs
     state = torch.empty(
@@ -112,7 +107,6 @@ def make_case(batch_size: int, num_verify_tokens: int) -> BenchmarkCase:
         seq_lens=seq_lens,
         req_pool_indices=req_pool_indices,
         req_to_token=req_to_token,
-        full_to_swa=full_to_swa,
         ape=ape,
         state=state,
         layer_bs=batch_size,
@@ -127,11 +121,9 @@ def call_write_prefix(module, case: BenchmarkCase) -> None:
         case.seq_lens,
         case.req_pool_indices,
         case.req_to_token,
-        case.full_to_swa,
         case.ape,
         case.state,
         case.layer_bs,
-        SWA_PAGE_SIZE,
         case.num_verify_tokens,
         case.state_slot_stride,
     )
@@ -153,8 +145,10 @@ def call_write_prefix(module, case: BenchmarkCase) -> None:
 def benchmark(
     batch_size: int, num_verify_tokens: int, launch_mode: str
 ) -> tuple[float, float, float]:
-    module = _jit_online_c128_mtp_module(HEAD_DIM)
     case = make_case(batch_size, num_verify_tokens)
+    module = _jit_online_c128_mtp_module(
+        HEAD_DIM, case.seq_lens.dtype, case.req_pool_indices.dtype
+    )
     fn = lambda: call_write_prefix(module, case)
 
     if launch_mode == "cuda_graph":

test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py

@@ -126,6 +126,7 @@ class TestUnifiedDeepSeekV4FlashHiCacheL3(AccuracyTwoPassMixin, CustomTestCase):
 
     l3_prefetch_page_size = 256
     l3_prefetch_prompt_pages = 4
+    max_running_requests = 4
 
     @classmethod
     def setUpClass(cls):
@@ -164,6 +165,8 @@ class TestUnifiedDeepSeekV4FlashHiCacheL3(AccuracyTwoPassMixin, CustomTestCase):
                 "file",
                 "--swa-full-tokens-ratio",
                 "0.25",
+                "--max-running-requests",
+                str(cls.max_running_requests),
             ],
             env={
                 "SGLANG_DSV4_FP4_EXPERTS": "0",

test/registered/unit/disaggregation/test_disaggregation_wire.py

@@ -9,6 +9,7 @@ from sglang.srt.disaggregation.common.utils import (
     unpack_int_lists,
     unpack_list_of_buffers,
 )
+from sglang.srt.disaggregation.utils import get_dsv4_c128_state_indices
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=2, suite="base-a-test-cpu")
@@ -90,5 +91,31 @@ class TestGroupConcurrentContiguous(unittest.TestCase):
             group_concurrent_contiguous(self._arr([1, 2, 3]), self._arr([1, 2]))
 
 
+class TestDSV4C128StateIndices(unittest.TestCase):
+    def test_online_aligned_boundary_has_no_partial_state(self):
+        np.testing.assert_array_equal(
+            get_dsv4_c128_state_indices(7, 256, online=True, ring_size=1),
+            np.empty((0,), dtype=np.int32),
+        )
+
+    def test_online_partial_boundary_uses_request_slot(self):
+        np.testing.assert_array_equal(
+            get_dsv4_c128_state_indices(7, 257, online=True, ring_size=1),
+            np.array([7], dtype=np.int32),
+        )
+
+    def test_offline_aligned_boundary_has_no_partial_state(self):
+        np.testing.assert_array_equal(
+            get_dsv4_c128_state_indices(7, 256, online=False, ring_size=128),
+            np.empty((0,), dtype=np.int32),
+        )
+
+    def test_offline_partial_boundary_uses_request_local_page(self):
+        np.testing.assert_array_equal(
+            get_dsv4_c128_state_indices(7, 129, online=False, ring_size=256),
+            np.array([15], dtype=np.int32),
+        )
+
+
 if __name__ == "__main__":
     unittest.main()