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

[DeepSeek-V4] Support BF16 Compress State for Online C128 (#29609)

Browse Source 
Co-authored-by: zhujunyu <zhujunyu.666@bytedance.com>
This commit is contained in:
Ryan Zzz
2026-07-15 23:17:29 -07:00
committed by GitHub
co-authored by zhujunyu
parent dc60f65661
commit 5af65d8542
7 changed files with 160 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/deepseek_v4/c128_online_v2.cuh
+76 -19
View File
@@ -20,6 +20,7 @@
#include <cstdint> #include <cstdint>
#include <cstdlib> #include <cstdlib>
#include <cstring> #include <cstring>
#include <type_traits>
namespace { namespace {
@@ -40,7 +41,7 @@ struct Compress128OnlineDecodeParams {
uint32_t batch_size; uint32_t batch_size;
}; };
template <int64_t kHeadDim, bool kUsePDL> template <int64_t kHeadDim, typename BufferFloat, bool kUsePDL>
__global__ void flash_c128_online_decode_v2(const __grid_constant__ Compress128OnlineDecodeParams params) { __global__ void flash_c128_online_decode_v2(const __grid_constant__ Compress128OnlineDecodeParams params) {
using namespace device; using namespace device;
constexpr uint32_t kVecSize = 4; constexpr uint32_t kVecSize = 4;
@@ -59,7 +60,7 @@ __global__ void flash_c128_online_decode_v2(const __grid_constant__ Compress128O
const auto plan = params.plan_d[batch_id]; const auto plan = params.plan_d[batch_id];
const auto pos_in_chunk = (plan.seq_len - 1) % 128; const auto pos_in_chunk = (plan.seq_len - 1) % 128;
const auto kv_score_buffer = static_cast<float*>(params.kv_score_buffer); const auto kv_score_buffer = static_cast<BufferFloat*>(params.kv_score_buffer);
const auto kv_score_input = static_cast<const float*>(params.kv_score_input); const auto kv_score_input = static_cast<const float*>(params.kv_score_input);
const auto kv_load_buf = kv_score_buffer + plan.read_page_0 * (kHeadDim * 3); const auto kv_load_buf = kv_score_buffer + plan.read_page_0 * (kHeadDim * 3);
const auto kv_store_buf = kv_score_buffer + plan.write_loc * (kHeadDim * 3); const auto kv_store_buf = kv_score_buffer + plan.write_loc * (kHeadDim * 3);
@@ -75,9 +76,24 @@ __global__ void flash_c128_online_decode_v2(const __grid_constant__ Compress128O
Vec out_sum_vec; Vec out_sum_vec;
if (pos_in_chunk != 0) { if (pos_in_chunk != 0) {
// Mid-chunk: combine prior partial state with the new token. // Mid-chunk: combine prior partial state with the new token.
const auto max_score_vec = gmem.load(kv_load_buf, 0); Vec max_score_vec, sum_score_vec, old_kv_vec;
const auto sum_score_vec = gmem.load(kv_load_buf, 1); if constexpr (std::is_same_v<BufferFloat, float>) {
const auto old_kv_vec = gmem.load(kv_load_buf, 2); max_score_vec = gmem.load(kv_load_buf, 0);
sum_score_vec = gmem.load(kv_load_buf, 1);
old_kv_vec = gmem.load(kv_load_buf, 2);
} else {
using BufferVec = AlignedVector<BufferFloat, kVecSize>;
const auto gmem_buffer = tile::Memory<BufferVec>::cta(kBlockSize);
const auto max_score_tmp = gmem_buffer.load(kv_load_buf, 0);
const auto sum_score_tmp = gmem_buffer.load(kv_load_buf, 1);
const auto old_kv_tmp = gmem_buffer.load(kv_load_buf, 2);
#pragma unroll
for (uint32_t i = 0; i < kVecSize; ++i) {
max_score_vec[i] = cast<float>(max_score_tmp[i]);
sum_score_vec[i] = cast<float>(sum_score_tmp[i]);
old_kv_vec[i] = cast<float>(old_kv_tmp[i]);
}
}
#pragma unroll #pragma unroll
for (uint32_t i = 0; i < kVecSize; ++i) { for (uint32_t i = 0; i < kVecSize; ++i) {
const auto old_max = max_score_vec[i]; const auto old_max = max_score_vec[i];
@@ -107,9 +123,24 @@ __global__ void flash_c128_online_decode_v2(const __grid_constant__ Compress128O
const auto kv_out = static_cast<float*>(params.kv_compressed_output) + batch_id * kHeadDim; const auto kv_out = static_cast<float*>(params.kv_compressed_output) + batch_id * kHeadDim;
gmem.store(kv_out, out_kv_vec); gmem.store(kv_out, out_kv_vec);
} else { } else {
if constexpr (std::is_same_v<BufferFloat, float>) {
gmem.store(kv_store_buf, out_max_vec, 0); gmem.store(kv_store_buf, out_max_vec, 0);
gmem.store(kv_store_buf, out_sum_vec, 1); gmem.store(kv_store_buf, out_sum_vec, 1);
gmem.store(kv_store_buf, out_kv_vec, 2); gmem.store(kv_store_buf, out_kv_vec, 2);
} else {
using BufferVec = AlignedVector<BufferFloat, kVecSize>;
const auto gmem_buffer = tile::Memory<BufferVec>::cta(kBlockSize);
BufferVec out_max_tmp, out_sum_tmp, out_kv_tmp;
#pragma unroll
for (uint32_t i = 0; i < kVecSize; ++i) {
out_max_tmp[i] = cast<BufferFloat>(out_max_vec[i]);
out_sum_tmp[i] = cast<BufferFloat>(out_sum_vec[i]);
out_kv_tmp[i] = cast<BufferFloat>(out_kv_vec[i]);
}
gmem_buffer.store(kv_store_buf, out_max_tmp, 0);
gmem_buffer.store(kv_store_buf, out_sum_tmp, 1);
gmem_buffer.store(kv_store_buf, out_kv_tmp, 2);
}
} }
} }
@@ -150,8 +181,7 @@ struct Compress128SharedBuffer {
/// \brief Sentinel score for padded positions in a 128-segment. /// \brief Sentinel score for padded positions in a 128-segment.
constexpr float kPadScore = -FLT_MAX; constexpr float kPadScore = -FLT_MAX;
[[maybe_unused]] [[maybe_unused]] SGL_DEVICE void c128_prefill_segment_softmax(
SGL_DEVICE void c128_prefill_segment_softmax(
const PrefillStorage (&kv)[kElementsPerWarp], const PrefillStorage (&kv)[kElementsPerWarp],
const PrefillStorage (&score)[kElementsPerWarp], const PrefillStorage (&score)[kElementsPerWarp],
float* seg_kv, float* seg_kv,
@@ -239,7 +269,7 @@ SGL_DEVICE void c128_prefill_segment_softmax(
/// `kWrite=true` (write pass) : handles trailing partial segments. /// `kWrite=true` (write pass) : handles trailing partial segments.
/// Reads optional prior state from `read_page_1` (-1 = fallback to /// Reads optional prior state from `read_page_1` (-1 = fallback to
/// `read_page_0`), writes new running state to `read_page_0`. /// `read_page_0`), writes new running state to `read_page_0`.
template <int64_t kHeadDim, bool kWrite, bool kUsePDL> template <int64_t kHeadDim, typename BufferFloat, bool kWrite, bool kUsePDL>
__global__ __launch_bounds__(kPrefillBlockSize, 2) // __global__ __launch_bounds__(kPrefillBlockSize, 2) //
void flash_c128_online_prefill_v2(const __grid_constant__ Compress128OnlinePrefillParams params) { void flash_c128_online_prefill_v2(const __grid_constant__ Compress128OnlinePrefillParams params) {
using namespace device; using namespace device;
@@ -272,7 +302,7 @@ __global__ __launch_bounds__(kPrefillBlockSize, 2) //
if (plan.is_invalid()) [[unlikely]] if (plan.is_invalid()) [[unlikely]]
return; return;
const auto kv_score_buffer = static_cast<float*>(params.kv_score_buffer); const auto kv_score_buffer = static_cast<BufferFloat*>(params.kv_score_buffer);
const auto kv_score_input = static_cast<const float*>(params.kv_score_input); const auto kv_score_input = static_cast<const float*>(params.kv_score_input);
const auto kv_compressed_output = static_cast<float*>(params.kv_compressed_output); const auto kv_compressed_output = static_cast<float*>(params.kv_compressed_output);
const auto score_bias_base = static_cast<const float*>(params.score_bias); const auto score_bias_base = static_cast<const float*>(params.score_bias);
@@ -340,9 +370,23 @@ __global__ __launch_bounds__(kPrefillBlockSize, 2) //
// Combine with prior partial state for this slot. // Combine with prior partial state for this slot.
const auto buf_load = kv_score_buffer + read_page * (kHeadDim * 3) + split_offset; const auto buf_load = kv_score_buffer + read_page * (kHeadDim * 3) + split_offset;
PrefillStorage buf_max_vec, buf_sum_vec, buf_kv_vec; PrefillStorage buf_max_vec, buf_sum_vec, buf_kv_vec;
if constexpr (std::is_same_v<BufferFloat, float>) {
buf_max_vec.load(buf_load + 0 * kHeadDim, lane_id); buf_max_vec.load(buf_load + 0 * kHeadDim, lane_id);
buf_sum_vec.load(buf_load + 1 * kHeadDim, lane_id); buf_sum_vec.load(buf_load + 1 * kHeadDim, lane_id);
buf_kv_vec.load(buf_load + 2 * kHeadDim, lane_id); buf_kv_vec.load(buf_load + 2 * kHeadDim, lane_id);
} else {
using BufferPrefillStorage = AlignedVector<BufferFloat, kTileElements>;
BufferPrefillStorage buf_max_tmp, buf_sum_tmp, buf_kv_tmp;
buf_max_tmp.load(buf_load + 0 * kHeadDim, lane_id);
buf_sum_tmp.load(buf_load + 1 * kHeadDim, lane_id);
buf_kv_tmp.load(buf_load + 2 * kHeadDim, lane_id);
#pragma unroll
for (uint32_t ii = 0; ii < kTileElements; ++ii) {
buf_max_vec[ii] = cast<float>(buf_max_tmp[ii]);
buf_sum_vec[ii] = cast<float>(buf_sum_tmp[ii]);
buf_kv_vec[ii] = cast<float>(buf_kv_tmp[ii]);
}
}
#pragma unroll #pragma unroll
for (uint32_t ii = 0; ii < kTileElements; ++ii) { for (uint32_t ii = 0; ii < kTileElements; ++ii) {
const float m1 = buf_max_vec[ii]; const float m1 = buf_max_vec[ii];
@@ -367,9 +411,23 @@ __global__ __launch_bounds__(kPrefillBlockSize, 2) //
// segment's own chunk slot (the request keeps a single in-progress // segment's own chunk slot (the request keeps a single in-progress
// chunk's running state at any time), so we reuse `read_page_0`. // chunk's running state at any time), so we reuse `read_page_0`.
const auto buf_store = kv_score_buffer + plan.read_page_0 * (kHeadDim * 3) + split_offset; const auto buf_store = kv_score_buffer + plan.read_page_0 * (kHeadDim * 3) + split_offset;
if constexpr (std::is_same_v<BufferFloat, float>) {
reinterpret_cast<PrefillStorage*>(buf_store + 0 * kHeadDim)[lane_id] = out_max_vec; reinterpret_cast<PrefillStorage*>(buf_store + 0 * kHeadDim)[lane_id] = out_max_vec;
reinterpret_cast<PrefillStorage*>(buf_store + 1 * kHeadDim)[lane_id] = out_sum_vec; reinterpret_cast<PrefillStorage*>(buf_store + 1 * kHeadDim)[lane_id] = out_sum_vec;
reinterpret_cast<PrefillStorage*>(buf_store + 2 * kHeadDim)[lane_id] = out_kv_vec; reinterpret_cast<PrefillStorage*>(buf_store + 2 * kHeadDim)[lane_id] = out_kv_vec;
} else {
using BufferPrefillStorage = AlignedVector<BufferFloat, kTileElements>;
BufferPrefillStorage out_max_tmp, out_sum_tmp, out_kv_tmp;
#pragma unroll
for (uint32_t ii = 0; ii < kTileElements; ++ii) {
out_max_tmp[ii] = cast<BufferFloat>(out_max_vec[ii]);
out_sum_tmp[ii] = cast<BufferFloat>(out_sum_vec[ii]);
out_kv_tmp[ii] = cast<BufferFloat>(out_kv_vec[ii]);
}
reinterpret_cast<BufferPrefillStorage*>(buf_store + 0 * kHeadDim)[lane_id] = out_max_tmp;
reinterpret_cast<BufferPrefillStorage*>(buf_store + 1 * kHeadDim)[lane_id] = out_sum_tmp;
reinterpret_cast<BufferPrefillStorage*>(buf_store + 2 * kHeadDim)[lane_id] = out_kv_tmp;
}
} else { } else {
// Compact output: one row per compress plan, indexed by `global_pid`. // Compact output: one row per compress plan, indexed by `global_pid`.
const auto out_ptr = kv_compressed_output + global_pid * kHeadDim + split_offset; const auto out_ptr = kv_compressed_output + global_pid * kHeadDim + split_offset;
@@ -381,14 +439,15 @@ __global__ __launch_bounds__(kPrefillBlockSize, 2) //
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
// Host wrapper: matches the c128_v2 / c4_v2 host API style (run_decode / // Host wrapper: matches the c128_v2 / c4_v2 host API style (run_decode /
// run_prefill methods on a kernel-class template). We only expose `kHeadDim` // run_prefill methods on a kernel-class template). We only expose `kHeadDim`
// + `kUsePDL`; the dtype is fixed to fp32 for the online state pool. // + state-buffer dtype + `kUsePDL`. Inputs, outputs, and APE remain fp32; only
// the online C128 running state buffer can be fp32/bf16.
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
template <int64_t kHeadDim, bool kUsePDL> template <int64_t kHeadDim, typename BufferFloat, bool kUsePDL>
struct FlashCompress128OnlineKernel { struct FlashCompress128OnlineKernel {
static constexpr auto decode_kernel = flash_c128_online_decode_v2<kHeadDim, kUsePDL>; static constexpr auto decode_kernel = flash_c128_online_decode_v2<kHeadDim, BufferFloat, kUsePDL>;
template <bool kWrite> template <bool kWrite>
static constexpr auto prefill_kernel = flash_c128_online_prefill_v2<kHeadDim, kWrite, kUsePDL>; static constexpr auto prefill_kernel = flash_c128_online_prefill_v2<kHeadDim, BufferFloat, kWrite, kUsePDL>;
static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads; // 64 static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads; // 64
static constexpr uint32_t kNumSplit = kHeadDim / kTileDim; static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
static constexpr uint32_t kDecodeBlockSize = kHeadDim / 4; static constexpr uint32_t kDecodeBlockSize = kHeadDim / 4;
@@ -406,7 +465,7 @@ struct FlashCompress128OnlineKernel {
device_.set_options<kDLCUDA>(); device_.set_options<kDLCUDA>();
TensorMatcher({-1, 1, kHeadDim * 3}) // kv score buffer (max, sum, kv) TensorMatcher({-1, 1, kHeadDim * 3}) // kv score buffer (max, sum, kv)
.with_dtype<float>() .with_dtype<BufferFloat>()
.with_device(device_) .with_device(device_)
.verify(kv_score_buffer); .verify(kv_score_buffer);
TensorMatcher({B, kHeadDim * 2}) // kv score input TensorMatcher({B, kHeadDim * 2}) // kv score input
@@ -453,7 +512,7 @@ struct FlashCompress128OnlineKernel {
device_.set_options<kDLCUDA>(); device_.set_options<kDLCUDA>();
TensorMatcher({-1, 1, kHeadDim * 3}) // kv score buffer TensorMatcher({-1, 1, kHeadDim * 3}) // kv score buffer
.with_dtype<float>() .with_dtype<BufferFloat>()
.with_device(device_) .with_device(device_)
.verify(kv_score_buffer); .verify(kv_score_buffer);
TensorMatcher({N, kHeadDim * 2}) // kv score input (ragged) TensorMatcher({N, kHeadDim * 2}) // kv score input (ragged)
@@ -868,9 +927,7 @@ inline OnlinePrefillPlan plan_online_prefill(
namespace { namespace {
[[maybe_unused]] [[maybe_unused]] constexpr auto& plan_compress_128_online_decode = host::compress::plan_online_decode;
constexpr auto& plan_compress_128_online_decode = host::compress::plan_online_decode; [[maybe_unused]] constexpr auto& plan_compress_128_online_prefill = host::compress::plan_online_prefill;
[[maybe_unused]]
constexpr auto& plan_compress_128_online_prefill = host::compress::plan_online_prefill;
} // namespace } // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/online_c128_mtp.cuh
+39 -21
View File
@@ -4,11 +4,13 @@
#include <sgl_kernel/utils.h> #include <sgl_kernel/utils.h>
#include <sgl_kernel/runtime.cuh> #include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/type.cuh>
#include <dlpack/dlpack.h> #include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h> #include <tvm/ffi/container/tensor.h>
#include <cstdint> #include <cstdint>
#include <type_traits>
namespace { namespace {
@@ -17,14 +19,14 @@ SGL_DEVICE int64_t clamp_accept_len(int64_t delta, int64_t max_accept) {
return delta < max_accept ? delta : max_accept; return delta < max_accept ? delta : max_accept;
} }
template <typename TSeq, typename TReq> template <typename TSeq, typename TReq, typename BufferFloat>
struct OnlineC128MTPWritePrefixParams { struct OnlineC128MTPWritePrefixParams {
const float* __restrict__ kv_score_input; const float* __restrict__ kv_score_input;
const TSeq* __restrict__ seq_lens; const TSeq* __restrict__ seq_lens;
const TReq* __restrict__ req_pool_indices; const TReq* __restrict__ req_pool_indices;
const int32_t* __restrict__ req_to_token; const int32_t* __restrict__ req_to_token;
const float* __restrict__ ape; const float* __restrict__ ape;
float* __restrict__ state; BufferFloat* __restrict__ state;
int64_t kv_score_stride_b; int64_t kv_score_stride_b;
int64_t req_to_token_stride_b; int64_t req_to_token_stride_b;
int64_t ape_stride_r; int64_t ape_stride_r;
@@ -43,13 +45,13 @@ struct OnlineC128MTPMarkPendingParams {
int64_t max_num_reqs; int64_t max_num_reqs;
}; };
template <typename TSeq, typename TReq> template <typename TSeq, typename TReq, typename BufferFloat>
struct OnlineC128MTPCommitPendingParams { struct OnlineC128MTPCommitPendingParams {
const TSeq* __restrict__ cur_seq_lens; const TSeq* __restrict__ cur_seq_lens;
const TReq* __restrict__ cur_req_pool_indices; const TReq* __restrict__ cur_req_pool_indices;
const int32_t* __restrict__ req_to_token; const int32_t* __restrict__ req_to_token;
const int64_t* __restrict__ pending_seq_lens; const int64_t* __restrict__ pending_seq_lens;
float* __restrict__ state; BufferFloat* __restrict__ state;
int64_t cur_bs; int64_t cur_bs;
int64_t req_to_token_stride_b; int64_t req_to_token_stride_b;
int64_t state_stride_b; int64_t state_stride_b;
@@ -73,8 +75,9 @@ __global__ void online_c128_mtp_mark_pending_kernel(const OnlineC128MTPMarkPendi
} }
} }
template <int64_t kHeadDim, typename TSeq, typename TReq> template <int64_t kHeadDim, typename TSeq, typename TReq, typename BufferFloat>
__global__ void online_c128_mtp_commit_pending_kernel(const OnlineC128MTPCommitPendingParams<TSeq, TReq> params) { __global__ void
online_c128_mtp_commit_pending_kernel(const OnlineC128MTPCommitPendingParams<TSeq, TReq, BufferFloat> params) {
const int64_t bid = static_cast<int64_t>(blockIdx.x); const int64_t bid = static_cast<int64_t>(blockIdx.x);
if (bid >= params.cur_bs) return; if (bid >= params.cur_bs) return;
@@ -91,16 +94,17 @@ __global__ void online_c128_mtp_commit_pending_kernel(const OnlineC128MTPCommitP
if ((final_seq & 127) == 0) return; if ((final_seq & 127) == 0) return;
const int64_t slot = req; const int64_t slot = req;
const float* const src = params.state + (slot + accept * params.state_slot_stride) * params.state_stride_b; const BufferFloat* const src = params.state + (slot + accept * params.state_slot_stride) * params.state_stride_b;
float* const dst = params.state + slot * params.state_stride_b; BufferFloat* 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) { for (int64_t d = static_cast<int64_t>(threadIdx.x); d < kHeadDim * 3; d += blockDim.x) {
dst[d] = src[d]; dst[d] = src[d];
} }
} }
template <int64_t kHeadDim, typename TSeq, typename TReq> template <int64_t kHeadDim, typename TSeq, typename TReq, typename BufferFloat>
__global__ void online_c128_mtp_write_prefix_kernel(const OnlineC128MTPWritePrefixParams<TSeq, TReq> params) { __global__ void
online_c128_mtp_write_prefix_kernel(const OnlineC128MTPWritePrefixParams<TSeq, TReq, BufferFloat> params) {
const int64_t bid = static_cast<int64_t>(blockIdx.x); const int64_t bid = static_cast<int64_t>(blockIdx.x);
if (bid >= params.layer_bs) return; if (bid >= params.layer_bs) return;
@@ -119,10 +123,17 @@ __global__ void online_c128_mtp_write_prefix_kernel(const OnlineC128MTPWritePref
float run_sum = 0.0f; float run_sum = 0.0f;
float run_kv = 0.0f; float run_kv = 0.0f;
if (has_partial) { if (has_partial) {
if constexpr (std::is_same_v<BufferFloat, float>) {
const float* const init = params.state + init_slot * params.state_stride_b; const float* const init = params.state + init_slot * params.state_stride_b;
run_max = init[d]; run_max = init[d];
run_sum = init[kHeadDim + d]; run_sum = init[kHeadDim + d];
run_kv = init[kHeadDim * 2 + d]; run_kv = init[kHeadDim * 2 + d];
} else {
const BufferFloat* const init = params.state + init_slot * params.state_stride_b;
run_max = device::cast<float>(init[d]);
run_sum = device::cast<float>(init[kHeadDim + d]);
run_kv = device::cast<float>(init[kHeadDim * 2 + d]);
}
} }
constexpr int kMaxVerifyTokens = 8; constexpr int kMaxVerifyTokens = 8;
@@ -163,10 +174,17 @@ __global__ void online_c128_mtp_write_prefix_kernel(const OnlineC128MTPWritePref
const int64_t final_seq = seq_before + step + 1; const int64_t final_seq = seq_before + step + 1;
if ((final_seq & 127) != 0) { if ((final_seq & 127) != 0) {
const int64_t slot = req_idx + (step + 1) * params.state_slot_stride; const int64_t slot = req_idx + (step + 1) * params.state_slot_stride;
if constexpr (std::is_same_v<BufferFloat, float>) {
float* const out = params.state + slot * params.state_stride_b; float* const out = params.state + slot * params.state_stride_b;
out[d] = run_max; out[d] = run_max;
out[kHeadDim + d] = run_sum; out[kHeadDim + d] = run_sum;
out[kHeadDim * 2 + d] = run_kv; out[kHeadDim * 2 + d] = run_kv;
} else {
BufferFloat* const out = params.state + slot * params.state_stride_b;
out[d] = device::cast<BufferFloat>(run_max);
out[kHeadDim + d] = device::cast<BufferFloat>(run_sum);
out[kHeadDim * 2 + d] = device::cast<BufferFloat>(run_kv);
}
} }
if (pos == 127) { if (pos == 127) {
@@ -177,7 +195,7 @@ __global__ void online_c128_mtp_write_prefix_kernel(const OnlineC128MTPWritePref
} }
} }
template <int64_t kHeadDim, typename TSeq, typename TReq> template <int64_t kHeadDim, typename TSeq, typename TReq, typename BufferFloat>
struct OnlineC128MTPWritePrefixKernel { struct OnlineC128MTPWritePrefixKernel {
static void launch( static void launch(
tvm::ffi::TensorView kv_score_input, tvm::ffi::TensorView kv_score_input,
@@ -192,13 +210,13 @@ struct OnlineC128MTPWritePrefixKernel {
DLDevice device) { DLDevice device) {
using namespace host; using namespace host;
const auto params = OnlineC128MTPWritePrefixParams<TSeq, TReq>{ const auto params = OnlineC128MTPWritePrefixParams<TSeq, TReq, BufferFloat>{
.kv_score_input = static_cast<const float*>(kv_score_input.data_ptr()), .kv_score_input = static_cast<const float*>(kv_score_input.data_ptr()),
.seq_lens = static_cast<const TSeq*>(seq_lens.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_pool_indices = static_cast<const TReq*>(req_pool_indices.data_ptr()),
.req_to_token = static_cast<const int32_t*>(req_to_token.data_ptr()), .req_to_token = static_cast<const int32_t*>(req_to_token.data_ptr()),
.ape = static_cast<const float*>(ape.data_ptr()), .ape = static_cast<const float*>(ape.data_ptr()),
.state = static_cast<float*>(state.data_ptr()), .state = static_cast<BufferFloat*>(state.data_ptr()),
.kv_score_stride_b = kv_score_input.stride(0), .kv_score_stride_b = kv_score_input.stride(0),
.req_to_token_stride_b = req_to_token.stride(0), .req_to_token_stride_b = req_to_token.stride(0),
.ape_stride_r = ape.stride(0), .ape_stride_r = ape.stride(0),
@@ -211,7 +229,7 @@ struct OnlineC128MTPWritePrefixKernel {
static_assert(kHeadDim == 512, "online c128 MTP write-prefix only supports head_dim=512"); static_assert(kHeadDim == 512, "online c128 MTP write-prefix only supports head_dim=512");
constexpr uint32_t kThreads = static_cast<uint32_t>(kHeadDim); constexpr uint32_t kThreads = static_cast<uint32_t>(kHeadDim);
LaunchKernel(static_cast<uint32_t>(layer_bs), kThreads, device)( LaunchKernel(static_cast<uint32_t>(layer_bs), kThreads, device)(
online_c128_mtp_write_prefix_kernel<kHeadDim, TSeq, TReq>, params); online_c128_mtp_write_prefix_kernel<kHeadDim, TSeq, TReq, BufferFloat>, params);
} }
static void static void
@@ -234,7 +252,7 @@ struct OnlineC128MTPWritePrefixKernel {
TensorMatcher({-1}).with_dtype<TReq>().with_device(device).verify(req_pool_indices); TensorMatcher({-1}).with_dtype<TReq>().with_device(device).verify(req_pool_indices);
TensorMatcher({-1, -1}).with_dtype<int32_t>().with_device(device).verify(req_to_token); TensorMatcher({-1, -1}).with_dtype<int32_t>().with_device(device).verify(req_to_token);
TensorMatcher({128, kHeadDim}).with_dtype<float>().with_device(device).verify(ape); TensorMatcher({128, kHeadDim}).with_dtype<float>().with_device(device).verify(ape);
TensorMatcher({-1, kHeadDim * 3}).with_dtype<float>().with_device(device).verify(state); TensorMatcher({-1, kHeadDim * 3}).with_dtype<BufferFloat>().with_device(device).verify(state);
if (layer_bs <= 0) return; if (layer_bs <= 0) return;
RuntimeCheck(num_verify_tokens > 0 && num_verify_tokens <= 8, "unsupported num_verify_tokens=", num_verify_tokens); RuntimeCheck(num_verify_tokens > 0 && num_verify_tokens <= 8, "unsupported num_verify_tokens=", num_verify_tokens);
@@ -257,7 +275,7 @@ struct OnlineC128MTPWritePrefixKernel {
} }
}; };
template <int64_t kHeadDim, typename TSeq, typename TReq> template <int64_t kHeadDim, typename TSeq, typename TReq, typename BufferFloat>
struct OnlineC128MTPMarkPendingKernel { struct OnlineC128MTPMarkPendingKernel {
static void launch( static void launch(
tvm::ffi::TensorView seq_lens, tvm::ffi::TensorView seq_lens,
@@ -308,7 +326,7 @@ struct OnlineC128MTPMarkPendingKernel {
} }
}; };
template <int64_t kHeadDim, typename TSeq, typename TReq> template <int64_t kHeadDim, typename TSeq, typename TReq, typename BufferFloat>
struct OnlineC128MTPCommitPendingKernel { struct OnlineC128MTPCommitPendingKernel {
static void launch( static void launch(
tvm::ffi::TensorView cur_seq_lens, tvm::ffi::TensorView cur_seq_lens,
@@ -323,12 +341,12 @@ struct OnlineC128MTPCommitPendingKernel {
DLDevice device) { DLDevice device) {
using namespace host; using namespace host;
const auto params = OnlineC128MTPCommitPendingParams<TSeq, TReq>{ const auto params = OnlineC128MTPCommitPendingParams<TSeq, TReq, BufferFloat>{
.cur_seq_lens = static_cast<const TSeq*>(cur_seq_lens.data_ptr()), .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()), .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()), .req_to_token = static_cast<const int32_t*>(req_to_token.data_ptr()),
.pending_seq_lens = static_cast<const int64_t*>(pending_seq_lens.data_ptr()), .pending_seq_lens = static_cast<const int64_t*>(pending_seq_lens.data_ptr()),
.state = static_cast<float*>(state.data_ptr()), .state = static_cast<BufferFloat*>(state.data_ptr()),
.cur_bs = cur_bs, .cur_bs = cur_bs,
.req_to_token_stride_b = req_to_token.stride(0), .req_to_token_stride_b = req_to_token.stride(0),
.state_stride_b = state.stride(0), .state_stride_b = state.stride(0),
@@ -339,7 +357,7 @@ struct OnlineC128MTPCommitPendingKernel {
constexpr uint32_t kThreads = 256; constexpr uint32_t kThreads = 256;
LaunchKernel(static_cast<uint32_t>(cur_bs), kThreads, device)( LaunchKernel(static_cast<uint32_t>(cur_bs), kThreads, device)(
online_c128_mtp_commit_pending_kernel<kHeadDim, TSeq, TReq>, params); online_c128_mtp_commit_pending_kernel<kHeadDim, TSeq, TReq, BufferFloat>, params);
} }
static void static void
@@ -361,7 +379,7 @@ struct OnlineC128MTPCommitPendingKernel {
TensorMatcher({-1}).with_dtype<TReq>().with_device(device).verify(cur_req_pool_indices); TensorMatcher({-1}).with_dtype<TReq>().with_device(device).verify(cur_req_pool_indices);
TensorMatcher({-1, -1}).with_dtype<int32_t>().with_device(device).verify(req_to_token); TensorMatcher({-1, -1}).with_dtype<int32_t>().with_device(device).verify(req_to_token);
TensorMatcher({-1}).with_dtype<int64_t>().with_device(device).verify(pending_seq_lens); 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); TensorMatcher({-1, kHeadDim * 3}).with_dtype<BufferFloat>().with_device(device).verify(state);
if (cur_bs <= 0) return; if (cur_bs <= 0) return;
RuntimeCheck(num_verify_tokens > 0 && num_verify_tokens <= 8, "unsupported num_verify_tokens=", num_verify_tokens); RuntimeCheck(num_verify_tokens > 0 && num_verify_tokens <= 8, "unsupported num_verify_tokens=", num_verify_tokens);
python/sglang/jit_kernel/dsv4/compress.py
+5 -3
View File
@@ -66,9 +66,11 @@ def _jit_compress_module(
@cache_once @cache_once
def _jit_compress_128_online_module(head_dim: int) -> Module: def _jit_compress_128_online_module(
head_dim: int, dtype_buffer: torch.dtype = torch.float32
) -> Module:
assert head_dim == 512 assert head_dim == 512
args = make_cpp_args(head_dim, is_arch_support_pdl()) args = make_cpp_args(head_dim, dtype_buffer, is_arch_support_pdl())
kernel_class = f"FlashCompress128OnlineKernel<{args}>" kernel_class = f"FlashCompress128OnlineKernel<{args}>"
return load_jit( return load_jit(
make_name(f"compress_128_online_v2"), make_name(f"compress_128_online_v2"),
@@ -327,7 +329,7 @@ def compress_forward(
assert plan.compress_ratio == compress_ratio assert plan.compress_ratio == compress_ratio
if is_online: if is_online:
assert compress_ratio == 128 and head_dim == 512 assert compress_ratio == 128 and head_dim == 512
module = _jit_compress_128_online_module(512) module = _jit_compress_128_online_module(512, kv_score_buffer.dtype)
else: else:
dtype_in, dtype_out = kv_score_input.dtype, out.dtype dtype_in, dtype_out = kv_score_input.dtype, out.dtype
module = _jit_compress_module( module = _jit_compress_module(
python/sglang/jit_kernel/dsv4/online_c128_mtp.py
+22 -7
View File
@@ -15,9 +15,12 @@ if TYPE_CHECKING:
@cache_once @cache_once
def _jit_online_c128_mtp_module( def _jit_online_c128_mtp_module(
head_dim: int, seq_dtype: torch.dtype, req_dtype: torch.dtype head_dim: int,
seq_dtype: torch.dtype,
req_dtype: torch.dtype,
dtype_buffer: torch.dtype,
) -> Module: ) -> Module:
args = make_cpp_args(head_dim, seq_dtype, req_dtype) args = make_cpp_args(head_dim, seq_dtype, req_dtype, dtype_buffer)
return load_jit( return load_jit(
make_name(f"online_c128_mtp_{head_dim}"), make_name(f"online_c128_mtp_{head_dim}"),
*args, *args,
@@ -35,6 +38,7 @@ def _jit_online_c128_mtp_module(
class _OnlineC128LayerRuntime: class _OnlineC128LayerRuntime:
head_dim: int head_dim: int
main_state: torch.Tensor main_state: torch.Tensor
state_dtype: torch.dtype
state_slot_offset: int state_slot_offset: int
@@ -76,11 +80,12 @@ class OnlineC128MTPController:
seq_lens=seq_lens.detach(), seq_lens=seq_lens.detach(),
) )
head_dim = self._head_dim() head_dim = self._head_dim()
if head_dim is None or self._num_verify_tokens() == 0: state_dtype = self._state_dtype()
if head_dim is None or state_dtype is None or self._num_verify_tokens() == 0:
return return
token_to_kv_pool = self.backend.token_to_kv_pool token_to_kv_pool = self.backend.token_to_kv_pool
_jit_online_c128_mtp_module( _jit_online_c128_mtp_module(
head_dim, seq_lens.dtype, req_pool_indices.dtype head_dim, seq_lens.dtype, req_pool_indices.dtype, state_dtype
).mark_pending( ).mark_pending(
seq_lens, seq_lens,
req_pool_indices, req_pool_indices,
@@ -156,20 +161,21 @@ class OnlineC128MTPController:
token_to_kv_pool = self.backend.token_to_kv_pool token_to_kv_pool = self.backend.token_to_kv_pool
head_dim = compressor.head_dim head_dim = compressor.head_dim
state_pool = token_to_kv_pool.get_attention_compress_states(layer_id) state_pool = token_to_kv_pool.get_attention_compress_states(layer_id)
state = state_pool.kv_score_buffer.kv_score
total_bs = kv_score_input.numel() // (num_verify_tokens * head_dim * 2) 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) layer_bs = min(ctx.seq_lens.shape[0], ctx.req_pool_indices.shape[0], total_bs)
if layer_bs <= 0: if layer_bs <= 0:
return return
_jit_online_c128_mtp_module( _jit_online_c128_mtp_module(
head_dim, ctx.seq_lens.dtype, ctx.req_pool_indices.dtype head_dim, ctx.seq_lens.dtype, ctx.req_pool_indices.dtype, state.dtype
).write_prefix_states( ).write_prefix_states(
kv_score_input, kv_score_input,
ctx.seq_lens, ctx.seq_lens,
ctx.req_pool_indices, ctx.req_pool_indices,
self.backend.req_to_token, self.backend.req_to_token,
compressor.ape.reshape(128, head_dim), compressor.ape.reshape(128, head_dim),
state_pool.kv_score_buffer.kv_score, state,
layer_bs, layer_bs,
num_verify_tokens, num_verify_tokens,
state_pool.online_mtp_state_slot_offset, state_pool.online_mtp_state_slot_offset,
@@ -200,7 +206,10 @@ class OnlineC128MTPController:
for runtime in self._iter_layer_runtimes(): for runtime in self._iter_layer_runtimes():
_jit_online_c128_mtp_module( _jit_online_c128_mtp_module(
runtime.head_dim, seq_lens.dtype, req_pool_indices.dtype runtime.head_dim,
seq_lens.dtype,
req_pool_indices.dtype,
runtime.state_dtype,
).commit_pending( ).commit_pending(
seq_lens, seq_lens,
req_pool_indices, req_pool_indices,
@@ -239,6 +248,11 @@ class OnlineC128MTPController:
return runtime.head_dim return runtime.head_dim
return None return None
def _state_dtype(self) -> Optional[torch.dtype]:
for runtime in self._iter_layer_runtimes():
return runtime.state_dtype
return None
def _iter_layer_runtimes(self): def _iter_layer_runtimes(self):
if self._layer_runtimes is None: if self._layer_runtimes is None:
runtimes = [] runtimes = []
@@ -255,6 +269,7 @@ class OnlineC128MTPController:
_OnlineC128LayerRuntime( _OnlineC128LayerRuntime(
head_dim=compressor.head_dim, head_dim=compressor.head_dim,
main_state=state_pool.kv_score_buffer.kv_score, main_state=state_pool.kv_score_buffer.kv_score,
state_dtype=state_pool.kv_score_buffer.kv_score.dtype,
state_slot_offset=state_pool.online_mtp_state_slot_offset, state_slot_offset=state_pool.online_mtp_state_slot_offset,
) )
) )
python/sglang/srt/mem_cache/kv_cache_configurator.py
-5
View File
@@ -74,11 +74,6 @@ def _get_dsv4_compress_state_dtypes() -> tuple[torch.dtype, torch.dtype]:
if dtype_name in ("float32", "fp32"): if dtype_name in ("float32", "fp32"):
return torch.float32, torch.float32 return torch.float32, torch.float32
if dtype_name in ("bfloat16", "bf16"): if dtype_name in ("bfloat16", "bf16"):
if envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get():
raise ValueError(
"SGLANG_DSV4_COMPRESS_STATE_DTYPE=bf16 is not supported when "
"SGLANG_OPT_USE_ONLINE_COMPRESS=1; online c128 state must stay float32."
)
return torch.bfloat16, torch.bfloat16 return torch.bfloat16, torch.bfloat16
raise ValueError( raise ValueError(
"Unsupported SGLANG_DSV4_COMPRESS_STATE_DTYPE=" "Unsupported SGLANG_DSV4_COMPRESS_STATE_DTYPE="
python/sglang/srt/model_executor/pool_configurator.py
-5
View File
@@ -78,11 +78,6 @@ def _get_dsv4_compress_state_dtype_sizes() -> tuple[int, int]:
if dtype_name in ("float32", "fp32"): if dtype_name in ("float32", "fp32"):
return 4, 4 return 4, 4
if dtype_name in ("bfloat16", "bf16"): if dtype_name in ("bfloat16", "bf16"):
if envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get():
raise ValueError(
"SGLANG_DSV4_COMPRESS_STATE_DTYPE=bf16 is not supported when "
"SGLANG_OPT_USE_ONLINE_COMPRESS=1; online c128 state must stay float32."
)
return 2, 2 return 2, 2
raise ValueError( raise ValueError(
"Unsupported SGLANG_DSV4_COMPRESS_STATE_DTYPE=" "Unsupported SGLANG_DSV4_COMPRESS_STATE_DTYPE="
test/registered/jit/benchmark/bench_online_c128_mtp.py
+1 -1
View File
@@ -148,7 +148,7 @@ def benchmark(
) -> tuple[float, float, float]: ) -> tuple[float, float, float]:
case = make_case(batch_size, num_verify_tokens) case = make_case(batch_size, num_verify_tokens)
module = _jit_online_c128_mtp_module( module = _jit_online_c128_mtp_module(
HEAD_DIM, case.seq_lens.dtype, case.req_pool_indices.dtype HEAD_DIM, case.seq_lens.dtype, case.req_pool_indices.dtype, case.state.dtype
) )
fn = lambda: call_write_prefix(module, case) fn = lambda: call_write_prefix(module, case)