diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c128_online_v2.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/c128_online_v2.cuh index 02f263078..534527ba3 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/c128_online_v2.cuh +++ b/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(params.seq_lens[idx]); const auto rid = params.req_pool_indices[idx]; - const int32_t chunk_start = static_cast((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(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(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() .with_device(device_) .verify(req_to_token); - TensorMatcher({-1}) // - .with_dtype() - .with_device(device_) - .verify(full_to_swa); TensorMatcher({B, sizeof(DecodePlan)}) // .with_dtype() .with_device(device_) .verify(plan_d_dev_); - RuntimeCheck(swa_page_size > 0); RuntimeCheck(state_slot_offset >= 0); const auto batch_size = static_cast(B.unwrap()); @@ -610,9 +598,7 @@ inline void plan_online_decode( .seq_lens = static_cast(seq_lens.data_ptr()), .req_pool_indices = static_cast(req_pool_indices.data_ptr()), .req_to_token = static_cast(req_to_token.data_ptr()), - .full_to_swa = static_cast(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(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(params.full_to_swa[full_loc]); - const int32_t main_slot = swa_loc / params.swa_page_size; + const int32_t main_slot = static_cast(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() .with_device(device_) .verify(req_to_token); - TensorMatcher({-1}) // - .with_dtype() - .with_device(device_) - .verify(full_to_swa); TensorMatcher({N, sizeof(CompressPlan)}) // .with_dtype() .with_device(cpu) @@ -878,9 +852,7 @@ inline OnlinePrefillPlan plan_online_prefill( .plan_w = plan_w_dev_ptr, .req_pool_indices = static_cast(req_pool_indices.data_ptr()), .req_to_token = static_cast(req_to_token.data_ptr()), - .full_to_swa = static_cast(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, diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh index 4e2f2ed28..b8ddb3787 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh +++ b/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(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(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(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(rid * params.ring_size + position % params.ring_size); + }; const auto seq_len = static_cast(params.seq_ptr[idx]); const auto position_1 = static_cast(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(compute_loc(state_loc_1)); + read_page_0 = static_cast(compute_loc(state_loc_0) / params.compress_ratio); + read_page_1 = static_cast(write_loc / params.compress_ratio); + } params.plan_d[idx] = { .seq_len = static_cast(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(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(seq_len), .write_loc = write_loc, @@ -443,7 +466,9 @@ using PrefillPlan = tvm::ffi::Tuple; * 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; 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() .with_device(device_) - .verify(full_to_swa); + .verify(full_to_state); TensorMatcher({B}) // .with_dtype() .with_device(cpu_or_gpu) @@ -500,7 +525,7 @@ inline PrefillPlan plan_compress_prefill( const auto ext_ptr = static_cast(extend_lens.data_ptr()); const auto rid_ptr = static_cast(req_pool_indices.data_ptr()); const auto r2t_ptr = static_cast(req_to_token.data_ptr()); - const auto f2s_ptr = static_cast(full_to_swa.data_ptr()); + const auto f2s_ptr = static_cast(full_to_state.data_ptr()); const auto batch_size = static_cast(B.unwrap()); constexpr auto kMaxTokens = static_cast(std::numeric_limits::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() .with_device(device_) - .verify(full_to_swa); + .verify(full_to_state); TensorMatcher({B}) // .with_dtype() .with_device(device_) @@ -670,7 +695,7 @@ inline tvm::ffi::Tensor plan_compress_decode( .plan_d = static_cast(D.data_ptr()), .rid_ptr = static_cast(req_pool_indices.data_ptr()), .r2t_ptr = static_cast(req_to_token.data_ptr()), - .f2s_ptr = static_cast(full_to_swa.data_ptr()), + .f2s_ptr = static_cast(full_to_state.data_ptr()), .seq_ptr = static_cast(seq_lens.data_ptr()), .stride_r2t = req_to_token.size(1), .batch_size = batch_size, diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/online_c128_mtp.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/online_c128_mtp.cuh index db1b94d89..72eeea609 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/online_c128_mtp.cuh +++ b/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(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(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(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(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 +template struct OnlineC128MTPWritePrefixKernel { - template 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(seq_lens.data_ptr()), .req_pool_indices = static_cast(req_pool_indices.data_ptr()), .req_to_token = static_cast(req_to_token.data_ptr()), - .full_to_swa = static_cast(full_to_swa.data_ptr()), .ape = static_cast(ape.data_ptr()), .state = static_cast(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(); TensorMatcher({-1, kHeadDim * 2}).with_dtype().with_device(device).verify(kv_score_input); - TensorMatcher({-1}).with_dtype(seq_dtype).with_device(device).verify(seq_lens); - TensorMatcher({-1}).with_dtype(req_dtype).with_device(device).verify(req_pool_indices); + TensorMatcher({-1}).with_dtype().with_device(device).verify(seq_lens); + TensorMatcher({-1}).with_dtype().with_device(device).verify(req_pool_indices); TensorMatcher({-1, -1}).with_dtype().with_device(device).verify(req_to_token); - TensorMatcher({-1}).with_dtype().with_device(device).verify(full_to_swa); TensorMatcher({128, kHeadDim}).with_dtype().with_device(device).verify(ape); TensorMatcher({-1, kHeadDim * 3}).with_dtype().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()) { - if (req_dtype.is_type()) { - launch( - 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( - 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()) { - launch( - 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( - 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 +template struct OnlineC128MTPMarkPendingKernel { - template 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(); - TensorMatcher({-1}).with_dtype(seq_dtype).with_device(device).verify(seq_lens); - TensorMatcher({-1}).with_dtype(req_dtype).with_device(device).verify(req_pool_indices); + TensorMatcher({-1}).with_dtype().with_device(device).verify(seq_lens); + TensorMatcher({-1}).with_dtype().with_device(device).verify(req_pool_indices); TensorMatcher({-1}).with_dtype().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()) { - if (req_dtype.is_type()) { - launch(seq_lens, req_pool_indices, pending_seq_lens, bs, max_num_reqs, device.unwrap()); - } else { - launch(seq_lens, req_pool_indices, pending_seq_lens, bs, max_num_reqs, device.unwrap()); - } - } else { - if (req_dtype.is_type()) { - launch(seq_lens, req_pool_indices, pending_seq_lens, bs, max_num_reqs, device.unwrap()); - } else { - launch(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 +template struct OnlineC128MTPCommitPendingKernel { - template 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(cur_seq_lens.data_ptr()), .cur_req_pool_indices = static_cast(cur_req_pool_indices.data_ptr()), .req_to_token = static_cast(req_to_token.data_ptr()), - .full_to_swa = static_cast(full_to_swa.data_ptr()), .pending_seq_lens = static_cast(pending_seq_lens.data_ptr()), .state = static_cast(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(); - TensorMatcher({-1}).with_dtype(seq_dtype).with_device(device).verify(cur_seq_lens); - TensorMatcher({-1}).with_dtype(req_dtype).with_device(device).verify(cur_req_pool_indices); + TensorMatcher({-1}).with_dtype().with_device(device).verify(cur_seq_lens); + TensorMatcher({-1}).with_dtype().with_device(device).verify(cur_req_pool_indices); TensorMatcher({-1, -1}).with_dtype().with_device(device).verify(req_to_token); - TensorMatcher({-1}).with_dtype().with_device(device).verify(full_to_swa); TensorMatcher({-1}).with_dtype().with_device(device).verify(pending_seq_lens); TensorMatcher({-1, kHeadDim * 3}).with_dtype().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()) { - if (req_dtype.is_type()) { - launch( - 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( - 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()) { - launch( - 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( - 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()); } }; diff --git a/python/sglang/jit_kernel/dsv4/__init__.py b/python/sglang/jit_kernel/dsv4/__init__.py index a9952ebfe..21a9c9b28 100644 --- a/python/sglang/jit_kernel/dsv4/__init__.py +++ b/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", diff --git a/python/sglang/jit_kernel/dsv4/attn.py b/python/sglang/jit_kernel/dsv4/attn.py index 87a265ecf..784711498 100644 --- a/python/sglang/jit_kernel/dsv4/attn.py +++ b/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 diff --git a/python/sglang/jit_kernel/dsv4/c128_cleanup.py b/python/sglang/jit_kernel/dsv4/c128_cleanup.py new file mode 100644 index 000000000..38b3a0ab2 --- /dev/null +++ b/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, + ) diff --git a/python/sglang/jit_kernel/dsv4/compress.py b/python/sglang/jit_kernel/dsv4/compress.py index ca8d2c342..7c41d5f74 100644 --- a/python/sglang/jit_kernel/dsv4/compress.py +++ b/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), ) diff --git a/python/sglang/jit_kernel/dsv4/online_c128_mtp.py b/python/sglang/jit_kernel/dsv4/online_c128_mtp.py index 2716d5641..d04c34f13 100644 --- a/python/sglang/jit_kernel/dsv4/online_c128_mtp.py +++ b/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() diff --git a/python/sglang/jit_kernel/tests/deepseek_v4/common.py b/python/sglang/jit_kernel/tests/deepseek_v4/common.py index 7288caa44..982745c67 100644 --- a/python/sglang/jit_kernel/tests/deepseek_v4/common.py +++ b/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, diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index 8945f1d49..1cbe4e7f7 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/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 diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 607f0b735..83fcde3f9 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/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 diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 993c04c16..d0293fba3 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/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) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 6054f3f7a..3e8153e96 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/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 ) diff --git a/python/sglang/srt/disaggregation/mori/conn.py b/python/sglang/srt/disaggregation/mori/conn.py index 7e29f408e..ef2d16a8e 100644 --- a/python/sglang/srt/disaggregation/mori/conn.py +++ b/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( diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 6c807c590..30965357e 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/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 diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index f23519560..8355a5097 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/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) diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 068b1a19f..9ffceda77 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/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() diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 7a47e0b19..303e914ed 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 8b504153f..f9571fe77 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index 88dbb5af3..152426276 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/dsv4/compress_hip.py b/python/sglang/srt/layers/attention/dsv4/compress_hip.py index 8c6b7df9b..004d49c5b 100644 --- a/python/sglang/srt/layers/attention/dsv4/compress_hip.py +++ b/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) diff --git a/python/sglang/srt/layers/attention/dsv4/compressor.py b/python/sglang/srt/layers/attention/dsv4/compressor.py index 53f824b08..5c685df35 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor.py +++ b/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, - ) diff --git a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py index 7a2d26617..908545ff9 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py +++ b/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, ) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index b37a78329..edd3ce869 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/deepseek_v4_compress_state.py b/python/sglang/srt/mem_cache/deepseek_v4_compress_state.py index d3e3d25df..735eb0430 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_compress_state.py +++ b/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] diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index 595412ff8..d8b0d756a 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/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] diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index 63877b48f..a87bf4b5f 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/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 diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index bc4fc63bd..50c37332e 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/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) diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index e5a1c151e..44686829c 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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( diff --git a/test/registered/jit/benchmark/bench_online_c128_mtp.py b/test/registered/jit/benchmark/bench_online_c128_mtp.py index cab1c159b..fdf2c7eb5 100644 --- a/test/registered/jit/benchmark/bench_online_c128_mtp.py +++ b/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": diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py index 6c35238d3..3e449bf13 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py +++ b/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", diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index 76ee5874b..76208a5cc 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/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()