From 187120ef540d1db5c0028b01c0032ac122e68f94 Mon Sep 17 00:00:00 2001 From: yhyang201 Date: Thu, 18 Jun 2026 09:39:52 +0800 Subject: [PATCH] Revert "DeepSeek-V4 Online Compress support MTP (#26471)" This reverts commit 063ab89ac1685e99e2252fdc95eea444dc83635e. --- .../csrc/deepseek_v4/c128_online_v2.cuh | 75 +-- .../csrc/deepseek_v4/online_c128_mtp.cuh | 537 ------------------ python/sglang/jit_kernel/dsv4/compress.py | 17 +- .../sglang/jit_kernel/dsv4/online_c128_mtp.py | 256 --------- python/sglang/srt/environ.py | 1 - .../layers/attention/deepseek_v4_backend.py | 170 +----- .../layers/attention/dsv4/compressor_v2.py | 19 +- .../mem_cache/deepseek_v4_compress_state.py | 13 +- .../srt/mem_cache/deepseek_v4_memory_pool.py | 29 +- .../model_runner_kv_cache_mixin.py | 3 - .../srt/model_executor/pool_configurator.py | 33 +- .../jit/benchmark/bench_online_c128_mtp.py | 166 ------ 12 files changed, 45 insertions(+), 1274 deletions(-) delete mode 100644 python/sglang/jit_kernel/csrc/deepseek_v4/online_c128_mtp.cuh delete mode 100644 python/sglang/jit_kernel/dsv4/online_c128_mtp.py delete mode 100644 test/registered/jit/benchmark/bench_online_c128_mtp.py 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 02f2630780fc..71e600dc39d3 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 @@ -237,8 +237,8 @@ SGL_DEVICE void c128_prefill_segment_softmax( /// Reads optional prior state from `read_page_0` (-1 = none), emits compressed /// kv to `kv_compressed_output[plan_id]` (compact). /// `kWrite=true` (write pass) : handles trailing partial segments. -/// Reads optional prior state from `read_page_1` (-1 = fallback to -/// `read_page_0`), writes new running state to `read_page_0`. +/// Reads optional prior state from `read_page_0` (-1 = none), writes new +/// running state to `read_page_1`. template __global__ __launch_bounds__(kPrefillBlockSize, 2) // void flash_c128_online_prefill_v2(const __grid_constant__ Compress128OnlinePrefillParams params) { @@ -279,13 +279,13 @@ __global__ __launch_bounds__(kPrefillBlockSize, 2) // constexpr int64_t kElementSize = kHeadDim * 2; // | kv | score | - // `j` below is a chunk-local offset. Convert it to the ragged-input row by - // anchoring on the last token in this segment: ragged_id - pos_in_chunk_end + 1 + j. + // The plan stores last-token coordinates; segment start is recoverable as + // ragged_id - window_len + 1. const uint32_t window_len = plan.buffer_len; const uint32_t position = plan.seq_len - 1; const uint32_t pos_in_chunk_end = (position % 128u) + 1u; // exclusive, in [1, 128] const uint32_t chunk_offset = pos_in_chunk_end - window_len; // in [0, 127] - const int32_t chunk_start_ragged = static_cast(plan.ragged_id) - static_cast(pos_in_chunk_end) + 1; + const int32_t segment_start_ragged = static_cast(plan.ragged_id) - static_cast(position % 128u); // --- Stage 1: load kv / score / bias for this warp's 8 chunk positions. PrefillStorage kv[kElementsPerWarp]; @@ -297,8 +297,7 @@ __global__ __launch_bounds__(kPrefillBlockSize, 2) // for (uint32_t i = 0; i < kElementsPerWarp; ++i) { const uint32_t j = i + warp_offset; if (j >= chunk_offset && j < pos_in_chunk_end) { - const int32_t ragged_id = chunk_start_ragged + static_cast(j); - const auto kv_src_ptr = kv_score_input + ragged_id * kElementSize + split_offset; + const auto kv_src_ptr = kv_score_input + (segment_start_ragged + j) * kElementSize + split_offset; const auto score_src_ptr = kv_src_ptr + kHeadDim; const auto bias_src_ptr = score_bias_base + j * kHeadDim + split_offset; kv[i].load(kv_src_ptr, lane_id); @@ -335,10 +334,9 @@ __global__ __launch_bounds__(kPrefillBlockSize, 2) // out_max_vec.load(seg_max, lane_id); out_sum_vec.load(seg_sum, lane_id); - const int32_t read_page = plan.read_page_1 >= 0 ? plan.read_page_1 : plan.read_page_0; - if (chunk_offset != 0 && read_page >= 0) { + if (chunk_offset != 0 && plan.read_page_0 >= 0) { // Combine with prior partial state for this slot. - const auto buf_load = kv_score_buffer + read_page * (kHeadDim * 3) + split_offset; + const auto buf_load = kv_score_buffer + plan.read_page_0 * (kHeadDim * 3) + split_offset; PrefillStorage buf_max_vec, buf_sum_vec, buf_kv_vec; buf_max_vec.load(buf_load + 0 * kHeadDim, lane_id); buf_sum_vec.load(buf_load + 1 * kHeadDim, lane_id); @@ -469,8 +467,8 @@ struct FlashCompress128OnlineKernel { .with_device(device_) .verify(ape); - // Both compress and write segments use PlanC layout. Stage 1 stores the - // committed-bank load slot in read_page_1 and the write slot in read_page_0. + // Both compress and write segments use PlanC layout. plan_c uses + // read_page_1=-1 (unused); plan_w uses read_page_1=store_slot. const auto plan_c = compress::verify_plan_c(plan_c_, C, device_); const auto plan_w = compress::verify_plan_c(plan_w_, W, device_); const auto device = device_.unwrap(); @@ -514,9 +512,8 @@ struct FlashCompress128OnlineKernel { // GPU kernel that finalizes `read_page_0` to `req_to_token[rid][chunk_start]`, // so the slot tensors never leave GPU memory. The online state pool keeps // a single in-progress chunk per request, so each segment's load and -// store slot collapse to one value (the slot for the segment's own chunk). -// For online-c128 MTP, stage 1 keeps that write slot in `read_page_0` and -// stores the committed-bank load slot in `read_page_1`. +// store slot collapse to one value (the slot for the segment's own chunk), +// and `read_page_1` is unused. // =========================================================================== namespace host::compress { @@ -536,7 +533,6 @@ struct OnlineDecodePlanParams { 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; }; @@ -548,7 +544,7 @@ __global__ void plan_c128_online_decode_kernel(const OnlineDecodePlanParams para 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 = swa_loc / params.swa_page_size; params.plan_d[idx] = DecodePlan{ .seq_len = seq_len, .write_loc = slot, @@ -568,8 +564,7 @@ inline void plan_online_decode( 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) { + const int32_t swa_page_size) { auto B = SymbolicSize{"batch_size"}; auto device_ = SymbolicDevice{}; device_.set_options(); @@ -596,7 +591,6 @@ inline void plan_online_decode( .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()); if (batch_size == 0) return; @@ -613,7 +607,6 @@ inline void plan_online_decode( .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, }; LaunchKernel(num_blocks, kBlockSize, device)(plan_c128_online_decode_kernel, params); @@ -661,7 +654,7 @@ inline std::tuple _plan_prefill_partial(const OnlinePrefillS .ragged_id = static_cast(last_ragged), .buffer_len = static_cast(seg_len), .read_page_0 = static_cast(i), // batch_id placeholder - .read_page_1 = -1, // filled by stage 1 with committed-bank slot + .read_page_1 = -1, // unused, kept so MSB layout is stable }; if (chunk_off + seg_len == 128u) { // close-chunk segment @@ -688,7 +681,6 @@ struct OnlinePrefillStage1Params { 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; }; @@ -701,16 +693,13 @@ __global__ void plan_c128_online_prefill_kernel(const OnlinePrefillStage1Params const bool is_compress = idx < params.num_c; CompressPlan* const plan_ptr = is_compress ? ¶ms.plan_c[idx] : ¶ms.plan_w[idx - params.num_c]; auto plan = *plan_ptr; - 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; - plan.read_page_0 = main_slot + params.state_slot_offset; - plan.read_page_1 = main_slot; + plan.read_page_0 = swa_loc / params.swa_page_size; *plan_ptr = plan; } @@ -726,9 +715,7 @@ inline OnlinePrefillPlan plan_online_prefill( 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) { + const int32_t swa_page_size) { auto B = SymbolicSize{"batch_size"}; auto N = SymbolicSize{"num_q_tokens"}; auto cpu = SymbolicDevice{}; @@ -763,7 +750,6 @@ inline OnlinePrefillPlan plan_online_prefill( .with_device(device_) .verify(plan_c_dev_) .verify(plan_w_dev_); - RuntimeCheck(state_slot_offset >= 0); const auto stage0_params = OnlinePrefillStage0Params{ .plan_c = static_cast(plan_c_pin.data_ptr()), @@ -798,8 +784,6 @@ inline OnlinePrefillPlan plan_online_prefill( } const auto [num_c, num_w] = _plan_prefill_partial(stage0_params); - const auto num_c_padded = use_cuda_graph ? static_cast(N.unwrap()) : num_c; - const auto num_w_padded = use_cuda_graph ? static_cast(N.unwrap()) : num_w; if (kGuard) { // Verify stage 0 wrote ONLY to the [0, num_c*16) and [0, num_w*16) prefix. @@ -837,19 +821,7 @@ inline OnlinePrefillPlan plan_online_prefill( auto* const plan_c_dev_ptr = static_cast(plan_c_dev_.data_ptr()); auto* const plan_w_dev_ptr = static_cast(plan_w_dev_.data_ptr()); - if (use_cuda_graph) { - const auto kInvalidPlan = CompressPlan::invalid(); - auto* const plan_c_pin_ptr = static_cast(plan_c_pin.data_ptr()); - auto* const plan_w_pin_ptr = static_cast(plan_w_pin.data_ptr()); - for (const auto i : irange(num_c, num_c_padded)) { - plan_c_pin_ptr[i] = kInvalidPlan; - } - for (const auto i : irange(num_w, num_w_padded)) { - plan_w_pin_ptr[i] = kInvalidPlan; - } - } - - if (const auto total = num_c_padded + num_w_padded) { + if (const auto total = num_c + num_w) { const auto stream = LaunchKernel::resolve_device(device); // SGLANG_DEBUG_C128_ONLINE_SYNC_H2D=1 forces a synchronous H2D copy. static const bool kSyncH2D = []() { @@ -870,8 +842,8 @@ inline OnlinePrefillPlan plan_online_prefill( RuntimeDeviceCheck(::cudaMemcpyAsync(dst, src, bytes, ::cudaMemcpyHostToDevice, stream)); } }; - if (num_c_padded) copy_to_device(plan_c_dev_ptr, plan_c_pin.data_ptr(), num_c_padded); - if (num_w_padded) copy_to_device(plan_w_dev_ptr, plan_w_pin.data_ptr(), num_w_padded); + if (num_c) copy_to_device(plan_c_dev_ptr, plan_c_pin.data_ptr(), num_c); + if (num_w) copy_to_device(plan_w_dev_ptr, plan_w_pin.data_ptr(), num_w); const auto stage1_params = OnlinePrefillStage1Params{ .plan_c = plan_c_dev_ptr, @@ -881,15 +853,14 @@ inline OnlinePrefillPlan plan_online_prefill( .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, + .num_c = num_c, + .num_w = num_w, }; constexpr uint32_t kBlockSize = 128; const auto num_blocks = host::div_ceil(total, kBlockSize); LaunchKernel(num_blocks, kBlockSize, device)(plan_c128_online_prefill_kernel, stage1_params); } - return OnlinePrefillPlan{num_c_padded, num_w_padded}; + return OnlinePrefillPlan{num_c, num_w}; } } // namespace host::compress 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 deleted file mode 100644 index db1b94d89344..000000000000 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/online_c128_mtp.cuh +++ /dev/null @@ -1,537 +0,0 @@ -#pragma once - -#include -#include - -#include - -#include -#include - -#include - -namespace { - -SGL_DEVICE int64_t clamp_accept_len(int64_t delta, int64_t max_accept) { - if (delta < 0) return 0; - return delta < max_accept ? delta : max_accept; -} - -template -struct OnlineC128MTPWritePrefixParams { - const float* __restrict__ kv_score_input; - 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; - int64_t req_to_token_stride_b; - 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; -}; - -template -struct OnlineC128MTPMarkPendingParams { - const TSeq* __restrict__ seq_lens; - const TReq* __restrict__ req_pool_indices; - int64_t* __restrict__ pending_seq_lens; - int64_t bs; - int64_t max_num_reqs; -}; - -template -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; -}; - -__global__ void online_c128_mtp_clear_all_pending_kernel(int64_t* pending_seq_lens, int64_t max_num_reqs) { - const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (idx < max_num_reqs) pending_seq_lens[idx] = -1; -} - -template -__global__ void online_c128_mtp_mark_pending_kernel(const OnlineC128MTPMarkPendingParams params) { - const int64_t bid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (bid >= params.bs) return; - const int64_t req = static_cast(params.req_pool_indices[bid]); - if (req >= 0 && req < params.max_num_reqs) { - params.pending_seq_lens[req] = static_cast(params.seq_lens[bid]); - } -} - -template -__global__ void online_c128_mtp_commit_pending_kernel(const OnlineC128MTPCommitPendingParams params) { - const int64_t bid = static_cast(blockIdx.x); - if (bid >= params.cur_bs) return; - - const int64_t req = static_cast(params.cur_req_pool_indices[bid]); - if (req < 0 || req >= params.max_num_reqs) return; - const int64_t old_seq = params.pending_seq_lens[req]; - if (old_seq < 0) return; - - const int64_t cur_seq = static_cast(params.cur_seq_lens[bid]); - const int64_t accept = clamp_accept_len(cur_seq - old_seq, params.num_verify_tokens); - if (accept <= 0) return; - - 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 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; - - for (int64_t d = static_cast(threadIdx.x); d < kHeadDim * 3; d += blockDim.x) { - dst[d] = src[d]; - } -} - -template -__global__ void online_c128_mtp_write_prefix_kernel(const OnlineC128MTPWritePrefixParams params) { - const int64_t bid = static_cast(blockIdx.x); - if (bid >= params.layer_bs) return; - - const int64_t seq_before = static_cast(params.seq_lens[bid]); - const int64_t req_idx = static_cast(params.req_pool_indices[bid]); - const int64_t start_pos = seq_before & 127; - const bool has_partial = seq_before > 0 && start_pos != 0; - - 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; - } - - const int64_t d = static_cast(threadIdx.x); - float run_max = 0.0f; - float run_sum = 0.0f; - float run_kv = 0.0f; - if (has_partial) { - const float* const init = params.state + init_slot * params.state_stride_b; - run_max = init[d]; - run_sum = init[kHeadDim + d]; - run_kv = init[kHeadDim * 2 + d]; - } - - constexpr int kMaxVerifyTokens = 8; - float kv_steps[kMaxVerifyTokens]; - float score_steps[kMaxVerifyTokens]; - -#pragma unroll - for (int step = 0; step < kMaxVerifyTokens; ++step) { - if (step >= params.num_verify_tokens) break; - - const int64_t pos = (start_pos + step) & 127; - const float* const kv = params.kv_score_input + (bid * params.num_verify_tokens + step) * params.kv_score_stride_b; - kv_steps[step] = kv[d]; - score_steps[step] = kv[kHeadDim + d] + params.ape[pos * params.ape_stride_r + d]; - } - -#pragma unroll - for (int step = 0; step < kMaxVerifyTokens; ++step) { - if (step >= params.num_verify_tokens) break; - - const int64_t pos = (start_pos + step) & 127; - const float kv_step = kv_steps[step]; - const float score_step = score_steps[step]; - if (pos == 0) { - run_kv = kv_step; - run_max = score_step; - run_sum = 1.0f; - } else { - const float new_max = fmaxf(run_max, score_step); - const float old_sum_scaled = run_sum * __expf(run_max - new_max); - const float new_exp = __expf(score_step - new_max); - const float new_sum = old_sum_scaled + new_exp; - run_kv = (run_kv * old_sum_scaled + kv_step * new_exp) / new_sum; - run_max = new_max; - run_sum = new_sum; - } - - 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; - float* const out = params.state + slot * params.state_stride_b; - out[d] = run_max; - out[kHeadDim + d] = run_sum; - out[kHeadDim * 2 + d] = run_kv; - } - - if (pos == 127) { - run_kv = 0.0f; - run_max = 0.0f; - run_sum = 0.0f; - } - } -} - -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) { - using namespace host; - - const auto params = OnlineC128MTPWritePrefixParams{ - .kv_score_input = static_cast(kv_score_input.data_ptr()), - .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), - .req_to_token_stride_b = req_to_token.stride(0), - .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, - }; - - static_assert(kHeadDim == 512, "online c128 MTP write-prefix only supports head_dim=512"); - constexpr uint32_t kThreads = static_cast(kHeadDim); - LaunchKernel(static_cast(layer_bs), kThreads, device)( - online_c128_mtp_write_prefix_kernel, params); - } - - static void - run(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) { - 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, -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); - - if (layer_bs <= 0) return; - RuntimeCheck(num_verify_tokens > 0 && num_verify_tokens <= 8, "unsupported num_verify_tokens=", num_verify_tokens); - RuntimeCheck(state_slot_stride > 0, "state_slot_stride must be positive"); - RuntimeCheck(layer_bs <= seq_lens.shape()[0], "layer_bs exceeds seq_lens rows"); - 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()); - } - } - } -}; - -template -struct OnlineC128MTPMarkPendingKernel { - template - static void launch( - tvm::ffi::TensorView seq_lens, - tvm::ffi::TensorView req_pool_indices, - tvm::ffi::TensorView pending_seq_lens, - int64_t bs, - int64_t max_num_reqs, - DLDevice device) { - using namespace host; - - const auto params = OnlineC128MTPMarkPendingParams{ - .seq_lens = static_cast(seq_lens.data_ptr()), - .req_pool_indices = static_cast(req_pool_indices.data_ptr()), - .pending_seq_lens = static_cast(pending_seq_lens.data_ptr()), - .bs = bs, - .max_num_reqs = max_num_reqs, - }; - - constexpr uint32_t kThreads = 256; - const uint32_t clear_blocks = host::div_ceil(static_cast(max_num_reqs), kThreads); - LaunchKernel(clear_blocks, kThreads, device)( - online_c128_mtp_clear_all_pending_kernel, params.pending_seq_lens, max_num_reqs); - const uint32_t mark_blocks = host::div_ceil(static_cast(bs), kThreads); - LaunchKernel(mark_blocks, kThreads, device)(online_c128_mtp_mark_pending_kernel, params); - } - - static void - run(tvm::ffi::TensorView seq_lens, - tvm::ffi::TensorView req_pool_indices, - tvm::ffi::TensorView pending_seq_lens, - int64_t bs, - 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(pending_seq_lens); - - if (bs <= 0) return; - RuntimeCheck(bs <= seq_lens.shape()[0], "bs exceeds seq_lens rows"); - 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()); - } - } - } -}; - -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, - DLDevice device) { - using namespace host; - - const auto params = OnlineC128MTPCommitPendingParams{ - .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, - }; - - constexpr uint32_t kThreads = 256; - LaunchKernel(static_cast(cur_bs), kThreads, device)( - online_c128_mtp_commit_pending_kernel, params); - } - - static void - 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, -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); - - if (cur_bs <= 0) return; - RuntimeCheck(num_verify_tokens > 0 && num_verify_tokens <= 8, "unsupported num_verify_tokens=", num_verify_tokens); - RuntimeCheck(state_slot_stride > 0, "state_slot_stride must be positive"); - RuntimeCheck(cur_bs <= cur_seq_lens.shape()[0], "cur_bs exceeds seq_lens rows"); - 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()); - } - } - } -}; - -} // namespace diff --git a/python/sglang/jit_kernel/dsv4/compress.py b/python/sglang/jit_kernel/dsv4/compress.py index ca8d2c3421ce..bf14a8da5870 100644 --- a/python/sglang/jit_kernel/dsv4/compress.py +++ b/python/sglang/jit_kernel/dsv4/compress.py @@ -159,7 +159,6 @@ def generate_online( 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]) module = _jit_compress_128_online_module(512) @@ -169,13 +168,7 @@ def generate_online( device=req_pool_indices.device, ) module.plan_decode( - seq_lens, - req_pool_indices, - req_to_token, - full_to_swa, - plan_d, - swa_page_size, - int(state_slot_offset), + seq_lens, req_pool_indices, req_to_token, full_to_swa, plan_d, swa_page_size ) return CompressorDecodePlan(128, plan_d) @@ -277,11 +270,9 @@ def generate_online( full_to_swa: torch.Tensor, num_q_tokens: int, swa_page_size: int, - use_cuda_graph: bool = False, - state_slot_offset: int = 0, ) -> CompressorPrefillPlan: - seq_lens_cpu = seq_lens.detach().to(torch.int64).cpu() - extend_lens_cpu = extend_lens.detach().to(torch.int64).cpu() + seq_lens_cpu = seq_lens.to(torch.int64) + extend_lens_cpu = extend_lens.to(torch.int64) rid_i64 = req_pool_indices.to(torch.int64) r2t_i32 = req_to_token.to(torch.int32) f2s_i64 = full_to_swa.to(torch.int64) @@ -304,8 +295,6 @@ def generate_online( plan_c_dev, plan_w_dev, int(swa_page_size), - int(state_slot_offset), - bool(use_cuda_graph), ) return CompressorPrefillPlan( 128, diff --git a/python/sglang/jit_kernel/dsv4/online_c128_mtp.py b/python/sglang/jit_kernel/dsv4/online_c128_mtp.py deleted file mode 100644 index d93c6f120816..000000000000 --- a/python/sglang/jit_kernel/dsv4/online_c128_mtp.py +++ /dev/null @@ -1,256 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, List, Optional - -import torch -from tvm_ffi.module import Module - -from sglang.jit_kernel.dsv4.utils import make_name -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args -from sglang.srt.environ import envs - - -@cache_once -def _jit_online_c128_mtp_module(head_dim: int) -> Module: - args = make_cpp_args(head_dim) - return load_jit( - make_name(f"online_c128_mtp_{head_dim}"), - *args, - cuda_files=["deepseek_v4/online_c128_mtp.cuh"], - cuda_wrappers=[ - ("write_prefix_states", f"OnlineC128MTPWritePrefixKernel<{args}>::run"), - ("mark_pending", f"OnlineC128MTPMarkPendingKernel<{args}>::run"), - ("commit_pending", f"OnlineC128MTPCommitPendingKernel<{args}>::run"), - ], - extra_cuda_cflags=["-use_fast_math"], - ) - - -@dataclass -class _OnlineC128LayerRuntime: - head_dim: int - main_state: torch.Tensor - state_slot_offset: int - - -@dataclass -class _OnlineC128VerifyContext: - req_pool_indices: torch.Tensor - seq_lens: torch.Tensor - - -class OnlineC128MTPController: - def __init__(self, backend: Any): - self.backend = backend - self._verify_ctx: Optional[_OnlineC128VerifyContext] = None - self._layer_runtimes: Optional[List[_OnlineC128LayerRuntime]] = None - - def enabled(self) -> bool: - return ( - envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get() - and envs.SGLANG_EXPERIMENTAL_ONLINE_C128_MTP.get() - and self.backend.mtp_enabled - ) - - def state_slot_offset(self) -> int: - if not self.enabled(): - return 0 - return self.backend.token_to_kv_pool.get_online_c128_mtp_state_slot_offset() - - def begin_verify( - self, - req_pool_indices: torch.Tensor, - seq_lens: torch.Tensor, - ) -> None: - if not self.enabled(): - self.clear() - return - - self._verify_ctx = _OnlineC128VerifyContext( - req_pool_indices=req_pool_indices.detach(), - seq_lens=seq_lens.detach(), - ) - head_dim = self._head_dim() - 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( - 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, - ) - - def clear(self) -> None: - self._verify_ctx = None - - def prepare_forward( - self, - logical_forward_mode, - req_pool_indices: torch.Tensor, - seq_lens: torch.Tensor, - *, - verify_bs: Optional[int] = None, - ) -> int: - if not self.enabled(): - self.clear() - return 0 - if logical_forward_mode is None or logical_forward_mode.is_idle(): - self.clear() - return 0 - - active_req_pool_indices = req_pool_indices - active_seq_lens = seq_lens - if logical_forward_mode.is_target_verify(): - if verify_bs is None: - verify_bs = req_pool_indices.shape[0] - active_req_pool_indices = req_pool_indices[:verify_bs] - active_seq_lens = seq_lens[:verify_bs] - if verify_bs == 0: - self.clear() - return 0 - - self.commit_pending( - req_pool_indices=active_req_pool_indices, - seq_lens=active_seq_lens, - ) - if not logical_forward_mode.is_target_verify(): - return 0 - - self.begin_verify( - req_pool_indices=active_req_pool_indices, - seq_lens=active_seq_lens, - ) - return self.state_slot_offset() - - def write_prefix_states( - self, - layer_id: int, - compressor: Any, - kv_score_input: torch.Tensor, - logical_forward_mode, - ) -> None: - if ( - not self.enabled() - or logical_forward_mode is None - or not logical_forward_mode.is_target_verify() - or compressor.is_in_indexer - or compressor.ratio != 128 - or kv_score_input.numel() == 0 - ): - return - - ctx = self._active_ctx() - num_verify_tokens = self._num_verify_tokens() - if ctx is None or num_verify_tokens == 0: - return - - token_to_kv_pool = self.backend.token_to_kv_pool - head_dim = compressor.head_dim - state_pool = token_to_kv_pool.get_attention_compress_states(layer_id) - total_bs = kv_score_input.numel() // (num_verify_tokens * head_dim * 2) - layer_bs = min(ctx.seq_lens.shape[0], ctx.req_pool_indices.shape[0], total_bs) - if layer_bs <= 0: - return - - _jit_online_c128_mtp_module(head_dim).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, - ) - - def commit_pending( - self, - req_pool_indices: torch.Tensor, - seq_lens: torch.Tensor, - ) -> None: - if self._verify_ctx is None: - return - if not self.enabled(): - self.clear() - return - if req_pool_indices.numel() == 0 or seq_lens.numel() == 0: - return - - num_verify_tokens = self._num_verify_tokens() - if num_verify_tokens == 0: - self.clear() - return - - backend = self.backend - token_to_kv_pool = backend.token_to_kv_pool - pending_seq_lens = token_to_kv_pool.get_online_c128_mtp_pending_seq_lens() - 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( - 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, - ) - - self.clear() - - def _num_verify_tokens(self) -> int: - if not self.enabled(): - return 0 - num_verify_tokens = int(self.backend.speculative_num_draft_tokens) - max_draft_tokens = ( - self.backend.token_to_kv_pool.get_online_c128_mtp_max_draft_tokens() - ) - return num_verify_tokens if 0 < num_verify_tokens <= max_draft_tokens else 0 - - def _active_ctx(self) -> Optional[_OnlineC128VerifyContext]: - ctx = self._verify_ctx - if ( - ctx is None - or ctx.seq_lens.numel() == 0 - or ctx.req_pool_indices.numel() == 0 - ): - return None - return ctx - - def _head_dim(self) -> Optional[int]: - for runtime in self._iter_layer_runtimes(): - return runtime.head_dim - return None - - def _iter_layer_runtimes(self): - if self._layer_runtimes is None: - runtimes = [] - token_to_kv_pool = self.backend.token_to_kv_pool - for layer in self.backend.model_runner.model.model.layers: - attn = getattr(layer, "self_attn", None) - compressor = getattr(attn, "compressor", None) - if compressor is None or compressor.ratio != 128: - continue - state_pool = token_to_kv_pool.get_attention_compress_states( - compressor.layer_id - ) - runtimes.append( - _OnlineC128LayerRuntime( - head_dim=compressor.head_dim, - main_state=state_pool.kv_score_buffer.kv_score, - state_slot_offset=state_pool.online_mtp_state_slot_offset, - ) - ) - self._layer_runtimes = runtimes - return iter(self._layer_runtimes) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index f67f4da2bced..4178d2ef79b4 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -787,7 +787,6 @@ class Envs: SGLANG_OPT_USE_AITER_INDEXER = EnvBool(False) SGLANG_OPT_USE_JIT_INDEXER_METADATA = EnvBool(True) SGLANG_OPT_USE_ONLINE_COMPRESS = EnvBool(False) - SGLANG_EXPERIMENTAL_ONLINE_C128_MTP = EnvBool(False) SGLANG_DSV4_COMPRESS_STATE_DTYPE = EnvStr("float32") SGLANG_OPT_USE_COMPRESSOR_V2 = EnvBool(True) SGLANG_FP8_PAGED_MQA_LOGITS_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 54f502442045..65f44fe0f25c 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -35,7 +35,6 @@ create_paged_compressor_data, ) -from sglang.jit_kernel.dsv4.online_c128_mtp import OnlineC128MTPController from sglang.srt.layers.attention.dsv4.dequant_k_cache import ( dequantize_k_cache_paged, ) @@ -80,37 +79,6 @@ PAGE_INDEX_ALIGNED_SIZE = 64 -def _get_logical_forward_mode(forward_batch: ForwardBatch) -> ForwardMode: - # IDLE is a real per-DP-rank mode. Do not let a stale _original_forward_mode - # from a reused/padded ForwardBatch turn an empty rank into TARGET_VERIFY. - if forward_batch.forward_mode.is_idle(): - return forward_batch.forward_mode - return ( - getattr(forward_batch, "_original_forward_mode", None) - or forward_batch.forward_mode - ) - - -def _get_target_verify_bs(forward_batch: ForwardBatch) -> int: - actual_forward_mode = getattr( - forward_batch, "actual_forward_mode", forward_batch.forward_mode - ) - if actual_forward_mode.is_idle(): - return 0 - - spec_info = getattr(forward_batch, "spec_info", None) - draft_token_num = getattr(spec_info, "draft_token_num", 0) - draft_token = getattr(spec_info, "draft_token", None) - if draft_token is None: - return forward_batch.batch_size - if draft_token_num <= 0: - return 0 - draft_count = len(draft_token) - if draft_count % draft_token_num != 0: - return 0 - return draft_count // draft_token_num - - T = TypeVar("T", bound=Optional[torch.Tensor]) @@ -134,13 +102,6 @@ def _create_dummy_paged_compress_data(compress_ratio: int): return None -def _copy_or_replace(dst, src): - if dst is not None and src is not None: - dst.copy_(src) - return dst - return src - - @dataclass class DSV4AttnMetadata: page_size: int @@ -419,8 +380,6 @@ class DSV4RawVerifyMetadata: out_cache_loc: torch.Tensor extend_seq_lens: Optional[torch.Tensor] = None - seq_lens_cpu: Optional[List[int]] = None - c128_compress_metadata: Optional[FusedCompressMetadata] = None def copy_(self, other: DSV4RawVerifyMetadata): self.req_pool_indices.copy_(other.req_pool_indices) @@ -428,10 +387,6 @@ def copy_(self, other: DSV4RawVerifyMetadata): self.out_cache_loc.copy_(other.out_cache_loc) self.extend_seq_lens = other.extend_seq_lens - self.seq_lens_cpu = other.seq_lens_cpu - self.c128_compress_metadata = _copy_or_replace( - self.c128_compress_metadata, other.c128_compress_metadata - ) @dataclass @@ -476,7 +431,6 @@ def __init__( speculative_num_steps=0, ): super().__init__() - self.model_runner = model_runner self.device = torch.device(model_runner.device) head_dim = model_runner.model_config.head_dim assert ( @@ -518,41 +472,11 @@ def __init__( DSV4RawVerifyMetadata, DSV4RawDecodeMetadata, ] = None - self.online_c128_mtp = OnlineC128MTPController(self) def _move_to_device(self, x: List[int]) -> torch.Tensor: pin_tensor = torch.tensor(x, dtype=torch.int32, pin_memory=True) return pin_tensor.to(self.device, non_blocking=True) - def _make_target_verify_c128_metadata( - self, - req_pool_indices: torch.Tensor, - seq_lens: torch.Tensor, - seq_lens_cpu: List[int], - extend_seq_lens: torch.Tensor, - use_prefill_cuda_graph: bool, - online_c128_state_slot_offset: int, - ) -> Optional[FusedCompressMetadata]: - if not self.online_c128_mtp.enabled(): - return None - - num_draft_tokens = self.speculative_num_draft_tokens - seq_lens_cpu = [int(x) + num_draft_tokens for x in seq_lens_cpu] - extend_lens_cpu = [num_draft_tokens] * len(seq_lens_cpu) - return create_paged_compressor_data( - compress_ratio=128, - is_prefill=True, - token_to_kv_pool=self.token_to_kv_pool, - req_to_token=self.req_to_token, - req_pool_indices=req_pool_indices, - seq_lens=seq_lens + self.speculative_num_draft_tokens, - seq_lens_cpu=seq_lens_cpu, - extend_lens=extend_seq_lens, - extend_lens_cpu=extend_lens_cpu, - use_prefill_cuda_graph=use_prefill_cuda_graph, - online_state_slot_offset=online_c128_state_slot_offset, - ) - def init_forward_metadata_indexer(self, core_attn_metadata: DSV4AttnMetadata): return PagedIndexerMetadata( page_size=self.page_size, @@ -618,7 +542,6 @@ def init_forward_metadata_prefill( extend_start_loc: Optional[torch.Tensor] = None, need_compress: bool = True, use_prefill_cuda_graph: bool = False, - online_c128_state_slot_offset: int = 0, ) -> DSV4Metadata: seq_lens_casual, req_pool_indices_repeated = self.expand_prefill_casually( num_tokens=num_tokens, @@ -668,7 +591,6 @@ def create(compress_ratio: Literal[4, 128]): extend_lens_cpu=None, use_prefill_cuda_graph=True, num_q_tokens=out_cache_loc.shape[0], - online_state_slot_offset=online_c128_state_slot_offset, ) return create_paged_compressor_data( compress_ratio=compress_ratio, @@ -681,7 +603,6 @@ def create(compress_ratio: Literal[4, 128]): extend_lens=extend_seq_lens, extend_lens_cpu=extend_seq_lens_cpu, use_prefill_cuda_graph=use_graph_plan, - online_state_slot_offset=online_c128_state_slot_offset, ) c4_compress_metadata = create(compress_ratio=4) @@ -698,18 +619,11 @@ def init_forward_metadata_target_verify( max_seq_len: int, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor, - seq_lens_cpu: Optional[torch.Tensor] = None, out_cache_loc: Optional[torch.Tensor] = None, use_prefill_cuda_graph: bool = False, - online_c128_state_slot_offset: int = 0, ) -> Union[DSV4Metadata, DSV4RawVerifyMetadata]: if envs.SGLANG_PREP_IN_CUDA_GRAPH.get(): assert out_cache_loc is not None - seq_lens_cpu_list = ( - seq_lens.detach().cpu().tolist() - if seq_lens_cpu is None - else seq_lens_cpu.tolist() - ) if not hasattr(self, "extend_seq_lens_buffer"): self.extend_seq_lens_buffer = torch.tensor( [self.speculative_num_draft_tokens] * 1025, device=self.device @@ -721,15 +635,6 @@ def init_forward_metadata_target_verify( seq_lens=seq_lens, out_cache_loc=out_cache_loc, extend_seq_lens=extend_seq_lens, - seq_lens_cpu=seq_lens_cpu_list, - c128_compress_metadata=self._make_target_verify_c128_metadata( - req_pool_indices, - seq_lens, - seq_lens_cpu_list, - extend_seq_lens, - use_prefill_cuda_graph, - online_c128_state_slot_offset, - ), ) else: seq_lens_cpu = seq_lens.tolist() @@ -740,7 +645,6 @@ def init_forward_metadata_target_verify( seq_lens_cpu=seq_lens_cpu, out_cache_loc=out_cache_loc, use_prefill_cuda_graph=use_prefill_cuda_graph, - online_c128_state_slot_offset=online_c128_state_slot_offset, ) def init_forward_metadata_target_verify_old( @@ -751,7 +655,6 @@ def init_forward_metadata_target_verify_old( seq_lens_cpu: Optional[List[int]] = None, out_cache_loc: Optional[torch.Tensor] = None, use_prefill_cuda_graph: bool = False, - online_c128_state_slot_offset: int = 0, ) -> DSV4Metadata: batch_size = len(seq_lens) seq_lens = seq_lens + self.speculative_num_draft_tokens @@ -773,13 +676,10 @@ def init_forward_metadata_target_verify_old( extend_start_loc=None, need_compress=True, use_prefill_cuda_graph=use_prefill_cuda_graph, - online_c128_state_slot_offset=online_c128_state_slot_offset, ) def make_forward_metadata_from_raw_verify( - self, - raw_metadata: DSV4RawVerifyMetadata, - online_c128_state_slot_offset: int = 0, + self, raw_metadata: DSV4RawVerifyMetadata ) -> DSV4Metadata: req_pool_indices = raw_metadata.req_pool_indices seq_lens = raw_metadata.seq_lens @@ -788,7 +688,6 @@ def make_forward_metadata_from_raw_verify( bs, num_draft_tokens = len(seq_lens), self.speculative_num_draft_tokens seq_lens = seq_lens + self.speculative_num_draft_tokens extend_seq_lens = raw_metadata.extend_seq_lens - assert extend_seq_lens is not None seq_lens_casual, req_pool_indices_repeated = ( self.expand_extend_with_same_length( @@ -816,21 +715,16 @@ def make_forward_metadata_from_raw_verify( extend_lens_cpu=None, use_prefill_cuda_graph=True, num_q_tokens=num_draft_tokens * bs, - online_state_slot_offset=online_c128_state_slot_offset, ) - c128_compress_metadata = raw_metadata.c128_compress_metadata - if c128_compress_metadata is None: - c128_compress_metadata = create(compress_ratio=128) return DSV4Metadata( core_attn_metadata, indexer_metadata, c4_compress_metadata=create(compress_ratio=4), - c128_compress_metadata=c128_compress_metadata, + c128_compress_metadata=create(compress_ratio=128), ) def make_forward_metadata_from_raw_decode( - self, - raw_metadata: DSV4RawDecodeMetadata, + self, raw_metadata: DSV4RawDecodeMetadata ) -> DSV4Metadata: req_pool_indices = raw_metadata.req_pool_indices seq_lens = raw_metadata.seq_lens @@ -899,7 +793,6 @@ def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch) -> None: if isinstance(self.forward_metadata, DSV4RawVerifyMetadata): self.forward_metadata = self.make_forward_metadata_from_raw_verify( raw_metadata=self.forward_metadata, - online_c128_state_slot_offset=self.online_c128_mtp.state_slot_offset(), ) elif isinstance(self.forward_metadata, DSV4RawDecodeMetadata): self.forward_metadata = self.make_forward_metadata_from_raw_decode( @@ -996,11 +889,6 @@ def init_forward_metadata_out_graph( if bucket == _GraphBucket.DECODE_OR_IDLE: assert out_cache_loc is not None assert len(out_cache_loc.shape) == 1, f"{out_cache_loc.shape=}" - self.online_c128_mtp.prepare_forward( - actual_forward_mode, - req_pool_indices, - seq_lens, - ) out_cache_loc_padded = torch.nn.functional.pad( out_cache_loc, pad=(0, bs - len(out_cache_loc)), @@ -1014,13 +902,6 @@ def init_forward_metadata_out_graph( out_cache_loc=out_cache_loc_padded, ) elif bucket == _GraphBucket.TARGET_VERIFY: - verify_bs = _get_target_verify_bs(forward_batch) - if self.online_c128_mtp.enabled() and verify_bs == 0: - self.online_c128_mtp.clear() - self.forward_metadata = self.cuda_graph_metadata_of_bucket_and_bs[ - bucket - ][bs] - return assert out_cache_loc is not None num_tokens_v = self.speculative_num_draft_tokens * bs out_cache_loc_padded = torch.nn.functional.pad( @@ -1029,27 +910,14 @@ def init_forward_metadata_out_graph( mode="constant", value=0, ) - online_c128_state_slot_offset = self.online_c128_mtp.prepare_forward( - actual_forward_mode, - req_pool_indices, - seq_lens, - verify_bs=verify_bs, - ) temp_metadata = self.init_forward_metadata_target_verify( max_seq_len=chosen_max_seq_len, req_pool_indices=req_pool_indices, seq_lens=seq_lens, - seq_lens_cpu=seq_lens_cpu, out_cache_loc=out_cache_loc_padded, use_prefill_cuda_graph=True, - online_c128_state_slot_offset=online_c128_state_slot_offset, ) elif bucket == _GraphBucket.DRAFT_EXTEND: - self.online_c128_mtp.prepare_forward( - actual_forward_mode, - req_pool_indices, - seq_lens, - ) num_tokens_per_bs = self.draft_extend_num_tokens_per_bs if out_cache_loc is not None: # Pad the real write locations to the captured token count so @@ -1070,7 +938,6 @@ def init_forward_metadata_out_graph( use_prefill_cuda_graph=True, ) else: - self.online_c128_mtp.clear() raise NotImplementedError self.replay_cuda_graph_metadata_from( @@ -1090,9 +957,7 @@ def init_forward_metadata_out_graph( ) def init_forward_metadata(self, forward_batch: ForwardBatch) -> None: - logical_forward_mode = _get_logical_forward_mode(forward_batch) - if self.mtp_enabled and logical_forward_mode.is_idle(): - self.online_c128_mtp.clear() + if self.mtp_enabled and forward_batch.forward_mode.is_idle(): return self.forward_metadata = self._build_forward_metadata(forward_batch) @@ -1105,7 +970,6 @@ def _build_forward_metadata( max_seq_len_override: Optional[int] = None, use_prefill_cuda_graph: bool = False, ): - logical_forward_mode = _get_logical_forward_mode(forward_batch) req_pool_indices = forward_batch.req_pool_indices seq_lens = forward_batch.seq_lens.to(torch.int32) seq_lens_cpu = forward_batch.seq_lens_cpu @@ -1113,22 +977,13 @@ def _build_forward_metadata( assert self.swa_page_size % SWA_WINDOW == 0 and self.page_size % 128 == 0 assert seq_lens_cpu is not None - if max_seq_len_override is None: - max_seq_len_override = getattr(forward_batch, "max_seq_len_override", None) max_seq_len = ( int(seq_lens_cpu.max().item()) if max_seq_len_override is None else max_seq_len_override ) - verify_bs = _get_target_verify_bs(forward_batch) - online_c128_state_slot_offset = self.online_c128_mtp.prepare_forward( - logical_forward_mode, - req_pool_indices, - seq_lens, - verify_bs=verify_bs, - ) - if logical_forward_mode.is_decode_or_idle(): + if forward_batch.forward_mode.is_decode_or_idle(): # DSv4 bakes this step's KV write target (c4/c128) into metadata, # so slice the shared multi-step out_cache_loc now, not at forward time. out_cache_loc = forward_batch.out_cache_loc @@ -1145,16 +1000,14 @@ def _build_forward_metadata( seq_lens=seq_lens, out_cache_loc=out_cache_loc, ) - elif logical_forward_mode.is_target_verify(): + elif forward_batch.forward_mode.is_target_verify(): metadata = self.init_forward_metadata_target_verify( max_seq_len=max_seq_len, req_pool_indices=req_pool_indices, seq_lens=seq_lens, - seq_lens_cpu=seq_lens_cpu, out_cache_loc=forward_batch.out_cache_loc, - online_c128_state_slot_offset=online_c128_state_slot_offset, ) - elif logical_forward_mode.is_prefill(include_draft_extend_v2=True): + elif forward_batch.forward_mode.is_prefill(include_draft_extend_v2=True): extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu extend_seq_lens = forward_batch.extend_seq_lens assert ( @@ -1237,12 +1090,12 @@ def replay_cuda_graph_metadata_from( ], bucket: _GraphBucket, ) -> None: - bucket_metadata = self.cuda_graph_metadata_of_bucket_and_bs[bucket] - chosen_metadata = bucket_metadata.get(bs) - if chosen_metadata is None: - bucket_metadata[bs] = temp_metadata + if bs not in self.cuda_graph_metadata_of_bucket_and_bs[bucket]: + # First call (from capture): store the new metadata directly. + self.cuda_graph_metadata_of_bucket_and_bs[bucket][bs] = temp_metadata self.forward_metadata = temp_metadata return + chosen_metadata = self.cuda_graph_metadata_of_bucket_and_bs[bucket][bs] chosen_metadata.copy_(temp_metadata) self.forward_metadata = chosen_metadata @@ -1750,7 +1603,6 @@ def __init__( self, model_runner: ModelRunner, topk: int, speculative_num_steps: int ): super().__init__(model_runner) - self.model_runner = model_runner self.topk = topk self.speculative_num_steps = speculative_num_steps self.attn_backends: List[DeepseekV4AttnBackend] = [] diff --git a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py index 7a2d266170a5..3d5c4b16f79a 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py @@ -541,17 +541,6 @@ def forward_unified( use_fp4_indexer=use_fp4_indexer, bf16_store=bf16_store, ) - online_c128_mtp = getattr(self, "online_c128_mtp", None) - if online_c128_mtp is not None: - online_c128_mtp.write_prefix_states( - layer_id=layer_id, - compressor=compressor, - kv_score_input=kv_score_input, - logical_forward_mode=getattr( - forward_batch, "_original_forward_mode", None - ) - or forward_batch.forward_mode, - ) def _forward_unified_hip( self, @@ -684,7 +673,6 @@ def create_paged_compressor_data( extend_lens_cpu: Optional[List[int]] = None, use_prefill_cuda_graph: bool = False, num_q_tokens: Optional[int] = None, - online_state_slot_offset: int = 0, ) -> CompressMetadata: """Build the paged compress metadata (= the plan). @@ -703,7 +691,6 @@ def create_paged_compressor_data( extend_lens_cpu=extend_lens_cpu, use_prefill_cuda_graph=use_prefill_cuda_graph, num_q_tokens=num_q_tokens, - online_state_slot_offset=online_state_slot_offset, ) swa_page_size = token_to_kv_pool.swa_page_size @@ -761,8 +748,9 @@ def _create_online_paged_compressor_data( extend_lens_cpu: Optional[List[int]], use_prefill_cuda_graph: bool, num_q_tokens: Optional[int], - online_state_slot_offset: int = 0, ) -> CompressMetadata: + assert not use_prefill_cuda_graph, "online c128 doesn't support cuda graph" + 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) @@ -790,8 +778,6 @@ def _create_online_paged_compressor_data( 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, ) else: return CompressorDecodePlan.generate_online( @@ -800,5 +786,4 @@ def _create_online_paged_compressor_data( 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/mem_cache/deepseek_v4_compress_state.py b/python/sglang/srt/mem_cache/deepseek_v4_compress_state.py index f43415a5a49d..4543c48fd7ab 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_compress_state.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_compress_state.py @@ -88,29 +88,18 @@ def __init__( ratio: int, online: bool = False, swa_page_size: int = 0, - online_mtp_max_draft_tokens: int = 0, ): - self.ratio = ratio self.ring_size = ring_size self.swa_page_size = swa_page_size self.enable_memory_saver = enable_memory_saver - self.online_mtp_state_slot_offset = 0 - self.online_mtp_max_draft_tokens = 0 if online: assert ring_size == 1, "online compress requires ring_size=1" - self._logical_size = size + self.ring_size + 1 - if online_mtp_max_draft_tokens > 0: - # Bank 0 is the committed state. Banks 1..N cache per-draft - # prefix states for lazy commit after target verify. - self.online_mtp_max_draft_tokens = online_mtp_max_draft_tokens - self.online_mtp_state_slot_offset = self._logical_size - self._size = self._logical_size * (1 + self.online_mtp_max_draft_tokens) + self._size = size + self.ring_size + 1 last_dim = 3 * head_dim else: self._size = size + self.ring_size + 1 self._size = (self._size + ratio - 1) // ratio * ratio - self._logical_size = self._size last_dim = 2 * (1 + overlap) * head_dim if _is_hip: 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 6d13fb6be9f2..24cc211791ea 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -35,8 +35,7 @@ def get_compress_state_ring_size( # 128-slot ring buffer of raw tokens, so ring_size collapses to 1. Online # is incompatible with speculative decode for now. if compress_ratio == 128 and ONLINE_C128: - if is_speculative and not envs.SGLANG_EXPERIMENTAL_ONLINE_C128_MTP.get(): - raise AssertionError("online c128 does not support MTP") + assert not is_speculative, "online c128 does not support MTP" return 1 if is_speculative: return 16 if compress_ratio == 4 else 256 @@ -461,7 +460,6 @@ def __init__( start_layer: Optional[int] = None, end_layer: Optional[int] = None, enable_hisparse: bool = False, - online_mtp_max_draft_tokens: int = 0, num_req_slots: Optional[int] = None, ): super().__init__( @@ -498,12 +496,6 @@ def __init__( 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_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 - ) # Determine this PP stage's absolute layer range if ( @@ -769,9 +761,6 @@ def _init_paged_compress_states(self, enable_memory_saver: bool): ratio=ratio, online=(ratio == 128 and ONLINE_C128), swa_page_size=self.swa_page_size, - online_mtp_max_draft_tokens=( - self.online_mtp_max_draft_tokens if ratio == 128 else 0 - ), ) if ratio == 4: @@ -829,22 +818,6 @@ def get_attention_compress_states(self, layer_id: int) -> CompressStatePool: ), "Only c4/c128 layers have attention states." return compress_state_pool - def get_online_c128_mtp_state_slot_offset(self) -> int: - for pool in self.compress_state_pools: - if pool is not None and pool.ratio == 128: - return int(pool.online_mtp_state_slot_offset) - return 0 - - def get_online_c128_mtp_max_draft_tokens(self) -> int: - for pool in self.compress_state_pools: - if pool is not None and pool.ratio == 128: - return int(pool.online_mtp_max_draft_tokens) - return 0 - - 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 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 a86fd9a2ccb8..91233ee1f8ae 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 @@ -448,9 +448,6 @@ def _init_pools(self: ModelRunner): start_layer=self.start_layer, end_layer=self.end_layer, enable_hisparse=self.enable_hisparse, - online_mtp_max_draft_tokens=( - self.server_args.max_speculative_num_draft_tokens or 0 - ), ) elif current_platform.is_out_of_tree() and not self.mambaish_config: if self.use_mla_backend and is_dsa_model: diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index fb1a7f0956df..8443ae4d4162 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -365,9 +365,6 @@ def __init__(self, mr: ModelRunner): self.swa_page_size = cfg.window_size self.swa_ratio = mr.server_args.swa_full_tokens_ratio self.is_speculative = mr.server_args.speculative_algorithm is not None - self.online_c128_mtp_max_draft_tokens = ( - mr.server_args.max_speculative_num_draft_tokens or 0 - ) if mr.enable_hisparse: from sglang.srt.mem_cache.sparsity import parse_hisparse_config @@ -402,30 +399,10 @@ def __init__(self, mr: ModelRunner): # would need rollback / replay across draft and verify, which the # online path doesn't support yet. if envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get(): - allow_experimental_online_c128_mtp = ( - envs.SGLANG_EXPERIMENTAL_ONLINE_C128_MTP.get() - and mr.spec_algorithm.is_eagle() - ) - assert mr.spec_algorithm.is_none() or allow_experimental_online_c128_mtp, ( - "SGLANG_OPT_USE_ONLINE_COMPRESS does not support speculative decode " - "(MTP) yet, except the experimental EAGLE topk=1 path gated by " - "SGLANG_EXPERIMENTAL_ONLINE_C128_MTP=1" - ) - if allow_experimental_online_c128_mtp: - assert self.online_c128_mtp_max_draft_tokens > 0, ( - "SGLANG_EXPERIMENTAL_ONLINE_C128_MTP requires " - "speculative_num_draft_tokens to be set." - ) - logger.warning( - "DSV4 compressed attention: experimental online c128 + MTP enabled " - f"(EAGLE topk=1 only, " - f"draft_banks={self.online_c128_mtp_max_draft_tokens}). " - "Validate correctness carefully." - ) - else: - logger.info( - "DSV4 compressed attention: online c128 enabled (ring_size=1)" - ) + assert ( + mr.spec_algorithm.is_none() + ), "SGLANG_OPT_USE_ONLINE_COMPRESS does not support speculative decode (MTP) yet" + logger.info("DSV4 compressed attention: online c128 enabled (ring_size=1)") def _get_bytes_per_full_token(self) -> float: kv_bytes = self.qk_nope_head_dim + self.qk_rope_head_dim * 2 + 8 @@ -451,8 +428,6 @@ def _get_bytes_per_full_token(self) -> float: 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 c4_frac = 1 / (4 * self.c4_shrink_factor) return ( diff --git a/test/registered/jit/benchmark/bench_online_c128_mtp.py b/test/registered/jit/benchmark/bench_online_c128_mtp.py deleted file mode 100644 index 9f7540125ad7..000000000000 --- a/test/registered/jit/benchmark/bench_online_c128_mtp.py +++ /dev/null @@ -1,166 +0,0 @@ -"""Benchmark online c128 MTP write-prefix kernel.""" - -from __future__ import annotations - -import itertools -from dataclasses import dataclass - -import torch -import triton -import triton.testing - -from sglang.jit_kernel.benchmark.utils import ( - DEFAULT_DEVICE, - get_benchmark_range, - run_benchmark, - run_benchmark_no_cudagraph, -) -from sglang.jit_kernel.dsv4.online_c128_mtp import _jit_online_c128_mtp_module -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci(est_time=10, suite="base-b-kernel-benchmark-1-gpu-large") - -HEAD_DIM = 512 -STATE_DIM = HEAD_DIM * 3 -SWA_PAGE_SIZE = 128 - -BATCH_SIZE_RANGE = get_benchmark_range( - full_range=[1, 2, 4, 8, 16, 32, 64, 128, 256], - ci_range=[8, 64], -) -NUM_VERIFY_TOKENS_RANGE = get_benchmark_range( - full_range=[1, 4, 8], - ci_range=[8], -) -BENCHMARK_CONFIGS = list(itertools.product(BATCH_SIZE_RANGE, NUM_VERIFY_TOKENS_RANGE)) - - -@dataclass -class BenchmarkCase: - kv_score_input: torch.Tensor - 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 - num_verify_tokens: int - state_slot_stride: int - - -def round_up_div(x: int, y: int) -> int: - return (x + y - 1) // y - - -def make_seq_lens(batch_size: int, num_verify_tokens: int) -> torch.Tensor: - # Cover chunk positions around the interesting boundaries. This exercises - # both the has-partial path and the final_seq % 128 == 0 skip-write path. - offsets = torch.tensor([0, 1, 2, 63, 120, 126, 127], dtype=torch.int64) - seq_offsets = offsets[torch.arange(batch_size, dtype=torch.int64) % offsets.numel()] - base = 8 * SWA_PAGE_SIZE - seq_lens = base + seq_offsets - assert int(seq_lens.max()) + num_verify_tokens < base + 2 * SWA_PAGE_SIZE - return seq_lens.to(device=DEFAULT_DEVICE) - - -def make_req_to_token( - batch_size: int, max_seq_len: int, num_chunks: int -) -> torch.Tensor: - chunk_ids = torch.arange(max_seq_len, dtype=torch.int32) // SWA_PAGE_SIZE - req_offsets = torch.arange(batch_size, dtype=torch.int32).unsqueeze(1) * num_chunks - req_to_token = req_offsets + chunk_ids.unsqueeze(0) - return req_to_token.contiguous().to(device=DEFAULT_DEVICE) - - -def make_case(batch_size: int, num_verify_tokens: int) -> BenchmarkCase: - seq_lens = make_seq_lens(batch_size, num_verify_tokens) - req_pool_indices = torch.arange( - batch_size, dtype=torch.int64, device=DEFAULT_DEVICE - ) - - max_seq_len = int(seq_lens.max().item()) + num_verify_tokens + SWA_PAGE_SIZE - num_chunks = round_up_div(max_seq_len, SWA_PAGE_SIZE) - 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( - (state_slot_stride * (1 + num_verify_tokens), STATE_DIM), - dtype=torch.float32, - device=DEFAULT_DEVICE, - ) - state.normal_(mean=0.0, std=0.01) - - kv_score_input = torch.randn( - batch_size * num_verify_tokens, - HEAD_DIM * 2, - dtype=torch.float32, - device=DEFAULT_DEVICE, - ) - ape = torch.randn(128, HEAD_DIM, dtype=torch.float32, device=DEFAULT_DEVICE) - - return BenchmarkCase( - kv_score_input=kv_score_input, - 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, - num_verify_tokens=num_verify_tokens, - state_slot_stride=state_slot_stride, - ) - - -def call_write_prefix(module, case: BenchmarkCase) -> None: - module.write_prefix_states( - case.kv_score_input, - 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, - ) - - -@triton.testing.perf_report( - triton.testing.Benchmark( - x_names=["batch_size", "num_verify_tokens"], - x_vals=BENCHMARK_CONFIGS, - line_arg="launch_mode", - line_vals=["cuda_graph", "eager"], - line_names=["CUDA graph", "Eager launch"], - styles=[("blue", "-"), ("orange", "--")], - ylabel="us", - plot_name="online-c128-mtp-write-prefix-performance", - args={}, - ) -) -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) - fn = lambda: call_write_prefix(module, case) - - if launch_mode == "cuda_graph": - return run_benchmark(fn) - if launch_mode == "eager": - return run_benchmark_no_cudagraph(fn) - raise ValueError(f"Unknown launch_mode: {launch_mode}") - - -if __name__ == "__main__": - benchmark.run(print_data=True)