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 71e600dc3..02f263078 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_0` (-1 = none), writes new -/// running state to `read_page_1`. +/// Reads optional prior state from `read_page_1` (-1 = fallback to +/// `read_page_0`), writes new running state to `read_page_0`. 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 | - // The plan stores last-token coordinates; segment start is recoverable as - // ragged_id - window_len + 1. + // `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. 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 segment_start_ragged = static_cast(plan.ragged_id) - static_cast(position % 128u); + const int32_t chunk_start_ragged = static_cast(plan.ragged_id) - static_cast(pos_in_chunk_end) + 1; // --- Stage 1: load kv / score / bias for this warp's 8 chunk positions. PrefillStorage kv[kElementsPerWarp]; @@ -297,7 +297,8 @@ __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 auto kv_src_ptr = kv_score_input + (segment_start_ragged + j) * kElementSize + split_offset; + 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 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); @@ -334,9 +335,10 @@ __global__ __launch_bounds__(kPrefillBlockSize, 2) // out_max_vec.load(seg_max, lane_id); out_sum_vec.load(seg_sum, lane_id); - if (chunk_offset != 0 && plan.read_page_0 >= 0) { + 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) { // Combine with prior partial state for this slot. - const auto buf_load = kv_score_buffer + plan.read_page_0 * (kHeadDim * 3) + split_offset; + const auto buf_load = kv_score_buffer + read_page * (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); @@ -467,8 +469,8 @@ struct FlashCompress128OnlineKernel { .with_device(device_) .verify(ape); - // Both compress and write segments use PlanC layout. plan_c uses - // read_page_1=-1 (unused); plan_w uses read_page_1=store_slot. + // 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. 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(); @@ -512,8 +514,9 @@ 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), -// and `read_page_1` is unused. +// 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`. // =========================================================================== namespace host::compress { @@ -533,6 +536,7 @@ 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; }; @@ -544,7 +548,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; + const int32_t slot = swa_loc / params.swa_page_size + params.state_slot_offset; params.plan_d[idx] = DecodePlan{ .seq_len = seq_len, .write_loc = slot, @@ -564,7 +568,8 @@ 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 swa_page_size, + const int32_t state_slot_offset) { auto B = SymbolicSize{"batch_size"}; auto device_ = SymbolicDevice{}; device_.set_options(); @@ -591,6 +596,7 @@ 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; @@ -607,6 +613,7 @@ 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); @@ -654,7 +661,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, // unused, kept so MSB layout is stable + .read_page_1 = -1, // filled by stage 1 with committed-bank slot }; if (chunk_off + seg_len == 128u) { // close-chunk segment @@ -681,6 +688,7 @@ 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; }; @@ -693,13 +701,16 @@ __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]); - plan.read_page_0 = swa_loc / params.swa_page_size; + 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_ptr = plan; } @@ -715,7 +726,9 @@ 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 swa_page_size, + const int32_t state_slot_offset, + const bool use_cuda_graph) { auto B = SymbolicSize{"batch_size"}; auto N = SymbolicSize{"num_q_tokens"}; auto cpu = SymbolicDevice{}; @@ -750,6 +763,7 @@ 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()), @@ -784,6 +798,8 @@ 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. @@ -821,7 +837,19 @@ 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 (const auto total = num_c + num_w) { + 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) { const auto stream = LaunchKernel::resolve_device(device); // SGLANG_DEBUG_C128_ONLINE_SYNC_H2D=1 forces a synchronous H2D copy. static const bool kSyncH2D = []() { @@ -842,8 +870,8 @@ inline OnlinePrefillPlan plan_online_prefill( RuntimeDeviceCheck(::cudaMemcpyAsync(dst, src, bytes, ::cudaMemcpyHostToDevice, stream)); } }; - 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); + 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); const auto stage1_params = OnlinePrefillStage1Params{ .plan_c = plan_c_dev_ptr, @@ -853,14 +881,15 @@ 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, - .num_c = num_c, - .num_w = num_w, + .state_slot_offset = state_slot_offset, + .num_c = num_c_padded, + .num_w = num_w_padded, }; 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, num_w}; + return OnlinePrefillPlan{num_c_padded, num_w_padded}; } } // 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 new file mode 100644 index 000000000..db1b94d89 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/online_c128_mtp.cuh @@ -0,0 +1,537 @@ +#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 ea3d8f729..cd62eaf29 100644 --- a/python/sglang/jit_kernel/dsv4/compress.py +++ b/python/sglang/jit_kernel/dsv4/compress.py @@ -1,8 +1,9 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Union +from typing import Literal, NamedTuple, Optional, Union import torch +from tvm_ffi.module import Module from sglang.jit_kernel.utils import ( cache_once, @@ -13,9 +14,6 @@ from sglang.jit_kernel.utils import ( from .utils import make_name -if TYPE_CHECKING: - from tvm_ffi.module import Module - @cache_once def _jit_compress_norm_rope_module( @@ -156,6 +154,7 @@ class CompressorDecodePlan(NamedTuple): 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) @@ -165,7 +164,13 @@ class CompressorDecodePlan(NamedTuple): device=req_pool_indices.device, ) module.plan_decode( - seq_lens, req_pool_indices, req_to_token, full_to_swa, plan_d, swa_page_size + 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) @@ -267,9 +272,11 @@ class CompressorPrefillPlan(NamedTuple): 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.to(torch.int64) - extend_lens_cpu = extend_lens.to(torch.int64) + seq_lens_cpu = seq_lens.detach().to(torch.int64).cpu() + 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) @@ -292,6 +299,8 @@ class CompressorPrefillPlan(NamedTuple): 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 new file mode 100644 index 000000000..d93c6f120 --- /dev/null +++ b/python/sglang/jit_kernel/dsv4/online_c128_mtp.py @@ -0,0 +1,256 @@ +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 b0f1923e8..300a76bef 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -782,6 +782,7 @@ 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_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 65f44fe0f..54f502442 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -35,6 +35,7 @@ else: 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, ) @@ -79,6 +80,37 @@ C4_TOPK = 512 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]) @@ -102,6 +134,13 @@ 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 @@ -380,6 +419,8 @@ 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) @@ -387,6 +428,10 @@ class 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 @@ -431,6 +476,7 @@ class DeepseekV4AttnBackend( 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 ( @@ -472,11 +518,41 @@ class DeepseekV4AttnBackend( 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, @@ -542,6 +618,7 @@ class DeepseekV4AttnBackend( 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, @@ -591,6 +668,7 @@ class DeepseekV4AttnBackend( 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, @@ -603,6 +681,7 @@ class DeepseekV4AttnBackend( 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) @@ -619,11 +698,18 @@ class DeepseekV4AttnBackend( 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 @@ -635,6 +721,15 @@ class DeepseekV4AttnBackend( 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() @@ -645,6 +740,7 @@ class DeepseekV4AttnBackend( 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( @@ -655,6 +751,7 @@ class DeepseekV4AttnBackend( 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 @@ -676,10 +773,13 @@ class DeepseekV4AttnBackend( 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 + self, + raw_metadata: DSV4RawVerifyMetadata, + online_c128_state_slot_offset: int = 0, ) -> DSV4Metadata: req_pool_indices = raw_metadata.req_pool_indices seq_lens = raw_metadata.seq_lens @@ -688,6 +788,7 @@ class DeepseekV4AttnBackend( 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( @@ -715,16 +816,21 @@ class DeepseekV4AttnBackend( 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=create(compress_ratio=128), + c128_compress_metadata=c128_compress_metadata, ) 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 @@ -793,6 +899,7 @@ class DeepseekV4AttnBackend( 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( @@ -889,6 +996,11 @@ class DeepseekV4AttnBackend( 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)), @@ -902,6 +1014,13 @@ class DeepseekV4AttnBackend( 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( @@ -910,14 +1029,27 @@ class DeepseekV4AttnBackend( 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 @@ -938,6 +1070,7 @@ class DeepseekV4AttnBackend( use_prefill_cuda_graph=True, ) else: + self.online_c128_mtp.clear() raise NotImplementedError self.replay_cuda_graph_metadata_from( @@ -957,7 +1090,9 @@ class DeepseekV4AttnBackend( ) def init_forward_metadata(self, forward_batch: ForwardBatch) -> None: - if self.mtp_enabled and forward_batch.forward_mode.is_idle(): + logical_forward_mode = _get_logical_forward_mode(forward_batch) + if self.mtp_enabled and logical_forward_mode.is_idle(): + self.online_c128_mtp.clear() return self.forward_metadata = self._build_forward_metadata(forward_batch) @@ -970,6 +1105,7 @@ class DeepseekV4AttnBackend( 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 @@ -977,13 +1113,22 @@ class DeepseekV4AttnBackend( 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 forward_batch.forward_mode.is_decode_or_idle(): + if logical_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 @@ -1000,14 +1145,16 @@ class DeepseekV4AttnBackend( seq_lens=seq_lens, out_cache_loc=out_cache_loc, ) - elif forward_batch.forward_mode.is_target_verify(): + elif logical_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 forward_batch.forward_mode.is_prefill(include_draft_extend_v2=True): + elif logical_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 ( @@ -1090,12 +1237,12 @@ class DeepseekV4AttnBackend( ], bucket: _GraphBucket, ) -> None: - 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 + 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 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 @@ -1603,6 +1750,7 @@ class DeepseekV4MultiStepBackend(DeepseekV4AttnBackend): 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 3d5c4b16f..7a2d26617 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py @@ -541,6 +541,17 @@ class CompressorBackendMixin: 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, @@ -673,6 +684,7 @@ 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). @@ -691,6 +703,7 @@ 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 @@ -748,9 +761,8 @@ 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) @@ -778,6 +790,8 @@ 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( @@ -786,4 +800,5 @@ 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 4543c48fd..f43415a5a 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_compress_state.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_compress_state.py @@ -88,18 +88,29 @@ class CompressStatePool: 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._size = size + self.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) 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 f914e6142..b66c9481e 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -35,7 +35,8 @@ 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: - assert not is_speculative, "online c128 does not support MTP" + if is_speculative and not envs.SGLANG_EXPERIMENTAL_ONLINE_C128_MTP.get(): + raise AssertionError("online c128 does not support MTP") return 1 if is_speculative: return 16 if compress_ratio == 4 else 256 @@ -458,6 +459,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): 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__( @@ -493,6 +495,12 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): self.c128_state_pool_size = c128_state_pool_size self.state_dtype = 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 ( @@ -758,6 +766,9 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): 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: @@ -815,6 +826,22 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): ), "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 d35144d94..0127dd1ce 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 @@ -429,6 +429,9 @@ class ModelRunnerKVCacheMixin: 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 166c973b4..b18a26b76 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -348,6 +348,9 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): 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 @@ -382,10 +385,30 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): # would need rollback / replay across draft and verify, which the # online path doesn't support yet. if envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get(): - 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)") + 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)" + ) def _get_bytes_per_full_token(self) -> float: kv_bytes = self.qk_nope_head_dim + self.qk_rope_head_dim * 2 + 8 @@ -409,6 +432,8 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): 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 new file mode 100644 index 000000000..9f7540125 --- /dev/null +++ b/test/registered/jit/benchmark/bench_online_c128_mtp.py @@ -0,0 +1,166 @@ +"""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)