minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

DeepSeek-V4 Online Compress support MTP (#26471)

Browse Source
This commit is contained in:
zhangxiaolei
2026-06-15 19:56:07 -07:00
committed by GitHub
parent b5bcd76a41
commit 063ab89ac1
12 changed files with 1276 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/deepseek_v4/c128_online_v2.cuh
+52 -23
View File
@@ -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 <int64_t kHeadDim, bool kWrite, bool kUsePDL>
__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<int32_t>(plan.ragged_id) - static_cast<int32_t>(position % 128u);
const int32_t chunk_start_ragged = static_cast<int32_t>(plan.ragged_id) - static_cast<int32_t>(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<int32_t>(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<int32_t>((seq_len - 1u) / 128u * 128u);
const int32_t full_loc = params.req_to_token[rid * params.stride_r2t + chunk_start];
const int32_t swa_loc = static_cast<int32_t>(params.full_to_swa[full_loc]);
const int32_t slot = swa_loc / params.swa_page_size;
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<kDLCUDA>();
@@ -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<uint32_t>(B.unwrap());
if (batch_size == 0) return;
@@ -607,6 +613,7 @@ inline void plan_online_decode(
.full_to_swa = static_cast<const int64_t*>(full_to_swa.data_ptr()),
.stride_r2t = stride_r2t,
.swa_page_size = swa_page_size,
.state_slot_offset = state_slot_offset,
.batch_size = batch_size,
};
LaunchKernel(num_blocks, kBlockSize, device)(plan_c128_online_decode_kernel, params);
@@ -654,7 +661,7 @@ inline std::tuple<uint32_t, uint32_t> _plan_prefill_partial(const OnlinePrefillS
.ragged_id = static_cast<uint16_t>(last_ragged),
.buffer_len = static_cast<uint16_t>(seg_len),
.read_page_0 = static_cast<int32_t>(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 ? &params.plan_c[idx] : &params.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<int32_t>(plan.seq_len - 1u);
const int32_t chunk_start = (position / 128) * 128;
const int32_t full_loc = params.req_to_token[rid * params.stride_r2t + chunk_start];
const int32_t swa_loc = static_cast<int32_t>(params.full_to_swa[full_loc]);
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<CompressPlan*>(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<uint32_t>(N.unwrap()) : num_c;
const auto num_w_padded = use_cuda_graph ? static_cast<uint32_t>(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<CompressPlan*>(plan_c_dev_.data_ptr());
auto* const plan_w_dev_ptr = static_cast<CompressPlan*>(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<CompressPlan*>(plan_c_pin.data_ptr());
auto* const plan_w_pin_ptr = static_cast<CompressPlan*>(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<const int64_t*>(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
python/sglang/jit_kernel/csrc/deepseek_v4/online_c128_mtp.cuh
+537
View File
@@ -0,0 +1,537 @@
#pragma once
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/runtime.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <cstdint>
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 <typename TSeq, typename TReq>
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 <typename TSeq, typename TReq>
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 <typename TSeq, typename TReq>
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<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
if (idx < max_num_reqs) pending_seq_lens[idx] = -1;
}
template <typename TSeq, typename TReq>
__global__ void online_c128_mtp_mark_pending_kernel(const OnlineC128MTPMarkPendingParams<TSeq, TReq> params) {
const int64_t bid = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
if (bid >= params.bs) return;
const int64_t req = static_cast<int64_t>(params.req_pool_indices[bid]);
if (req >= 0 && req < params.max_num_reqs) {
params.pending_seq_lens[req] = static_cast<int64_t>(params.seq_lens[bid]);
}
}
template <int64_t kHeadDim, typename TSeq, typename TReq>
__global__ void online_c128_mtp_commit_pending_kernel(const OnlineC128MTPCommitPendingParams<TSeq, TReq> params) {
const int64_t bid = static_cast<int64_t>(blockIdx.x);
if (bid >= params.cur_bs) return;
const int64_t req = static_cast<int64_t>(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<int64_t>(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<int64_t>(params.req_to_token[req * params.req_to_token_stride_b + chunk_start]);
const int64_t swa_loc = params.full_to_swa[full_loc];
const int64_t slot = swa_loc / params.swa_page_size;
const 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<int64_t>(threadIdx.x); d < kHeadDim * 3; d += blockDim.x) {
dst[d] = src[d];
}
}
template <int64_t kHeadDim, typename TSeq, typename TReq>
__global__ void online_c128_mtp_write_prefix_kernel(const OnlineC128MTPWritePrefixParams<TSeq, TReq> params) {
const int64_t bid = static_cast<int64_t>(blockIdx.x);
if (bid >= params.layer_bs) return;
const int64_t seq_before = static_cast<int64_t>(params.seq_lens[bid]);
const int64_t req_idx = static_cast<int64_t>(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<int64_t>(params.req_to_token[req_idx * params.req_to_token_stride_b + chunk_start]);
const int64_t swa_loc = params.full_to_swa[full_loc];
init_slot = swa_loc / params.swa_page_size;
}
const int64_t d = static_cast<int64_t>(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<int64_t>(params.req_to_token[req_idx * params.req_to_token_stride_b + chunk_start]);
const int64_t swa_loc = params.full_to_swa[full_loc];
const int64_t slot = swa_loc / params.swa_page_size + (step + 1) * params.state_slot_stride;
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 <int64_t kHeadDim>
struct OnlineC128MTPWritePrefixKernel {
template <typename TSeq, typename TReq>
static void launch(
tvm::ffi::TensorView kv_score_input,
tvm::ffi::TensorView seq_lens,
tvm::ffi::TensorView req_pool_indices,
tvm::ffi::TensorView req_to_token,
tvm::ffi::TensorView full_to_swa,
tvm::ffi::TensorView ape,
tvm::ffi::TensorView state,
int64_t layer_bs,
int64_t swa_page_size,
int64_t num_verify_tokens,
int64_t state_slot_stride,
DLDevice device) {
using namespace host;
const auto params = OnlineC128MTPWritePrefixParams<TSeq, TReq>{
.kv_score_input = static_cast<const float*>(kv_score_input.data_ptr()),
.seq_lens = static_cast<const TSeq*>(seq_lens.data_ptr()),
.req_pool_indices = static_cast<const TReq*>(req_pool_indices.data_ptr()),
.req_to_token = static_cast<const int32_t*>(req_to_token.data_ptr()),
.full_to_swa = static_cast<const int64_t*>(full_to_swa.data_ptr()),
.ape = static_cast<const float*>(ape.data_ptr()),
.state = static_cast<float*>(state.data_ptr()),
.kv_score_stride_b = kv_score_input.stride(0),
.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<uint32_t>(kHeadDim);
LaunchKernel(static_cast<uint32_t>(layer_bs), kThreads, device)(
online_c128_mtp_write_prefix_kernel<kHeadDim, TSeq, TReq>, 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<kDLCUDA>();
TensorMatcher({-1, kHeadDim * 2}).with_dtype<float>().with_device(device).verify(kv_score_input);
TensorMatcher({-1}).with_dtype<int32_t, int64_t>(seq_dtype).with_device(device).verify(seq_lens);
TensorMatcher({-1}).with_dtype<int32_t, int64_t>(req_dtype).with_device(device).verify(req_pool_indices);
TensorMatcher({-1, -1}).with_dtype<int32_t>().with_device(device).verify(req_to_token);
TensorMatcher({-1}).with_dtype<int64_t>().with_device(device).verify(full_to_swa);
TensorMatcher({128, kHeadDim}).with_dtype<float>().with_device(device).verify(ape);
TensorMatcher({-1, kHeadDim * 3}).with_dtype<float>().with_device(device).verify(state);
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<int32_t>()) {
if (req_dtype.is_type<int32_t>()) {
launch<int32_t, int32_t>(
kv_score_input,
seq_lens,
req_pool_indices,
req_to_token,
full_to_swa,
ape,
state,
layer_bs,
swa_page_size,
num_verify_tokens,
state_slot_stride,
device.unwrap());
} else {
launch<int32_t, int64_t>(
kv_score_input,
seq_lens,
req_pool_indices,
req_to_token,
full_to_swa,
ape,
state,
layer_bs,
swa_page_size,
num_verify_tokens,
state_slot_stride,
device.unwrap());
}
} else {
if (req_dtype.is_type<int32_t>()) {
launch<int64_t, int32_t>(
kv_score_input,
seq_lens,
req_pool_indices,
req_to_token,
full_to_swa,
ape,
state,
layer_bs,
swa_page_size,
num_verify_tokens,
state_slot_stride,
device.unwrap());
} else {
launch<int64_t, int64_t>(
kv_score_input,
seq_lens,
req_pool_indices,
req_to_token,
full_to_swa,
ape,
state,
layer_bs,
swa_page_size,
num_verify_tokens,
state_slot_stride,
device.unwrap());
}
}
}
};
template <int64_t kHeadDim>
struct OnlineC128MTPMarkPendingKernel {
template <typename TSeq, typename TReq>
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<TSeq, TReq>{
.seq_lens = static_cast<const TSeq*>(seq_lens.data_ptr()),
.req_pool_indices = static_cast<const TReq*>(req_pool_indices.data_ptr()),
.pending_seq_lens = static_cast<int64_t*>(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<uint32_t>(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<uint32_t>(bs), kThreads);
LaunchKernel(mark_blocks, kThreads, device)(online_c128_mtp_mark_pending_kernel<TSeq, TReq>, 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<kDLCUDA>();
TensorMatcher({-1}).with_dtype<int32_t, int64_t>(seq_dtype).with_device(device).verify(seq_lens);
TensorMatcher({-1}).with_dtype<int32_t, int64_t>(req_dtype).with_device(device).verify(req_pool_indices);
TensorMatcher({-1}).with_dtype<int64_t>().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<int32_t>()) {
if (req_dtype.is_type<int32_t>()) {
launch<int32_t, int32_t>(seq_lens, req_pool_indices, pending_seq_lens, bs, max_num_reqs, device.unwrap());
} else {
launch<int32_t, int64_t>(seq_lens, req_pool_indices, pending_seq_lens, bs, max_num_reqs, device.unwrap());
}
} else {
if (req_dtype.is_type<int32_t>()) {
launch<int64_t, int32_t>(seq_lens, req_pool_indices, pending_seq_lens, bs, max_num_reqs, device.unwrap());
} else {
launch<int64_t, int64_t>(seq_lens, req_pool_indices, pending_seq_lens, bs, max_num_reqs, device.unwrap());
}
}
}
};
template <int64_t kHeadDim>
struct OnlineC128MTPCommitPendingKernel {
template <typename TSeq, typename TReq>
static void launch(
tvm::ffi::TensorView cur_seq_lens,
tvm::ffi::TensorView cur_req_pool_indices,
tvm::ffi::TensorView req_to_token,
tvm::ffi::TensorView full_to_swa,
tvm::ffi::TensorView pending_seq_lens,
tvm::ffi::TensorView state,
int64_t cur_bs,
int64_t swa_page_size,
int64_t num_verify_tokens,
int64_t state_slot_stride,
int64_t max_num_reqs,
DLDevice device) {
using namespace host;
const auto params = OnlineC128MTPCommitPendingParams<TSeq, TReq>{
.cur_seq_lens = static_cast<const TSeq*>(cur_seq_lens.data_ptr()),
.cur_req_pool_indices = static_cast<const TReq*>(cur_req_pool_indices.data_ptr()),
.req_to_token = static_cast<const int32_t*>(req_to_token.data_ptr()),
.full_to_swa = static_cast<const int64_t*>(full_to_swa.data_ptr()),
.pending_seq_lens = static_cast<const int64_t*>(pending_seq_lens.data_ptr()),
.state = static_cast<float*>(state.data_ptr()),
.cur_bs = cur_bs,
.req_to_token_stride_b = req_to_token.stride(0),
.state_stride_b = state.stride(0),
.swa_page_size = swa_page_size,
.num_verify_tokens = num_verify_tokens,
.state_slot_stride = state_slot_stride,
.max_num_reqs = max_num_reqs,
};
constexpr uint32_t kThreads = 256;
LaunchKernel(static_cast<uint32_t>(cur_bs), kThreads, device)(
online_c128_mtp_commit_pending_kernel<kHeadDim, TSeq, TReq>, 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<kDLCUDA>();
TensorMatcher({-1}).with_dtype<int32_t, int64_t>(seq_dtype).with_device(device).verify(cur_seq_lens);
TensorMatcher({-1}).with_dtype<int32_t, int64_t>(req_dtype).with_device(device).verify(cur_req_pool_indices);
TensorMatcher({-1, -1}).with_dtype<int32_t>().with_device(device).verify(req_to_token);
TensorMatcher({-1}).with_dtype<int64_t>().with_device(device).verify(full_to_swa);
TensorMatcher({-1}).with_dtype<int64_t>().with_device(device).verify(pending_seq_lens);
TensorMatcher({-1, kHeadDim * 3}).with_dtype<float>().with_device(device).verify(state);
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<int32_t>()) {
if (req_dtype.is_type<int32_t>()) {
launch<int32_t, int32_t>(
cur_seq_lens,
cur_req_pool_indices,
req_to_token,
full_to_swa,
pending_seq_lens,
state,
cur_bs,
swa_page_size,
num_verify_tokens,
state_slot_stride,
max_num_reqs,
device.unwrap());
} else {
launch<int32_t, int64_t>(
cur_seq_lens,
cur_req_pool_indices,
req_to_token,
full_to_swa,
pending_seq_lens,
state,
cur_bs,
swa_page_size,
num_verify_tokens,
state_slot_stride,
max_num_reqs,
device.unwrap());
}
} else {
if (req_dtype.is_type<int32_t>()) {
launch<int64_t, int32_t>(
cur_seq_lens,
cur_req_pool_indices,
req_to_token,
full_to_swa,
pending_seq_lens,
state,
cur_bs,
swa_page_size,
num_verify_tokens,
state_slot_stride,
max_num_reqs,
device.unwrap());
} else {
launch<int64_t, int64_t>(
cur_seq_lens,
cur_req_pool_indices,
req_to_token,
full_to_swa,
pending_seq_lens,
state,
cur_bs,
swa_page_size,
num_verify_tokens,
state_slot_stride,
max_num_reqs,
device.unwrap());
}
}
}
};
} // namespace
python/sglang/jit_kernel/dsv4/compress.py
+16 -7
View File
@@ -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,
python/sglang/jit_kernel/dsv4/online_c128_mtp.py
+256
View File
@@ -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)
python/sglang/srt/environ.py
+1
View File
@@ -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)
python/sglang/srt/layers/attention/deepseek_v4_backend.py
+159 -11
View File
@@ -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] = []
python/sglang/srt/layers/attention/dsv4/compressor_v2.py
+17 -2
View File
@@ -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,
)
python/sglang/srt/mem_cache/deepseek_v4_compress_state.py
+12 -1
View File
@@ -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:
python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
+28 -1
View File
@@ -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]
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
+3
View File
@@ -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:
python/sglang/srt/model_executor/pool_configurator.py
+29 -4
View File
@@ -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 (
test/registered/jit/benchmark/bench_online_c128_mtp.py
+166
View File
@@ -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)