Port KV Compression V2 from deepseek_v4_dev (#24890)
Co-authored-by: Cheng Wan <chwan@rice.edu> Co-authored-by: DarkSharpness <2040703891@qq.com>
This commit is contained in:
Yuhao Yang
co-authored by
Cheng Wan
DarkSharpness
Cheng Wan
DarkSharpness
parent
d0913fca8d
commit
e2290b155a
@@ -0,0 +1,875 @@
|
||||
#include <sgl_kernel/ffi.h>
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/runtime.cuh>
|
||||
#include <sgl_kernel/tile.cuh>
|
||||
#include <sgl_kernel/type.cuh>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
#include <sgl_kernel/warp.cuh>
|
||||
|
||||
#include <sgl_kernel/deepseek_v4/compress_v2.cuh>
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
#include <tvm/ffi/container/tuple.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cfloat>
|
||||
#include <cstdint>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
|
||||
namespace {
|
||||
|
||||
using PlanD = device::compress::DecodePlan;
|
||||
using PlanC = device::compress::CompressPlan;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Decode kernel: 1 token / batch. Each block handles one batch.
|
||||
// 4 elements per thread -> kBlockSize = head_dim / 4.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
struct Compress128OnlineDecodeParams {
|
||||
void* __restrict__ kv_score_buffer; // [num_slots, 1, head_dim * 3]
|
||||
const void* __restrict__ kv_score_input; // [batch_size, head_dim * 2]
|
||||
void* __restrict__ kv_compressed_output; // [batch_size, head_dim]
|
||||
const void* __restrict__ score_bias; // [128, head_dim]
|
||||
const PlanD* __restrict__ plan_d;
|
||||
uint32_t batch_size;
|
||||
};
|
||||
|
||||
template <int64_t kHeadDim, bool kUsePDL>
|
||||
__global__ void flash_c128_online_decode_v2(const __grid_constant__ Compress128OnlineDecodeParams params) {
|
||||
using namespace device;
|
||||
constexpr uint32_t kVecSize = 4;
|
||||
constexpr uint32_t kBlockSize = kHeadDim / kVecSize;
|
||||
using Vec = AlignedVector<float, kVecSize>;
|
||||
const auto gmem = tile::Memory<Vec>::cta(kBlockSize);
|
||||
const auto batch_id = blockIdx.x;
|
||||
if (batch_id >= params.batch_size) return;
|
||||
|
||||
// Wait for the plan-finalize kernel to publish `plan.read_page_0 / write_loc`
|
||||
// before reading the plan. The plan kernel runs on the same stream and does
|
||||
// NOT issue a PDL trigger, so launching this kernel with PDL means our
|
||||
// pre-wait global reads can race with the plan kernel's writes.
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
|
||||
const auto plan = params.plan_d[batch_id];
|
||||
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_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_store_buf = kv_score_buffer + plan.write_loc * (kHeadDim * 3);
|
||||
const auto kv_src = kv_score_input + batch_id * (kHeadDim * 2);
|
||||
|
||||
// Buffer layout: [max | sum | kv] (slot 0 / 1 / 2 of the head_dim*3 row).
|
||||
const auto new_kv_vec = gmem.load(kv_src, 0);
|
||||
const auto new_score_raw_vec = gmem.load(kv_src, 1);
|
||||
const auto bias_vec = gmem.load(params.score_bias, pos_in_chunk);
|
||||
|
||||
Vec out_kv_vec;
|
||||
Vec out_max_vec;
|
||||
Vec out_sum_vec;
|
||||
if (pos_in_chunk != 0) {
|
||||
// Mid-chunk: combine prior partial state with the new token.
|
||||
const auto max_score_vec = gmem.load(kv_load_buf, 0);
|
||||
const auto sum_score_vec = gmem.load(kv_load_buf, 1);
|
||||
const auto old_kv_vec = gmem.load(kv_load_buf, 2);
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kVecSize; ++i) {
|
||||
const auto old_max = max_score_vec[i];
|
||||
const auto old_kv = old_kv_vec[i];
|
||||
const auto new_score = new_score_raw_vec[i] + bias_vec[i];
|
||||
const auto new_kv = new_kv_vec[i];
|
||||
const auto new_max = fmaxf(old_max, new_score);
|
||||
const auto old_sum = sum_score_vec[i] * expf(old_max - new_max);
|
||||
const auto new_exp = expf(new_score - new_max);
|
||||
const auto new_sum = old_sum + new_exp;
|
||||
out_kv_vec[i] = (old_kv * old_sum + new_kv * new_exp) / new_sum;
|
||||
out_max_vec[i] = new_max;
|
||||
out_sum_vec[i] = new_sum;
|
||||
}
|
||||
} else {
|
||||
// First token of a new chunk: state == this token alone.
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kVecSize; ++i) {
|
||||
out_kv_vec[i] = new_kv_vec[i];
|
||||
out_max_vec[i] = new_score_raw_vec[i] + bias_vec[i];
|
||||
out_sum_vec[i] = 1.0f;
|
||||
}
|
||||
}
|
||||
|
||||
if (pos_in_chunk == 127) {
|
||||
// Chunk just closed: emit compressed kv, no buffer update.
|
||||
const auto kv_out = static_cast<float*>(params.kv_compressed_output) + batch_id * kHeadDim;
|
||||
gmem.store(kv_out, out_kv_vec);
|
||||
} else {
|
||||
gmem.store(kv_store_buf, out_max_vec, 0);
|
||||
gmem.store(kv_store_buf, out_sum_vec, 1);
|
||||
gmem.store(kv_store_buf, out_kv_vec, 2);
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Prefill kernel: 1 segment / block. Two passes (compress + write) share the
|
||||
// kernel template, parameterized by `kWrite`.
|
||||
// 16 warps per block; each warp handles 8 of the 128 chunk positions.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
constexpr int32_t kTileElements = 2; // split along head-dim
|
||||
constexpr int32_t kElementsPerWarp = 8; // split along the 128-chunk
|
||||
constexpr uint32_t kNumWarps = 128 / kElementsPerWarp;
|
||||
constexpr uint32_t kPrefillBlockSize = device::kWarpThreads * kNumWarps;
|
||||
using PrefillStorage = device::AlignedVector<float, kTileElements>;
|
||||
|
||||
struct Compress128OnlinePrefillParams {
|
||||
void* __restrict__ kv_score_buffer; // [num_slots, 1, head_dim * 3]
|
||||
const void* __restrict__ kv_score_input; // [num_q_tokens, head_dim * 2]
|
||||
void* __restrict__ kv_compressed_output; // [num_compress, head_dim]
|
||||
const void* __restrict__ score_bias; // [128, head_dim]
|
||||
const PlanC* __restrict__ plan_c; // close-chunk segments
|
||||
const PlanC* __restrict__ plan_w; // trailing partial segments
|
||||
uint32_t num_compress;
|
||||
uint32_t num_write;
|
||||
};
|
||||
|
||||
struct Compress128SharedBuffer {
|
||||
using Storage = device::AlignedVector<float, 4>;
|
||||
Storage data[kNumWarps][device::kWarpThreads + 1]; // +1 to avoid bank conflict
|
||||
SGL_DEVICE Storage& operator()(uint32_t warp_id, uint32_t lane_id) {
|
||||
return data[warp_id][lane_id];
|
||||
}
|
||||
SGL_DEVICE float& operator()(uint32_t warp_id, uint32_t lane_id, uint32_t tile_id) {
|
||||
return data[warp_id][lane_id][tile_id];
|
||||
}
|
||||
};
|
||||
|
||||
/// \brief Sentinel score for padded positions in a 128-segment.
|
||||
constexpr float kPadScore = -FLT_MAX;
|
||||
|
||||
[[maybe_unused]]
|
||||
SGL_DEVICE void c128_prefill_segment_softmax(
|
||||
const PrefillStorage (&kv)[kElementsPerWarp],
|
||||
const PrefillStorage (&score)[kElementsPerWarp],
|
||||
float* seg_kv,
|
||||
float* seg_max,
|
||||
float* seg_sum,
|
||||
const uint32_t warp_id,
|
||||
const uint32_t lane_id) {
|
||||
using namespace device;
|
||||
|
||||
// Per-warp running state (max, sum, kv) for kTileElements head-dim slots.
|
||||
using TmpStorage = typename Compress128SharedBuffer::Storage;
|
||||
__shared__ Compress128SharedBuffer s_local_val_max;
|
||||
__shared__ Compress128SharedBuffer s_local_exp_sum;
|
||||
__shared__ Compress128SharedBuffer s_local_product;
|
||||
|
||||
TmpStorage tmp_val_max;
|
||||
TmpStorage tmp_exp_sum;
|
||||
TmpStorage tmp_product;
|
||||
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < kTileElements; ++i) {
|
||||
float score_fp32[kElementsPerWarp];
|
||||
#pragma unroll
|
||||
for (int32_t j = 0; j < kElementsPerWarp; ++j) {
|
||||
score_fp32[j] = score[j][i];
|
||||
}
|
||||
float max_value = score_fp32[0];
|
||||
#pragma unroll
|
||||
for (int32_t j = 1; j < kElementsPerWarp; ++j) {
|
||||
max_value = fmaxf(max_value, score_fp32[j]);
|
||||
}
|
||||
float sum_exp_value = 0.0f;
|
||||
float sum_product = 0.0f;
|
||||
#pragma unroll
|
||||
for (int32_t j = 0; j < kElementsPerWarp; ++j) {
|
||||
const auto exp_score = expf(score_fp32[j] - max_value);
|
||||
sum_product += kv[j][i] * exp_score;
|
||||
sum_exp_value += exp_score;
|
||||
}
|
||||
tmp_val_max[i] = max_value;
|
||||
tmp_exp_sum[i] = sum_exp_value;
|
||||
tmp_product[i] = sum_product;
|
||||
}
|
||||
|
||||
// Aligned writes (no bank conflict thanks to `+1` padding).
|
||||
s_local_val_max(warp_id, lane_id) = tmp_val_max;
|
||||
s_local_exp_sum(warp_id, lane_id) = tmp_exp_sum;
|
||||
s_local_product(warp_id, lane_id) = tmp_product;
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// Cross-warp reduction. Same recipe as c128_online.cuh: each block-thread
|
||||
// pair reduces a (tile_id, lane_id) slot using a kNumWarps-wide warp shuffle.
|
||||
constexpr uint32_t kReductionCount = kTileElements * kWarpThreads * kNumWarps;
|
||||
constexpr uint32_t kIteration = kReductionCount / kPrefillBlockSize;
|
||||
static_assert(kTileElements * kNumWarps == kWarpThreads, "TODO: support other configs");
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kIteration; ++i) {
|
||||
const uint32_t j = i * kPrefillBlockSize + warp_id * kWarpThreads + lane_id;
|
||||
const uint32_t local_warp_id = j % kNumWarps;
|
||||
const uint32_t local_elem_id = j / kNumWarps;
|
||||
const uint32_t local_tile_id = local_elem_id % kTileElements;
|
||||
const uint32_t local_lane_id = local_elem_id / kTileElements;
|
||||
const auto local_val_max = s_local_val_max(local_warp_id, local_lane_id, local_tile_id);
|
||||
const auto local_exp_sum = s_local_exp_sum(local_warp_id, local_lane_id, local_tile_id);
|
||||
const auto local_product = s_local_product(local_warp_id, local_lane_id, local_tile_id);
|
||||
const auto global_val_max = warp::reduce_max<kNumWarps>(local_val_max);
|
||||
const auto rescale = expf(local_val_max - global_val_max);
|
||||
const auto global_exp_sum = warp::reduce_sum<kNumWarps>(local_exp_sum * rescale);
|
||||
const auto final_scale = rescale / global_exp_sum;
|
||||
const auto global_product = warp::reduce_sum<kNumWarps>(local_product * final_scale);
|
||||
seg_kv[local_elem_id] = global_product;
|
||||
seg_max[local_elem_id] = global_val_max;
|
||||
seg_sum[local_elem_id] = global_exp_sum;
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
/// \brief Online compress 128 prefill v2.
|
||||
///
|
||||
/// `kWrite=false` (compress pass): handles segments that close a 128-chunk.
|
||||
/// 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`.
|
||||
template <int64_t kHeadDim, bool kWrite, bool kUsePDL>
|
||||
__global__ __launch_bounds__(kPrefillBlockSize, 2) //
|
||||
void flash_c128_online_prefill_v2(const __grid_constant__ Compress128OnlinePrefillParams params) {
|
||||
using namespace device;
|
||||
|
||||
constexpr int64_t kTileDim = kTileElements * kWarpThreads; // 64
|
||||
constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
|
||||
static_assert(kHeadDim % kTileDim == 0);
|
||||
|
||||
// Compile-time fold to the right plan list.
|
||||
const auto num_plans = kWrite ? params.num_write : params.num_compress;
|
||||
const auto plan_ptr = kWrite ? params.plan_w : params.plan_c;
|
||||
const uint32_t global_id = blockIdx.x;
|
||||
const uint32_t global_pid = global_id / kNumSplit;
|
||||
const uint32_t global_sid = global_id % kNumSplit;
|
||||
if (global_pid >= num_plans) return;
|
||||
|
||||
const uint32_t warp_id = threadIdx.x / kWarpThreads;
|
||||
const uint32_t lane_id = threadIdx.x % kWarpThreads;
|
||||
const int32_t split_offset = global_sid * kTileDim;
|
||||
|
||||
// The previous kernel (plan-finalize stage 1) does NOT issue a PDL trigger,
|
||||
// so PDLWaitPrimary effectively waits for stage 1 to complete. Read the plan
|
||||
// AFTER the wait so the freshly-written `read_page_0` (= state-pool slot) is
|
||||
// visible. Reading it before the wait is a real race -- with PDL enabled the
|
||||
// kernel can begin executing before stage 1's stores propagate, and we'd see
|
||||
// the stage-0 batch_id placeholder in `read_page_0` instead of the slot.
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
|
||||
const auto plan = plan_ptr[global_pid];
|
||||
if (plan.is_invalid()) [[unlikely]]
|
||||
return;
|
||||
|
||||
const auto kv_score_buffer = static_cast<float*>(params.kv_score_buffer);
|
||||
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 score_bias_base = static_cast<const float*>(params.score_bias);
|
||||
|
||||
constexpr int64_t kElementSize = kHeadDim * 2; // | kv | score |
|
||||
|
||||
// The plan stores last-token coordinates; segment start is recoverable as
|
||||
// ragged_id - window_len + 1.
|
||||
const uint32_t window_len = plan.buffer_len;
|
||||
const uint32_t position = plan.seq_len - 1;
|
||||
const uint32_t pos_in_chunk_end = (position % 128u) + 1u; // exclusive, in [1, 128]
|
||||
const uint32_t chunk_offset = pos_in_chunk_end - window_len; // in [0, 127]
|
||||
const int32_t segment_start_ragged = static_cast<int32_t>(plan.ragged_id) - static_cast<int32_t>(position % 128u);
|
||||
|
||||
// --- Stage 1: load kv / score / bias for this warp's 8 chunk positions.
|
||||
PrefillStorage kv[kElementsPerWarp];
|
||||
PrefillStorage score[kElementsPerWarp];
|
||||
PrefillStorage bias[kElementsPerWarp];
|
||||
const uint32_t warp_offset = warp_id * kElementsPerWarp;
|
||||
|
||||
#pragma unroll
|
||||
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 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);
|
||||
score[i].load(score_src_ptr, lane_id);
|
||||
bias[i].load(bias_src_ptr, lane_id);
|
||||
}
|
||||
}
|
||||
|
||||
// --- Stage 2: pad invalid positions. score = -FLT_MAX, kv = 0 (so that
|
||||
// kv * exp(score-max) ??? 0 / 0 cleanly without producing NaN/inf).
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kElementsPerWarp; ++i) {
|
||||
const uint32_t j = i + warp_offset;
|
||||
const bool is_valid = (j >= chunk_offset && j < pos_in_chunk_end);
|
||||
#pragma unroll
|
||||
for (uint32_t ii = 0; ii < kTileElements; ++ii) {
|
||||
score[i][ii] = is_valid ? score[i][ii] + bias[i][ii] : kPadScore;
|
||||
kv[i][ii] = is_valid ? kv[i][ii] : 0.0f;
|
||||
}
|
||||
}
|
||||
|
||||
// --- Stage 3: warp-tile online softmax over the 128-position chunk.
|
||||
__shared__ alignas(16) float seg_kv[kTileDim];
|
||||
__shared__ alignas(16) float seg_max[kTileDim];
|
||||
__shared__ alignas(16) float seg_sum[kTileDim];
|
||||
c128_prefill_segment_softmax(kv, score, seg_kv, seg_max, seg_sum, warp_id, lane_id);
|
||||
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
|
||||
// --- Stage 4: warp 0 folds with prior partial state (if any) and writes.
|
||||
if (warp_id == 0) {
|
||||
PrefillStorage out_kv_vec, out_max_vec, out_sum_vec;
|
||||
out_kv_vec.load(seg_kv, lane_id);
|
||||
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) {
|
||||
// Combine with prior partial state for this slot.
|
||||
const auto buf_load = kv_score_buffer + plan.read_page_0 * (kHeadDim * 3) + split_offset;
|
||||
PrefillStorage buf_max_vec, buf_sum_vec, buf_kv_vec;
|
||||
buf_max_vec.load(buf_load + 0 * kHeadDim, lane_id);
|
||||
buf_sum_vec.load(buf_load + 1 * kHeadDim, lane_id);
|
||||
buf_kv_vec.load(buf_load + 2 * kHeadDim, lane_id);
|
||||
#pragma unroll
|
||||
for (uint32_t ii = 0; ii < kTileElements; ++ii) {
|
||||
const float m1 = buf_max_vec[ii];
|
||||
const float s1 = buf_sum_vec[ii];
|
||||
const float k1 = buf_kv_vec[ii];
|
||||
const float m2 = out_max_vec[ii];
|
||||
const float s2 = out_sum_vec[ii];
|
||||
const float k2 = out_kv_vec[ii];
|
||||
const float new_max = fmaxf(m1, m2);
|
||||
const float new_s1 = s1 * expf(m1 - new_max);
|
||||
const float new_s2 = s2 * expf(m2 - new_max);
|
||||
const float new_sum = new_s1 + new_s2;
|
||||
const float new_kv = (k1 * new_s1 + k2 * new_s2) / new_sum;
|
||||
out_max_vec[ii] = new_max;
|
||||
out_sum_vec[ii] = new_sum;
|
||||
out_kv_vec[ii] = new_kv;
|
||||
}
|
||||
}
|
||||
|
||||
if constexpr (kWrite) {
|
||||
// For trailing-partial segments the load and store slots collapse to the
|
||||
// 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`.
|
||||
const auto buf_store = kv_score_buffer + plan.read_page_0 * (kHeadDim * 3) + split_offset;
|
||||
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 + 2 * kHeadDim)[lane_id] = out_kv_vec;
|
||||
} else {
|
||||
// Compact output: one row per compress plan, indexed by `global_pid`.
|
||||
const auto out_ptr = kv_compressed_output + global_pid * kHeadDim + split_offset;
|
||||
reinterpret_cast<PrefillStorage*>(out_ptr)[lane_id] = out_kv_vec;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 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`
|
||||
// + `kUsePDL`; the dtype is fixed to fp32 for the online state pool.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
template <int64_t kHeadDim, bool kUsePDL>
|
||||
struct FlashCompress128OnlineKernel {
|
||||
static constexpr auto decode_kernel = flash_c128_online_decode_v2<kHeadDim, kUsePDL>;
|
||||
template <bool kWrite>
|
||||
static constexpr auto prefill_kernel = flash_c128_online_prefill_v2<kHeadDim, kWrite, kUsePDL>;
|
||||
static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads; // 64
|
||||
static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
|
||||
static constexpr uint32_t kDecodeBlockSize = kHeadDim / 4;
|
||||
|
||||
static void run_decode(
|
||||
const tvm::ffi::TensorView kv_score_buffer,
|
||||
const tvm::ffi::TensorView kv_score_input,
|
||||
const tvm::ffi::TensorView kv_compressed_output,
|
||||
const tvm::ffi::TensorView ape,
|
||||
const tvm::ffi::TensorView plan_d_) {
|
||||
using namespace host;
|
||||
|
||||
auto B = SymbolicSize{"batch_size"};
|
||||
auto device_ = SymbolicDevice{};
|
||||
device_.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({-1, 1, kHeadDim * 3}) // kv score buffer (max, sum, kv)
|
||||
.with_dtype<float>()
|
||||
.with_device(device_)
|
||||
.verify(kv_score_buffer);
|
||||
TensorMatcher({B, kHeadDim * 2}) // kv score input
|
||||
.with_dtype<float>()
|
||||
.with_device(device_)
|
||||
.verify(kv_score_input);
|
||||
TensorMatcher({B, kHeadDim}) // kv compressed output (sparse by batch_id)
|
||||
.with_dtype<float>()
|
||||
.with_device(device_)
|
||||
.verify(kv_compressed_output);
|
||||
TensorMatcher({128, kHeadDim}) // ape
|
||||
.with_dtype<float>()
|
||||
.with_device(device_)
|
||||
.verify(ape);
|
||||
|
||||
const auto plan_d = compress::verify_plan_d(plan_d_, B, device_);
|
||||
const auto batch_size = static_cast<uint32_t>(B.unwrap());
|
||||
if (batch_size == 0) return;
|
||||
const auto params = Compress128OnlineDecodeParams{
|
||||
.kv_score_buffer = kv_score_buffer.data_ptr(),
|
||||
.kv_score_input = kv_score_input.data_ptr(),
|
||||
.kv_compressed_output = kv_compressed_output.data_ptr(),
|
||||
.score_bias = ape.data_ptr(),
|
||||
.plan_d = plan_d,
|
||||
.batch_size = batch_size,
|
||||
};
|
||||
LaunchKernel(batch_size, kDecodeBlockSize, device_.unwrap()) //
|
||||
.enable_pdl(kUsePDL)(decode_kernel, params);
|
||||
}
|
||||
|
||||
static void run_prefill(
|
||||
const tvm::ffi::TensorView kv_score_buffer,
|
||||
const tvm::ffi::TensorView kv_score_input,
|
||||
const tvm::ffi::TensorView kv_compressed_output,
|
||||
const tvm::ffi::TensorView ape,
|
||||
const tvm::ffi::TensorView plan_c_,
|
||||
const tvm::ffi::TensorView plan_w_) {
|
||||
using namespace host;
|
||||
|
||||
auto N = SymbolicSize{"num_q_tokens"};
|
||||
auto C = SymbolicSize{"num_c_plans"};
|
||||
auto W = SymbolicSize{"num_w_plans"};
|
||||
auto device_ = SymbolicDevice{};
|
||||
device_.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({-1, 1, kHeadDim * 3}) // kv score buffer
|
||||
.with_dtype<float>()
|
||||
.with_device(device_)
|
||||
.verify(kv_score_buffer);
|
||||
TensorMatcher({N, kHeadDim * 2}) // kv score input (ragged)
|
||||
.with_dtype<float>()
|
||||
.with_device(device_)
|
||||
.verify(kv_score_input);
|
||||
TensorMatcher({C, kHeadDim}) // kv compressed output (compact, by plan_c index)
|
||||
.with_dtype<float>()
|
||||
.with_device(device_)
|
||||
.verify(kv_compressed_output);
|
||||
TensorMatcher({128, kHeadDim}) // ape
|
||||
.with_dtype<float>()
|
||||
.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.
|
||||
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();
|
||||
const auto num_q_tokens = static_cast<uint32_t>(N.unwrap());
|
||||
const auto num_c = static_cast<uint32_t>(C.unwrap());
|
||||
const auto num_w = static_cast<uint32_t>(W.unwrap());
|
||||
RuntimeCheck(num_q_tokens >= num_w, "invalid prefill plan: num_q < num_w");
|
||||
const auto params = Compress128OnlinePrefillParams{
|
||||
.kv_score_buffer = kv_score_buffer.data_ptr(),
|
||||
.kv_score_input = kv_score_input.data_ptr(),
|
||||
.kv_compressed_output = kv_compressed_output.data_ptr(),
|
||||
.score_bias = ape.data_ptr(),
|
||||
.plan_c = plan_c,
|
||||
.plan_w = plan_w,
|
||||
.num_compress = num_c,
|
||||
.num_write = num_w,
|
||||
};
|
||||
|
||||
// The two passes MUST be serialized in stream order: pass 1 reads slots
|
||||
// that pass 2 may write to; running them in parallel would race.
|
||||
if (const auto num_c_blocks = num_c * kNumSplit) {
|
||||
LaunchKernel(num_c_blocks, kPrefillBlockSize, device) //
|
||||
.enable_pdl(kUsePDL)(prefill_kernel</*kWrite=*/false>, params);
|
||||
}
|
||||
if (const auto num_w_blocks = num_w * kNumSplit) {
|
||||
LaunchKernel(num_w_blocks, kPrefillBlockSize, device) //
|
||||
.enable_pdl(kUsePDL)(prefill_kernel</*kWrite=*/true>, params);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
// ===========================================================================
|
||||
// Plan builders. Mirrors the offline v2 pattern (`c_plan.cuh`):
|
||||
// - Decode: a single GPU kernel reads seq_lens / req_to_token /
|
||||
// req_pool_indices on device and emits the final PlanD tensor in one go.
|
||||
// - Prefill: stage 0 (host, on CPU pinned memory) splits each batch's
|
||||
// extend range into per-chunk segments and emits PlanC entries with the
|
||||
// batch_id stashed in `read_page_0` as a placeholder. Stage 1 is a tiny
|
||||
// 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.
|
||||
// ===========================================================================
|
||||
|
||||
namespace host::compress {
|
||||
|
||||
using device::compress::CompressPlan;
|
||||
using device::compress::DecodePlan;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Decode plan builder.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
struct OnlineDecodePlanParams {
|
||||
DecodePlan* __restrict__ plan_d;
|
||||
const int64_t* __restrict__ seq_lens;
|
||||
const int64_t* __restrict__ req_pool_indices;
|
||||
const int32_t* __restrict__ req_to_token;
|
||||
const int64_t* __restrict__ full_to_swa; // (full_cache_size,) int64
|
||||
int64_t stride_r2t;
|
||||
int32_t swa_page_size;
|
||||
uint32_t batch_size;
|
||||
};
|
||||
|
||||
__global__ void plan_c128_online_decode_kernel(const OnlineDecodePlanParams params) {
|
||||
const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= params.batch_size) return;
|
||||
const auto seq_len = static_cast<uint32_t>(params.seq_lens[idx]);
|
||||
const auto rid = params.req_pool_indices[idx];
|
||||
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;
|
||||
params.plan_d[idx] = DecodePlan{
|
||||
.seq_len = seq_len,
|
||||
.write_loc = slot,
|
||||
.read_page_0 = slot,
|
||||
.read_page_1 = -1,
|
||||
};
|
||||
}
|
||||
|
||||
/// \brief Build the decode plan tensor. Caller (Python) pre-allocates
|
||||
/// `plan_d_dev` as a `(batch_size, 16)` device uint8 tensor; this routine
|
||||
/// only fills it. See `plan_online_prefill` for the rationale (avoid
|
||||
/// `ffi::empty` + dlpack roundtrip / PyTorch caching-allocator stream
|
||||
/// tracking issue that surfaces as IMA in unrelated downstream kernels).
|
||||
inline void plan_online_decode(
|
||||
const tvm::ffi::TensorView seq_lens,
|
||||
const tvm::ffi::TensorView req_pool_indices,
|
||||
const tvm::ffi::TensorView req_to_token,
|
||||
const tvm::ffi::TensorView full_to_swa,
|
||||
const tvm::ffi::TensorView plan_d_dev_,
|
||||
const int32_t swa_page_size) {
|
||||
auto B = SymbolicSize{"batch_size"};
|
||||
auto device_ = SymbolicDevice{};
|
||||
device_.set_options<kDLCUDA>();
|
||||
|
||||
auto seq_dtype = SymbolicDType{};
|
||||
TensorMatcher({B}) //
|
||||
.with_dtype<int64_t>(seq_dtype)
|
||||
.with_device(device_)
|
||||
.verify(seq_lens);
|
||||
TensorMatcher({B}) //
|
||||
.with_dtype<int64_t>()
|
||||
.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({B, sizeof(DecodePlan)}) //
|
||||
.with_dtype<uint8_t>()
|
||||
.with_device(device_)
|
||||
.verify(plan_d_dev_);
|
||||
RuntimeCheck(swa_page_size > 0);
|
||||
|
||||
const auto batch_size = static_cast<uint32_t>(B.unwrap());
|
||||
if (batch_size == 0) return;
|
||||
|
||||
const auto device = device_.unwrap();
|
||||
constexpr uint32_t kBlockSize = 256;
|
||||
const uint32_t num_blocks = host::div_ceil(batch_size, kBlockSize);
|
||||
const auto stride_r2t = req_to_token.stride(0);
|
||||
const auto params = OnlineDecodePlanParams{
|
||||
.plan_d = static_cast<DecodePlan*>(plan_d_dev_.data_ptr()),
|
||||
.seq_lens = static_cast<const int64_t*>(seq_lens.data_ptr()),
|
||||
.req_pool_indices = static_cast<const int64_t*>(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()),
|
||||
.stride_r2t = stride_r2t,
|
||||
.swa_page_size = swa_page_size,
|
||||
.batch_size = batch_size,
|
||||
};
|
||||
LaunchKernel(num_blocks, kBlockSize, device)(plan_c128_online_decode_kernel, params);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Prefill plan builder: host stage 0 + GPU stage 1.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
struct OnlinePrefillStage0Params {
|
||||
CompressPlan* __restrict__ plan_c;
|
||||
CompressPlan* __restrict__ plan_w;
|
||||
const int64_t* __restrict__ seq_lens;
|
||||
const int64_t* __restrict__ extend_lens;
|
||||
uint32_t batch_size;
|
||||
uint32_t num_q_tokens;
|
||||
};
|
||||
|
||||
inline std::tuple<uint32_t, uint32_t> _plan_prefill_partial(const OnlinePrefillStage0Params& p) {
|
||||
uint32_t counter = 0;
|
||||
uint32_t compress_count = 0;
|
||||
uint32_t write_count = 0;
|
||||
for (const auto i : irange(p.batch_size)) {
|
||||
const uint32_t seq_len = static_cast<uint32_t>(p.seq_lens[i]);
|
||||
const uint32_t extend_len = static_cast<uint32_t>(p.extend_lens[i]);
|
||||
RuntimeCheck(0 < extend_len && extend_len <= seq_len);
|
||||
const uint32_t prefix_len = seq_len - extend_len;
|
||||
const uint32_t end_pos = prefix_len + extend_len;
|
||||
|
||||
uint32_t pos = prefix_len;
|
||||
while (pos < end_pos) {
|
||||
const uint32_t chunk_start = (pos / 128u) * 128u;
|
||||
const uint32_t seg_end = std::min(end_pos, chunk_start + 128u); // exclusive
|
||||
const uint32_t seg_len = seg_end - pos;
|
||||
const uint32_t chunk_off = pos - chunk_start;
|
||||
const uint32_t last_pos = seg_end - 1;
|
||||
const uint32_t last_ragged = counter + (last_pos - prefix_len);
|
||||
RuntimeCheck(last_ragged < (1u << 16), "PlanC.ragged_id is uint16; ragged ", last_ragged, " overflows");
|
||||
RuntimeCheck(seg_len <= 128u);
|
||||
// Stash batch_id in `read_page_0` for stage 1 to translate. A
|
||||
// chunk-aligned segment never loads, so we still need stage 1 to fill
|
||||
// a slot in -- the kernel keys the load on `chunk_offset != 0`.
|
||||
const auto plan = CompressPlan{
|
||||
.seq_len = last_pos + 1u,
|
||||
.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
|
||||
};
|
||||
if (chunk_off + seg_len == 128u) {
|
||||
// close-chunk segment
|
||||
RuntimeCheck(compress_count < p.num_q_tokens);
|
||||
p.plan_c[compress_count++] = plan;
|
||||
} else {
|
||||
// trailing partial segment
|
||||
RuntimeCheck(write_count < p.num_q_tokens);
|
||||
p.plan_w[write_count++] = plan;
|
||||
}
|
||||
pos = seg_end;
|
||||
}
|
||||
counter += extend_len;
|
||||
}
|
||||
RuntimeCheck(counter == p.num_q_tokens, "input size ", counter, " != num_q_tokens ", p.num_q_tokens);
|
||||
return std::tuple<uint32_t, uint32_t>{compress_count, write_count};
|
||||
}
|
||||
|
||||
struct OnlinePrefillStage1Params {
|
||||
CompressPlan* __restrict__ plan_c;
|
||||
CompressPlan* __restrict__ plan_w;
|
||||
const int64_t* __restrict__ req_pool_indices; // (batch_size,)
|
||||
const int32_t* __restrict__ req_to_token; // (num_reqs, max_tokens)
|
||||
const int64_t* __restrict__ full_to_swa; // (full_cache_size,)
|
||||
int64_t stride_r2t;
|
||||
int32_t swa_page_size;
|
||||
uint32_t num_c;
|
||||
uint32_t num_w;
|
||||
};
|
||||
|
||||
__global__ void plan_c128_online_prefill_kernel(const OnlinePrefillStage1Params params) {
|
||||
const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const uint32_t total = params.num_c + params.num_w;
|
||||
if (idx >= total) return;
|
||||
|
||||
const bool is_compress = idx < params.num_c;
|
||||
CompressPlan* const plan_ptr = is_compress ? ¶ms.plan_c[idx] : ¶ms.plan_w[idx - params.num_c];
|
||||
auto plan = *plan_ptr;
|
||||
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;
|
||||
*plan_ptr = plan;
|
||||
}
|
||||
|
||||
using OnlinePrefillPlan = tvm::ffi::Tuple<uint32_t, uint32_t>;
|
||||
|
||||
inline OnlinePrefillPlan plan_online_prefill(
|
||||
const tvm::ffi::TensorView seq_lens,
|
||||
const tvm::ffi::TensorView extend_lens,
|
||||
const tvm::ffi::TensorView req_pool_indices,
|
||||
const tvm::ffi::TensorView req_to_token,
|
||||
const tvm::ffi::TensorView full_to_swa,
|
||||
const tvm::ffi::TensorView plan_c_pin,
|
||||
const tvm::ffi::TensorView plan_w_pin,
|
||||
const tvm::ffi::TensorView plan_c_dev_,
|
||||
const tvm::ffi::TensorView plan_w_dev_,
|
||||
const int32_t swa_page_size) {
|
||||
auto B = SymbolicSize{"batch_size"};
|
||||
auto N = SymbolicSize{"num_q_tokens"};
|
||||
auto cpu = SymbolicDevice{};
|
||||
auto device_ = SymbolicDevice{};
|
||||
cpu.set_options<kDLCPU, kDLCUDAHost>();
|
||||
device_.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({B}) //
|
||||
.with_dtype<int64_t>()
|
||||
.with_device(cpu)
|
||||
.verify(seq_lens)
|
||||
.verify(extend_lens);
|
||||
TensorMatcher({B}) //
|
||||
.with_dtype<int64_t>()
|
||||
.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({N, sizeof(CompressPlan)}) //
|
||||
.with_dtype<uint8_t>()
|
||||
.with_device(cpu)
|
||||
.verify(plan_c_pin)
|
||||
.verify(plan_w_pin);
|
||||
TensorMatcher({N, sizeof(CompressPlan)}) //
|
||||
.with_dtype<uint8_t>()
|
||||
.with_device(device_)
|
||||
.verify(plan_c_dev_)
|
||||
.verify(plan_w_dev_);
|
||||
|
||||
const auto stage0_params = OnlinePrefillStage0Params{
|
||||
.plan_c = static_cast<CompressPlan*>(plan_c_pin.data_ptr()),
|
||||
.plan_w = static_cast<CompressPlan*>(plan_w_pin.data_ptr()),
|
||||
.seq_lens = static_cast<const int64_t*>(seq_lens.data_ptr()),
|
||||
.extend_lens = static_cast<const int64_t*>(extend_lens.data_ptr()),
|
||||
.batch_size = static_cast<uint32_t>(B.unwrap()),
|
||||
.num_q_tokens = static_cast<uint32_t>(N.unwrap()),
|
||||
};
|
||||
|
||||
// Debug instrumentation: SGLANG_DEBUG_C128_ONLINE_GUARD=1 wraps stage 0
|
||||
// with redzone + post-write magic-check on the pin buffers, plus a strict
|
||||
// upper-bound check on `batch_size` and `num_q_tokens`. If stage 0 has a
|
||||
// CPU OOB this trips a clear panic at the offending byte instead of a
|
||||
// delayed CUDA IMA from corrupted heap memory.
|
||||
static const bool kGuard = []() {
|
||||
const char* v = std::getenv("SGLANG_DEBUG_C128_ONLINE_GUARD");
|
||||
return v != nullptr && v[0] == '1';
|
||||
}();
|
||||
if (kGuard) {
|
||||
RuntimeCheck(stage0_params.batch_size <= 65536u, "batch_size out of bound: ", stage0_params.batch_size);
|
||||
RuntimeCheck(stage0_params.num_q_tokens <= 65536u, "num_q_tokens out of bound: ", stage0_params.num_q_tokens);
|
||||
// Stamp the pin buffers with 0xAB so we can detect any byte still 0xAB
|
||||
// beyond what stage 0 should have written (= OOB never reached, that's fine)
|
||||
// or any byte BEYOND num_q_tokens*16 written to (= true OOB into
|
||||
// adjacent allocation).
|
||||
auto* pc = static_cast<uint8_t*>(plan_c_pin.data_ptr());
|
||||
auto* pw = static_cast<uint8_t*>(plan_w_pin.data_ptr());
|
||||
const auto bytes = static_cast<size_t>(N.unwrap()) * sizeof(CompressPlan);
|
||||
std::memset(pc, 0xAB, bytes);
|
||||
std::memset(pw, 0xAB, bytes);
|
||||
}
|
||||
|
||||
const auto [num_c, num_w] = _plan_prefill_partial(stage0_params);
|
||||
|
||||
if (kGuard) {
|
||||
// Verify stage 0 wrote ONLY to the [0, num_c*16) and [0, num_w*16) prefix.
|
||||
auto* pc = static_cast<const uint8_t*>(plan_c_pin.data_ptr());
|
||||
auto* pw = static_cast<const uint8_t*>(plan_w_pin.data_ptr());
|
||||
const auto end_c = static_cast<size_t>(num_c) * sizeof(CompressPlan);
|
||||
const auto end_w = static_cast<size_t>(num_w) * sizeof(CompressPlan);
|
||||
const auto pin_bytes = static_cast<size_t>(N.unwrap()) * sizeof(CompressPlan);
|
||||
for (size_t k = end_c; k < pin_bytes; ++k) {
|
||||
RuntimeCheck(
|
||||
pc[k] == 0xAB,
|
||||
"GUARD: plan_c_pin OOB write at byte ",
|
||||
k,
|
||||
" (num_c=",
|
||||
num_c,
|
||||
", num_q_tokens=",
|
||||
N.unwrap(),
|
||||
")");
|
||||
}
|
||||
for (size_t k = end_w; k < pin_bytes; ++k) {
|
||||
RuntimeCheck(
|
||||
pw[k] == 0xAB,
|
||||
"GUARD: plan_w_pin OOB write at byte ",
|
||||
k,
|
||||
" (num_w=",
|
||||
num_w,
|
||||
", num_q_tokens=",
|
||||
N.unwrap(),
|
||||
")");
|
||||
}
|
||||
}
|
||||
|
||||
const auto device = device_.unwrap();
|
||||
// Out-params pre-allocated by Python. Cast to typed pointers for use.
|
||||
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) {
|
||||
const auto stream = LaunchKernel::resolve_device(device);
|
||||
// SGLANG_DEBUG_C128_ONLINE_SYNC_H2D=1 forces a synchronous H2D copy.
|
||||
static const bool kSyncH2D = []() {
|
||||
const char* v = std::getenv("SGLANG_DEBUG_C128_ONLINE_SYNC_H2D");
|
||||
return v != nullptr && v[0] == '1';
|
||||
}();
|
||||
// SGLANG_DEBUG_C128_ONLINE_NO_H2D=1 skips the H2D copy entirely (debug only).
|
||||
static const bool kNoH2D = []() {
|
||||
const char* v = std::getenv("SGLANG_DEBUG_C128_ONLINE_NO_H2D");
|
||||
return v != nullptr && v[0] == '1';
|
||||
}();
|
||||
const auto copy_to_device = [stream](void* dst, void* src, int64_t count) {
|
||||
if (kNoH2D) return;
|
||||
const auto bytes = count * sizeof(CompressPlan);
|
||||
if (kSyncH2D) {
|
||||
RuntimeDeviceCheck(::cudaMemcpy(dst, src, bytes, ::cudaMemcpyHostToDevice));
|
||||
} else {
|
||||
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);
|
||||
|
||||
const auto stage1_params = OnlinePrefillStage1Params{
|
||||
.plan_c = plan_c_dev_ptr,
|
||||
.plan_w = plan_w_dev_ptr,
|
||||
.req_pool_indices = static_cast<const int64_t*>(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()),
|
||||
.stride_r2t = req_to_token.stride(0),
|
||||
.swa_page_size = swa_page_size,
|
||||
.num_c = num_c,
|
||||
.num_w = num_w,
|
||||
};
|
||||
constexpr uint32_t kBlockSize = 128;
|
||||
const auto num_blocks = host::div_ceil(total, kBlockSize);
|
||||
LaunchKernel(num_blocks, kBlockSize, device)(plan_c128_online_prefill_kernel, stage1_params);
|
||||
}
|
||||
return OnlinePrefillPlan{num_c, num_w};
|
||||
}
|
||||
|
||||
} // namespace host::compress
|
||||
|
||||
namespace {
|
||||
|
||||
[[maybe_unused]]
|
||||
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;
|
||||
|
||||
} // namespace
|
||||
@@ -1,3 +1,16 @@
|
||||
/**
|
||||
* \brief Here's some dimension info for the main buffer used in C128 prefill and decode.
|
||||
*
|
||||
* kv_buffer: [num_indices, 128, head_dim * 2]
|
||||
* - last dimension layout: | kv | score |
|
||||
* kv_input: [batch_size, head_dim * 2]
|
||||
* kv_output: [batch_size, head_dim]
|
||||
* score_bias (ape): [128, head_dim]
|
||||
* plan_c/plan_w: [variable length]
|
||||
*
|
||||
* For prefill, batch_size = num_q_tokens
|
||||
*/
|
||||
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
@@ -8,7 +21,7 @@
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
#include <sgl_kernel/warp.cuh>
|
||||
|
||||
#include <sgl_kernel/deepseek_v4/compress.cuh>
|
||||
#include <sgl_kernel/deepseek_v4/compress_v2.cuh>
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
@@ -18,8 +31,9 @@
|
||||
|
||||
namespace {
|
||||
|
||||
using Plan128 = device::compress::PrefillPlan;
|
||||
using IndiceT = int32_t;
|
||||
using PlanD = device::compress::DecodePlan;
|
||||
using PlanC = device::compress::CompressPlan;
|
||||
using PlanW = device::compress::WritePlan;
|
||||
|
||||
/// \brief Each thread will handle this many elements (split along head_dim)
|
||||
constexpr int32_t kTileElements = 2;
|
||||
@@ -27,59 +41,30 @@ constexpr int32_t kTileElements = 2;
|
||||
constexpr int32_t kElementsPerWarp = 8;
|
||||
constexpr uint32_t kNumWarps = 128 / kElementsPerWarp;
|
||||
constexpr uint32_t kBlockSize = device::kWarpThreads * kNumWarps;
|
||||
constexpr uint32_t kWriteBlockSize = 128; // one warp per write
|
||||
|
||||
/// \brief Need to reduce register usage to increase occupancy
|
||||
#define C128_KERNEL __global__ __launch_bounds__(kBlockSize, 2)
|
||||
#define WRITE_KERNEL __global__ __launch_bounds__(kWriteBlockSize, 16)
|
||||
|
||||
struct Compress128DecodeParams {
|
||||
/**
|
||||
* \brief Shape: `[num_indices, 128, head_dim * 2]` \n
|
||||
* last dimension layout:
|
||||
* | kv current | score current |
|
||||
*/
|
||||
void* __restrict__ kv_score_buffer;
|
||||
/** \brief Shape: `[batch_size, head_dim * 2]` */
|
||||
const void* __restrict__ kv_score_input;
|
||||
/** \brief Shape: `[batch_size, head_dim]` */
|
||||
void* __restrict__ kv_compressed_output;
|
||||
/** \brief Shape: `[128, head_dim]` (called `ape`) */
|
||||
void* __restrict__ kv_buffer;
|
||||
const void* __restrict__ kv_input;
|
||||
void* __restrict__ kv_output;
|
||||
const void* __restrict__ score_bias;
|
||||
/** \brief Shape: `[batch_size, ]`*/
|
||||
const IndiceT* __restrict__ indices;
|
||||
/** \brief Shape: `[batch_size, ]` */
|
||||
const IndiceT* __restrict__ seq_lens;
|
||||
/** \NOTE: `batch_size` <= `num_indices` */
|
||||
const PlanD* __restrict__ plan_d;
|
||||
uint32_t batch_size;
|
||||
};
|
||||
|
||||
struct Compress128PrefillParams {
|
||||
/**
|
||||
* \brief Shape: `[num_indices, 128, head_dim * 2]` \n
|
||||
* last dimension layout:
|
||||
* | kv current | score current |
|
||||
*/
|
||||
void* __restrict__ kv_score_buffer;
|
||||
/** \brief Shape: `[batch_size, head_dim * 2]` */
|
||||
const void* __restrict__ kv_score_input;
|
||||
/** \brief Shape: `[batch_size, head_dim]` */
|
||||
void* __restrict__ kv_compressed_output;
|
||||
/** \brief Shape: `[128, head_dim]` (called `ape`) */
|
||||
void* __restrict__ kv_buffer;
|
||||
const void* __restrict__ kv_input;
|
||||
void* __restrict__ kv_output;
|
||||
const void* __restrict__ score_bias;
|
||||
/** \brief Shape: `[batch_size, ]`*/
|
||||
const IndiceT* __restrict__ indices;
|
||||
/** \brief Shape: `[batch_size, ]`*/
|
||||
const int32_t* __restrict__ load_indices;
|
||||
/** \brief The following part is plan info. */
|
||||
|
||||
const Plan128* __restrict__ compress_plan;
|
||||
const Plan128* __restrict__ write_plan;
|
||||
|
||||
const PlanC* __restrict__ plan_c;
|
||||
const PlanW* __restrict__ plan_w;
|
||||
uint32_t num_compress;
|
||||
uint32_t num_write;
|
||||
|
||||
uint32_t num_q_tokens;
|
||||
uint32_t batch_size;
|
||||
uint32_t num_indices;
|
||||
};
|
||||
|
||||
struct Compress128SharedBuffer {
|
||||
@@ -93,46 +78,28 @@ struct Compress128SharedBuffer {
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
SGL_DEVICE void c128_write(
|
||||
T* kv_score_buf, //
|
||||
const T* kv_score_src,
|
||||
const int64_t head_dim,
|
||||
const int32_t write_pos,
|
||||
const uint32_t lane_id) {
|
||||
using namespace device;
|
||||
template <int64_t kHeadDim_>
|
||||
struct C128Trait {
|
||||
static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads; // 64
|
||||
static constexpr int64_t kHeadDim = kHeadDim_;
|
||||
static constexpr int64_t kScoreOffset = kHeadDim;
|
||||
static constexpr int64_t kElementSize = kHeadDim * 2;
|
||||
static constexpr int64_t kPageElementSize = 128 * kElementSize; // page size = 128
|
||||
static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
|
||||
static_assert(kHeadDim % kTileDim == 0);
|
||||
};
|
||||
|
||||
using Storage = AlignedVector<T, kTileElements>;
|
||||
const auto element_size = head_dim * 2;
|
||||
const auto gmem = tile::Memory<Storage>{lane_id, kWarpThreads};
|
||||
kv_score_buf += write_pos * element_size;
|
||||
|
||||
/// NOTE: Layout | [0] = kv | [1] = score |
|
||||
Storage kv_score[2];
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < 2; ++i) {
|
||||
kv_score[i] = gmem.load(kv_score_src + head_dim * i);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < 2; ++i) {
|
||||
gmem.store(kv_score_buf + head_dim * i, kv_score[i]);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename InFloat, typename OutFloat>
|
||||
template <typename Trait, bool kUsePDL, typename InFloat, typename OutFloat>
|
||||
SGL_DEVICE void c128_forward(
|
||||
const InFloat* kv_score_buf,
|
||||
const InFloat* kv_score_src,
|
||||
const InFloat* kv_buf, // [128n, 128n + 127]
|
||||
const InFloat* kv_src, // ragged pointer at position = 128n + 127
|
||||
OutFloat* kv_out,
|
||||
const InFloat* score_bias,
|
||||
const int64_t head_dim,
|
||||
const int32_t window_len,
|
||||
const uint32_t warp_id,
|
||||
const uint32_t lane_id) {
|
||||
const int32_t buffer_len) {
|
||||
using namespace device;
|
||||
|
||||
const auto element_size = head_dim * 2;
|
||||
const auto score_offset = head_dim;
|
||||
const auto warp_id = threadIdx.x / kWarpThreads;
|
||||
const auto lane_id = threadIdx.x % kWarpThreads;
|
||||
|
||||
/// NOTE: part 1: load kv + score
|
||||
using StorageIn = AlignedVector<InFloat, kTileElements>;
|
||||
@@ -145,23 +112,18 @@ SGL_DEVICE void c128_forward(
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < 8; ++i) {
|
||||
const int32_t j = i + warp_offset;
|
||||
bias[i] = gmem_in.load(score_bias + j * head_dim);
|
||||
bias[i] = gmem_in.load(score_bias + j * Trait::kHeadDim);
|
||||
}
|
||||
|
||||
const auto kv_start = kv_src - 127 * Trait::kElementSize; // point to start
|
||||
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < kElementsPerWarp; ++i) {
|
||||
const int32_t j = i + warp_offset;
|
||||
const InFloat* src;
|
||||
__builtin_assume(j < 128);
|
||||
if (j < window_len) {
|
||||
src = kv_score_buf + j * element_size;
|
||||
} else {
|
||||
/// NOTE: k in [-127, 0]. We'll load from the ragged `kv_score_src`
|
||||
const int32_t k = j - 127;
|
||||
src = kv_score_src + k * element_size;
|
||||
}
|
||||
kv[i] = gmem_in.load(src);
|
||||
score[i] = gmem_in.load(src + score_offset);
|
||||
const auto src = j < buffer_len ? kv_buf : kv_start;
|
||||
kv[i] = gmem_in.load(src + j * Trait::kElementSize);
|
||||
score[i] = gmem_in.load(src + j * Trait::kElementSize + Trait::kScoreOffset);
|
||||
}
|
||||
|
||||
/// NOTE: part 2: safe online softmax + weighted sum
|
||||
@@ -174,28 +136,32 @@ SGL_DEVICE void c128_forward(
|
||||
TmpStorage tmp_exp_sum;
|
||||
TmpStorage tmp_product;
|
||||
|
||||
float score_fp32[kTileElements][kElementsPerWarp];
|
||||
|
||||
// convert to fp32 and apply bias first
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < kTileElements; ++i) {
|
||||
float score_fp32[kElementsPerWarp];
|
||||
for (int32_t j = 0; j < kElementsPerWarp; ++j) {
|
||||
score_fp32[i][j] = cast<float>(score[j][i]) + cast<float>(bias[j][i]);
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int32_t j = 0; j < kElementsPerWarp; ++j) {
|
||||
score_fp32[j] = cast<float>(score[j][i]) + cast<float>(bias[j][i]);
|
||||
}
|
||||
|
||||
float max_value = score_fp32[0];
|
||||
for (int32_t i = 0; i < kTileElements; ++i) {
|
||||
const auto& score = score_fp32[i];
|
||||
float max_value = score[0];
|
||||
float sum_exp_value = 0.0f;
|
||||
|
||||
#pragma unroll
|
||||
for (int32_t j = 1; j < kElementsPerWarp; ++j) {
|
||||
const auto fp32_score = score_fp32[j];
|
||||
const auto fp32_score = score[j];
|
||||
max_value = fmaxf(max_value, fp32_score);
|
||||
}
|
||||
|
||||
float sum_product = 0.0f;
|
||||
#pragma unroll
|
||||
for (int32_t j = 0; j < 8; ++j) {
|
||||
const auto fp32_score = score_fp32[j];
|
||||
const auto fp32_score = score[j];
|
||||
const auto exp_score = expf(fp32_score - max_value);
|
||||
sum_product += cast<float>(kv[j][i]) * exp_score;
|
||||
sum_exp_value += exp_score;
|
||||
@@ -219,6 +185,8 @@ SGL_DEVICE void c128_forward(
|
||||
constexpr uint32_t kReductionCount = kTileElements * kWarpThreads * kNumWarps;
|
||||
constexpr uint32_t kIteration = kReductionCount / kBlockSize;
|
||||
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kIteration; ++i) {
|
||||
/// NOTE: Range `[0, kTileElements * kWarpThreads * kNumWarps)`
|
||||
@@ -247,293 +215,230 @@ SGL_DEVICE void c128_forward(
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Trait, typename InFloat>
|
||||
SGL_DEVICE void c128_write_decode(InFloat* kv_buf, const InFloat* kv_src) {
|
||||
using namespace device;
|
||||
|
||||
using Storage = AlignedVector<InFloat, kTileElements>;
|
||||
const auto gmem = tile::Memory<Storage>::warp();
|
||||
|
||||
Storage data[2];
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < 2; ++i) {
|
||||
data[i] = gmem.load(kv_src + Trait::kHeadDim * i);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < 2; ++i) {
|
||||
gmem.store(kv_buf + Trait::kHeadDim * i, data[i]);
|
||||
}
|
||||
}
|
||||
|
||||
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
|
||||
C128_KERNEL void flash_c128_decode(const __grid_constant__ Compress128DecodeParams params) {
|
||||
using namespace device;
|
||||
using Trait = C128Trait<kHeadDim>;
|
||||
|
||||
constexpr int64_t kTileDim = kTileElements * kWarpThreads; // 64
|
||||
constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
|
||||
constexpr int64_t kElementSize = kHeadDim * 2;
|
||||
static_assert(kHeadDim % kTileDim == 0, "Head dim must be multiple of tile dim");
|
||||
|
||||
const auto& [
|
||||
_kv_score_buffer, _kv_score_input, _kv_compressed_output, _score_bias, // kv score
|
||||
indices, seq_lens, batch_size // decode info
|
||||
] = params;
|
||||
const uint32_t warp_id = threadIdx.x / kWarpThreads;
|
||||
const uint32_t lane_id = threadIdx.x % kWarpThreads;
|
||||
const uint32_t global_bid = blockIdx.x / Trait::kNumSplit; // batch id
|
||||
const uint32_t global_sid = blockIdx.x % Trait::kNumSplit; // split id
|
||||
const int64_t split_offset = global_sid * Trait::kTileDim;
|
||||
if (global_bid >= params.batch_size) return;
|
||||
|
||||
const uint32_t global_bid = blockIdx.x / kNumSplit; // batch id
|
||||
const uint32_t global_sid = blockIdx.x % kNumSplit; // split id
|
||||
if (global_bid >= batch_size) return;
|
||||
const auto plan = params.plan_d[global_bid];
|
||||
const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
|
||||
const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
|
||||
const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
|
||||
const auto score_bias = static_cast<const InFloat*>(params.score_bias) + split_offset;
|
||||
|
||||
const int32_t index = indices[global_bid];
|
||||
const int32_t seq_len = seq_lens[global_bid];
|
||||
const int64_t split_offset = global_sid * kTileDim;
|
||||
|
||||
// kv score
|
||||
const auto kv_score_buffer = static_cast<InFloat*>(_kv_score_buffer);
|
||||
const auto kv_buf = kv_score_buffer + index * (kElementSize * 128) + split_offset;
|
||||
|
||||
// kv input
|
||||
const auto kv_score_input = static_cast<const InFloat*>(_kv_score_input);
|
||||
const auto kv_src = kv_score_input + global_bid * kElementSize + split_offset;
|
||||
|
||||
// kv output
|
||||
const auto kv_compressed_output = static_cast<OutFloat*>(_kv_compressed_output);
|
||||
const auto kv_out = kv_compressed_output + global_bid * kHeadDim + split_offset;
|
||||
|
||||
// score bias (ape)
|
||||
const auto score_bias = static_cast<const InFloat*>(_score_bias) + split_offset;
|
||||
const auto kv_src = kv_input + global_bid * Trait::kElementSize;
|
||||
const auto kv_out = kv_output + global_bid * Trait::kHeadDim;
|
||||
const auto kv_buf = kv_buffer + plan.read_page_1 * Trait::kPageElementSize;
|
||||
const auto kv_dst = kv_buffer + plan.write_loc * Trait::kElementSize;
|
||||
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
|
||||
/// NOTE: the write must be visible to the subsequent c128_forward,
|
||||
/// so only the last warp can write to HBM
|
||||
/// In addition, `position` = `seq_len - 1`. To avoid underflow, we use `seq_len + 127`
|
||||
// the write warp must match the load warp in the following `c128_forward`
|
||||
if (warp_id == kNumWarps - 1) {
|
||||
c128_write(kv_buf, kv_src, kHeadDim, /*write_pos=*/(seq_len + 127) % 128, lane_id);
|
||||
c128_write_decode<Trait>(kv_dst, kv_src);
|
||||
}
|
||||
if (seq_len % 128 == 0) {
|
||||
c128_forward(kv_buf, kv_src, kv_out, score_bias, kHeadDim, /*window_len=*/128, warp_id, lane_id);
|
||||
if (plan.write_loc % 128 == 127) {
|
||||
c128_forward<Trait, kUsePDL>(kv_buf, kv_src, kv_out, score_bias, 128);
|
||||
}
|
||||
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
}
|
||||
|
||||
// compress kernel
|
||||
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kWrite, bool kUsePDL>
|
||||
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
|
||||
C128_KERNEL void flash_c128_prefill(const __grid_constant__ Compress128PrefillParams params) {
|
||||
using namespace device;
|
||||
using Trait = C128Trait<kHeadDim>;
|
||||
|
||||
constexpr int64_t kTileDim = kTileElements * kWarpThreads; // 64
|
||||
constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
|
||||
constexpr int64_t kElementSize = kHeadDim * 2;
|
||||
static_assert(kHeadDim % kTileDim == 0, "Head dim must be multiple of tile dim");
|
||||
const uint32_t global_pid = blockIdx.x / Trait::kNumSplit; // plan id
|
||||
const uint32_t global_sid = blockIdx.x % Trait::kNumSplit; // split id
|
||||
const int64_t split_offset = global_sid * Trait::kTileDim;
|
||||
if (global_pid >= params.num_compress) return;
|
||||
|
||||
const auto& [
|
||||
_kv_score_buffer, _kv_score_input, _kv_compressed_output, _score_bias, // kv score
|
||||
indices, load_indices, compress_plan, write_plan, num_compress, num_write, // prefill plan
|
||||
_num_q_tokens, _batch_size, _num_indices
|
||||
] = params;
|
||||
const uint32_t warp_id = threadIdx.x / kWarpThreads;
|
||||
const uint32_t lane_id = threadIdx.x % kWarpThreads;
|
||||
const auto plan = params.plan_c[global_pid];
|
||||
const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
|
||||
const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
|
||||
const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
|
||||
const auto score_bias = static_cast<const InFloat*>(params.score_bias) + split_offset;
|
||||
if (plan.is_invalid()) return;
|
||||
|
||||
uint32_t global_id;
|
||||
if constexpr (kWrite) {
|
||||
// for write kernel, we use global warp_id to dispatch work
|
||||
global_id = (blockIdx.x * blockDim.x + threadIdx.x) / kWarpThreads;
|
||||
} else {
|
||||
// for compress kernel, we use block id to dispatch work
|
||||
global_id = blockIdx.x; // block id
|
||||
}
|
||||
const uint32_t global_pid = global_id / kNumSplit; // plan id
|
||||
const uint32_t global_sid = global_id % kNumSplit; // split id
|
||||
const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize;
|
||||
// Compact output: one row per compress plan, indexed by `global_pid`.
|
||||
const auto kv_out = kv_output + global_pid * Trait::kHeadDim;
|
||||
const auto kv_buf = kv_buffer + plan.read_page_1 * Trait::kPageElementSize;
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
c128_forward<Trait, kUsePDL>(kv_buf, kv_src, kv_out, score_bias, plan.buffer_len);
|
||||
}
|
||||
|
||||
/// NOTE: compiler can optimize this if-else at compile time
|
||||
const auto num_plans = kWrite ? num_write : num_compress;
|
||||
const auto plan_ptr = kWrite ? write_plan : compress_plan;
|
||||
if (global_pid >= num_plans) return;
|
||||
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
|
||||
WRITE_KERNEL void write_c128_prefill(const __grid_constant__ Compress128PrefillParams params) {
|
||||
using namespace device;
|
||||
using Trait = C128Trait<kHeadDim>;
|
||||
using StorageIn = AlignedVector<InFloat, kTileElements>;
|
||||
|
||||
const auto& [ragged_id, global_bid, position, window_len] = plan_ptr[global_pid];
|
||||
const auto indices_ptr = kWrite ? indices : load_indices;
|
||||
const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const uint32_t global_wid = global_tid / kWarpThreads; // warp id
|
||||
const uint32_t global_pid = global_wid / Trait::kNumSplit; // plan id
|
||||
const uint32_t global_sid = global_wid % Trait::kNumSplit; // split id
|
||||
// split the contiguous `kHeadDim * 2` into `kNumSplit` tiles
|
||||
// each warp handles 1 contiguous tile (in contrast, decode handle the strided head_dim)
|
||||
const int64_t split_offset = global_sid * (Trait::kTileDim * 2);
|
||||
if (global_pid >= params.num_write) return;
|
||||
|
||||
const int64_t split_offset = global_sid * kTileDim;
|
||||
const auto plan = params.plan_w[global_pid];
|
||||
const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
|
||||
const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
|
||||
if (plan.is_invalid()) return;
|
||||
|
||||
// kv input
|
||||
const auto kv_score_input = static_cast<const InFloat*>(_kv_score_input);
|
||||
const auto kv_src = kv_score_input + ragged_id * kElementSize + split_offset;
|
||||
|
||||
// kv output
|
||||
const auto kv_compressed_output = static_cast<OutFloat*>(_kv_compressed_output);
|
||||
const auto kv_out = kv_compressed_output + ragged_id * kHeadDim + split_offset;
|
||||
|
||||
// score bias (ape)
|
||||
const auto score_bias = static_cast<const InFloat*>(_score_bias) + split_offset;
|
||||
|
||||
if (ragged_id == 0xFFFFFFFF) [[unlikely]]
|
||||
return;
|
||||
|
||||
if (ragged_id >= _num_q_tokens) [[unlikely]]
|
||||
return;
|
||||
if (global_bid >= _batch_size) [[unlikely]]
|
||||
return;
|
||||
|
||||
const int32_t index = indices_ptr[global_bid];
|
||||
|
||||
if (index < 0 || static_cast<uint32_t>(index) >= _num_indices) [[unlikely]]
|
||||
return;
|
||||
|
||||
// kv score
|
||||
const auto kv_score_buffer = static_cast<InFloat*>(_kv_score_buffer);
|
||||
const auto kv_buf = kv_score_buffer + index * (kElementSize * 128) + split_offset;
|
||||
// each warp will handle a contiguous region
|
||||
const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize;
|
||||
const auto kv_buf = kv_buffer + plan.write_loc * Trait::kElementSize;
|
||||
const auto gmem = tile::Memory<StorageIn>::warp();
|
||||
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
|
||||
// only responsible for the compress part
|
||||
if constexpr (kWrite) {
|
||||
c128_write(kv_buf, kv_src, kHeadDim, /*write_pos=*/position % 128, lane_id);
|
||||
} else {
|
||||
c128_forward(kv_buf, kv_src, kv_out, score_bias, kHeadDim, window_len, warp_id, lane_id);
|
||||
StorageIn data[2];
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < 2; ++i) {
|
||||
data[i] = gmem.load(kv_src, i);
|
||||
}
|
||||
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < 2; ++i) {
|
||||
gmem.store(kv_buf, data[i], i);
|
||||
}
|
||||
}
|
||||
|
||||
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
|
||||
struct FlashCompress128Kernel {
|
||||
static constexpr auto decode_kernel = flash_c128_decode<kHeadDim, InFloat, OutFloat, kUsePDL>;
|
||||
template <bool kWrite>
|
||||
static constexpr auto prefill_kernel = flash_c128_prefill<kHeadDim, InFloat, OutFloat, kWrite, kUsePDL>;
|
||||
static constexpr auto prefill_c_kernel = prefill_kernel</*kWrite=*/false>;
|
||||
static constexpr auto prefill_w_kernel = prefill_kernel</*kWrite=*/true>;
|
||||
static constexpr auto prefill_c_kernel = flash_c128_prefill<kHeadDim, InFloat, OutFloat, kUsePDL>;
|
||||
static constexpr auto prefill_w_kernel = write_c128_prefill<kHeadDim, InFloat, OutFloat, kUsePDL>;
|
||||
static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads; // 64
|
||||
static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
|
||||
static constexpr uint32_t kWriteBlockSize = 128;
|
||||
static constexpr uint32_t kWarpsPerWriteBlock = kWriteBlockSize / device::kWarpThreads;
|
||||
using Trait = C128Trait<kHeadDim>;
|
||||
|
||||
static void run_decode(
|
||||
const tvm::ffi::TensorView kv_score_buffer,
|
||||
const tvm::ffi::TensorView kv_score_input,
|
||||
const tvm::ffi::TensorView kv_compressed_output,
|
||||
const tvm::ffi::TensorView kv_buffer,
|
||||
const tvm::ffi::TensorView kv_input,
|
||||
const tvm::ffi::TensorView kv_output,
|
||||
const tvm::ffi::TensorView ape,
|
||||
const tvm::ffi::TensorView indices,
|
||||
const tvm::ffi::TensorView seq_lens,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> /* UNUSED */) {
|
||||
const tvm::ffi::TensorView plan_d_) {
|
||||
using namespace host;
|
||||
|
||||
// this should not happen in practice
|
||||
auto B = SymbolicSize{"batch_size"};
|
||||
auto device = SymbolicDevice{};
|
||||
device.set_options<kDLCUDA>();
|
||||
auto N = SymbolicSize{"batch_size"};
|
||||
auto device_ = SymbolicDevice{};
|
||||
device_.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({-1, 128, kHeadDim * 2}) // kv score
|
||||
TensorMatcher({-1, 128, Trait::kElementSize}) // kv score
|
||||
.with_dtype<InFloat>()
|
||||
.with_device(device)
|
||||
.verify(kv_score_buffer);
|
||||
TensorMatcher({B, kHeadDim * 2}) // kv score input
|
||||
.with_device(device_)
|
||||
.verify(kv_buffer);
|
||||
TensorMatcher({N, Trait::kElementSize}) // kv score input
|
||||
.with_dtype<InFloat>()
|
||||
.with_device(device)
|
||||
.verify(kv_score_input);
|
||||
TensorMatcher({B, kHeadDim}) // kv compressed output
|
||||
.with_device(device_)
|
||||
.verify(kv_input);
|
||||
TensorMatcher({N, kHeadDim}) // kv compressed output
|
||||
.with_dtype<OutFloat>()
|
||||
.with_device(device)
|
||||
.verify(kv_compressed_output);
|
||||
.with_device(device_)
|
||||
.verify(kv_output);
|
||||
TensorMatcher({128, kHeadDim}) // ape
|
||||
.with_dtype<InFloat>()
|
||||
.with_device(device)
|
||||
.with_device(device_)
|
||||
.verify(ape);
|
||||
TensorMatcher({B}) // indices
|
||||
.with_dtype<IndiceT>()
|
||||
.with_device(device)
|
||||
.verify(indices);
|
||||
TensorMatcher({B}) // seq lens
|
||||
.with_dtype<IndiceT>()
|
||||
.with_device(device)
|
||||
.verify(seq_lens);
|
||||
|
||||
const auto batch_size = static_cast<uint32_t>(B.unwrap());
|
||||
const auto plan_d = compress::verify_plan_d(plan_d_, N, device_);
|
||||
const auto batch_size = static_cast<uint32_t>(N.unwrap());
|
||||
const auto params = Compress128DecodeParams{
|
||||
.kv_score_buffer = kv_score_buffer.data_ptr(),
|
||||
.kv_score_input = kv_score_input.data_ptr(),
|
||||
.kv_compressed_output = kv_compressed_output.data_ptr(),
|
||||
.kv_buffer = kv_buffer.data_ptr(),
|
||||
.kv_input = kv_input.data_ptr(),
|
||||
.kv_output = kv_output.data_ptr(),
|
||||
.score_bias = ape.data_ptr(),
|
||||
.indices = static_cast<const IndiceT*>(indices.data_ptr()),
|
||||
.seq_lens = static_cast<const IndiceT*>(seq_lens.data_ptr()),
|
||||
.plan_d = plan_d,
|
||||
.batch_size = batch_size,
|
||||
};
|
||||
|
||||
const uint32_t num_blocks = batch_size * kNumSplit;
|
||||
LaunchKernel(num_blocks, kBlockSize, device.unwrap()) //
|
||||
LaunchKernel(num_blocks, kBlockSize, device_.unwrap()) //
|
||||
.enable_pdl(kUsePDL)(decode_kernel, params);
|
||||
}
|
||||
|
||||
static void run_prefill(
|
||||
const tvm::ffi::TensorView kv_score_buffer,
|
||||
const tvm::ffi::TensorView kv_score_input,
|
||||
const tvm::ffi::TensorView kv_compressed_output,
|
||||
const tvm::ffi::TensorView kv_buffer,
|
||||
const tvm::ffi::TensorView kv_input,
|
||||
const tvm::ffi::TensorView kv_output,
|
||||
const tvm::ffi::TensorView ape,
|
||||
const tvm::ffi::TensorView indices,
|
||||
const tvm::ffi::TensorView compress_plan,
|
||||
const tvm::ffi::TensorView write_plan,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> extra) {
|
||||
const tvm::ffi::TensorView plan_c_,
|
||||
const tvm::ffi::TensorView plan_w_) {
|
||||
using namespace host;
|
||||
|
||||
auto B = SymbolicSize{"batch_size"};
|
||||
auto N = SymbolicSize{"num_q_tokens"};
|
||||
auto X = SymbolicSize{"compress_tokens"};
|
||||
auto Y = SymbolicSize{"write_tokens"};
|
||||
auto K = SymbolicSize{"num_indices"};
|
||||
auto C = SymbolicSize{"num_c_plans"};
|
||||
auto W = SymbolicSize{"num_w_plans"};
|
||||
auto device_ = SymbolicDevice{};
|
||||
device_.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({K, 128, kHeadDim * 2}) // kv score
|
||||
TensorMatcher({-1, 128, Trait::kElementSize}) // kv score
|
||||
.with_dtype<InFloat>()
|
||||
.with_device(device_)
|
||||
.verify(kv_score_buffer);
|
||||
TensorMatcher({N, kHeadDim * 2}) // kv score input
|
||||
.verify(kv_buffer);
|
||||
TensorMatcher({N, Trait::kElementSize}) // kv score input (ragged)
|
||||
.with_dtype<InFloat>()
|
||||
.with_device(device_)
|
||||
.verify(kv_score_input);
|
||||
TensorMatcher({N, kHeadDim}) // kv compressed output
|
||||
.verify(kv_input);
|
||||
TensorMatcher({C, kHeadDim}) // kv compressed output (compact)
|
||||
.with_dtype<OutFloat>()
|
||||
.with_device(device_)
|
||||
.verify(kv_compressed_output);
|
||||
.verify(kv_output);
|
||||
TensorMatcher({128, kHeadDim}) // ape
|
||||
.with_dtype<InFloat>()
|
||||
.with_device(device_)
|
||||
.verify(ape);
|
||||
TensorMatcher({B}) // indices
|
||||
.with_dtype<IndiceT>()
|
||||
.with_device(device_)
|
||||
.verify(indices);
|
||||
TensorMatcher({X, compress::kPrefillPlanDim}) // compress plan
|
||||
.with_dtype<compress::PrefillPlanTensorDtype>()
|
||||
.with_device(device_)
|
||||
.verify(compress_plan);
|
||||
TensorMatcher({Y, compress::kPrefillPlanDim}) // write plan
|
||||
.with_dtype<compress::PrefillPlanTensorDtype>()
|
||||
.with_device(device_)
|
||||
.verify(write_plan);
|
||||
|
||||
// might be needed for prefill write
|
||||
const auto load_indices = extra.value_or(indices);
|
||||
TensorMatcher({B}) // [read_positions]
|
||||
.with_dtype<IndiceT>()
|
||||
.with_device(device_)
|
||||
.verify(load_indices);
|
||||
|
||||
const auto plan_c = compress::verify_plan_c(plan_c_, C, device_);
|
||||
const auto plan_w = compress::verify_plan_w(plan_w_, W, device_);
|
||||
const auto device = device_.unwrap();
|
||||
const auto batch_size = static_cast<uint32_t>(B.unwrap());
|
||||
const auto num_q_tokens = static_cast<uint32_t>(N.unwrap());
|
||||
const auto num_c = static_cast<uint32_t>(X.unwrap());
|
||||
const auto num_w = static_cast<uint32_t>(Y.unwrap());
|
||||
const auto num_indices = static_cast<uint32_t>(K.unwrap());
|
||||
const auto num_c = static_cast<uint32_t>(C.unwrap());
|
||||
const auto num_w = static_cast<uint32_t>(W.unwrap());
|
||||
const auto params = Compress128PrefillParams{
|
||||
.kv_score_buffer = kv_score_buffer.data_ptr(),
|
||||
.kv_score_input = kv_score_input.data_ptr(),
|
||||
.kv_compressed_output = kv_compressed_output.data_ptr(),
|
||||
.kv_buffer = kv_buffer.data_ptr(),
|
||||
.kv_input = kv_input.data_ptr(),
|
||||
.kv_output = kv_output.data_ptr(),
|
||||
.score_bias = ape.data_ptr(),
|
||||
.indices = static_cast<const IndiceT*>(indices.data_ptr()),
|
||||
.load_indices = static_cast<const IndiceT*>(load_indices.data_ptr()),
|
||||
.compress_plan = static_cast<const Plan128*>(compress_plan.data_ptr()),
|
||||
.write_plan = static_cast<const Plan128*>(write_plan.data_ptr()),
|
||||
.plan_c = plan_c,
|
||||
.plan_w = plan_w,
|
||||
.num_compress = num_c,
|
||||
.num_write = num_w,
|
||||
.num_q_tokens = num_q_tokens,
|
||||
.batch_size = batch_size,
|
||||
.num_indices = num_indices,
|
||||
};
|
||||
RuntimeCheck(num_q_tokens >= batch_size, "num_q_tokens must be >= batch_size");
|
||||
RuntimeCheck(num_q_tokens >= std::max(num_c, num_w), "invalid prefill plan");
|
||||
|
||||
constexpr auto kBlockSize_C = kBlockSize;
|
||||
constexpr auto kBlockSize_W = kWriteBlockSize;
|
||||
RuntimeCheck(num_q_tokens >= num_w, "invalid prefill plan: num_q < num_w");
|
||||
if (const auto num_c_blocks = num_c * kNumSplit) {
|
||||
constexpr auto kBlockSize_C = kBlockSize;
|
||||
LaunchKernel(num_c_blocks, kBlockSize_C, device) //
|
||||
.enable_pdl(kUsePDL)(prefill_c_kernel, params);
|
||||
}
|
||||
constexpr uint32_t kWarpsPerWriteBlock = kWriteBlockSize / device::kWarpThreads;
|
||||
if (const auto num_w_blocks = div_ceil(num_w * kNumSplit, kWarpsPerWriteBlock)) {
|
||||
constexpr auto kBlockSize_W = kWriteBlockSize;
|
||||
LaunchKernel(num_w_blocks, kBlockSize_W, device) //
|
||||
.enable_pdl(kUsePDL)(prefill_w_kernel, params);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,405 @@
|
||||
/**
|
||||
* \brief Here's some dimension info for the main buffer used in C4 prefill and decode.
|
||||
*
|
||||
* kv_buffer: [num_indices, 8, head_dim * 4]
|
||||
* - last dimension layout: | kv overlap | kv | score overlap | score |
|
||||
* kv_input: [batch_size, head_dim * 4]
|
||||
* kv_output: [batch_size, head_dim]
|
||||
* score_bias (ape): [8, head_dim]
|
||||
* plan_c/plan_w: [variable length]
|
||||
*
|
||||
* For prefill, batch_size = num_q_tokens
|
||||
*/
|
||||
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/tile.cuh>
|
||||
#include <sgl_kernel/type.cuh>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
#include <sgl_kernel/warp.cuh>
|
||||
|
||||
#include <sgl_kernel/deepseek_v4/compress_v2.cuh>
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
#include <tvm/ffi/object.h>
|
||||
|
||||
#include <cfloat>
|
||||
#include <cstdint>
|
||||
|
||||
namespace {
|
||||
|
||||
using PlanD = device::compress::DecodePlan;
|
||||
using PlanC = device::compress::CompressPlan;
|
||||
using PlanW = device::compress::WritePlan;
|
||||
|
||||
/// \brief Each thread will handle this many elements (split along head_dim)
|
||||
constexpr int32_t kTileElements = 4;
|
||||
|
||||
/// \brief Need to improve register usage to reduce latency
|
||||
#define C4_KERNEL __global__ __launch_bounds__(128, 4)
|
||||
#define WRITE_KERNEL __global__ __launch_bounds__(128, 16)
|
||||
|
||||
struct Compress4DecodeParams {
|
||||
void* __restrict__ kv_buffer;
|
||||
const void* __restrict__ kv_input;
|
||||
void* __restrict__ kv_output;
|
||||
const void* __restrict__ score_bias;
|
||||
const PlanD* __restrict__ plan_d;
|
||||
uint32_t batch_size;
|
||||
};
|
||||
|
||||
struct Compress4PrefillParams {
|
||||
void* __restrict__ kv_buffer;
|
||||
const void* __restrict__ kv_input;
|
||||
void* __restrict__ kv_output;
|
||||
const void* __restrict__ score_bias;
|
||||
const PlanC* __restrict__ plan_c;
|
||||
const PlanW* __restrict__ plan_w;
|
||||
uint32_t num_compress;
|
||||
uint32_t num_write;
|
||||
};
|
||||
|
||||
template <int64_t kHeadDim_>
|
||||
struct C4Trait {
|
||||
static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads; // 128
|
||||
static constexpr int64_t kHeadDim = kHeadDim_;
|
||||
static constexpr int64_t kOverlapOffset = kHeadDim;
|
||||
static constexpr int64_t kScoreOffset = kHeadDim * 2;
|
||||
static constexpr int64_t kElementSize = kHeadDim * 4;
|
||||
static constexpr int64_t kPageElementSize = 4 * kElementSize; // page size = 4
|
||||
static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
|
||||
static_assert(kHeadDim % kTileDim == 0);
|
||||
};
|
||||
|
||||
template <typename Trait, bool kUsePDL, typename InFloat, typename OutFloat>
|
||||
SGL_DEVICE void c4_forward(
|
||||
const InFloat* kv_buf_0, // overlap [4n - 4, 4n - 1]
|
||||
const InFloat* kv_buf_1, // normal [4n + 0, 4n + 3]
|
||||
const InFloat* kv_src, // ragged pointer at position = 4n + 3
|
||||
OutFloat* kv_out,
|
||||
const InFloat* score_bias,
|
||||
const bool should_overlap,
|
||||
const int32_t buffer_len) {
|
||||
using namespace device;
|
||||
|
||||
/// NOTE: part 1: load kv + score
|
||||
using StorageIn = AlignedVector<InFloat, kTileElements>;
|
||||
/// NOTE: load one tile_dim (< head_dim) at at time
|
||||
const auto gmem_in = tile::Memory<StorageIn>::warp();
|
||||
StorageIn kv[8];
|
||||
StorageIn score[8];
|
||||
StorageIn bias[8];
|
||||
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < 8; ++i) {
|
||||
bias[i] = gmem_in.load(score_bias + i * Trait::kHeadDim);
|
||||
}
|
||||
|
||||
if (should_overlap) {
|
||||
const auto kv_start = kv_src - 7 * Trait::kElementSize; // point to start
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < 4; ++i) {
|
||||
const auto src = i < buffer_len ? kv_buf_0 : kv_start;
|
||||
const auto base = src + i * Trait::kElementSize;
|
||||
kv[i] = gmem_in.load(base);
|
||||
score[i] = gmem_in.load(base + Trait::kScoreOffset);
|
||||
}
|
||||
} else {
|
||||
[[unlikely]];
|
||||
constexpr float kFloatNegInf = -FLT_MAX;
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < 4; ++i) {
|
||||
kv[i].fill(cast<InFloat>(0.0f));
|
||||
score[i].fill(cast<InFloat>(kFloatNegInf));
|
||||
}
|
||||
}
|
||||
|
||||
const auto kv_start = kv_src - 3 * Trait::kElementSize; // point to start
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < 4; ++i) {
|
||||
const auto src = i + 4 < buffer_len ? kv_buf_1 : kv_start;
|
||||
const auto base = src + i * Trait::kElementSize + Trait::kOverlapOffset;
|
||||
kv[i + 4] = gmem_in.load(base);
|
||||
score[i + 4] = gmem_in.load(base + Trait::kScoreOffset);
|
||||
}
|
||||
|
||||
/// NOTE: part 2: safe online softmax + weighted sum
|
||||
using StorageOut = AlignedVector<OutFloat, kTileElements>;
|
||||
const auto gmem_out = tile::Memory<StorageOut>::warp();
|
||||
StorageOut result;
|
||||
|
||||
// consume 32 fp registers
|
||||
float score_fp32[kTileElements][8];
|
||||
|
||||
// convert to fp32 and apply bias first
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < kTileElements; ++i) {
|
||||
for (int32_t j = 0; j < 8; ++j) {
|
||||
score_fp32[i][j] = cast<float>(score[j][i]) + cast<float>(bias[j][i]);
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < kTileElements; ++i) {
|
||||
const auto& score = score_fp32[i];
|
||||
float max_value = score[0];
|
||||
float sum_exp_value = 0.0f;
|
||||
|
||||
#pragma unroll
|
||||
for (int32_t j = 1; j < 8; ++j) {
|
||||
const auto fp32_score = score[j];
|
||||
max_value = fmaxf(max_value, fp32_score);
|
||||
}
|
||||
|
||||
float sum_product = 0.0f;
|
||||
#pragma unroll
|
||||
for (int32_t j = 0; j < 8; ++j) {
|
||||
const auto fp32_score = score[j];
|
||||
const auto exp_score = expf(fp32_score - max_value);
|
||||
sum_product += cast<float>(kv[j][i]) * exp_score;
|
||||
sum_exp_value += exp_score;
|
||||
}
|
||||
|
||||
result[i] = cast<OutFloat>(sum_product / sum_exp_value);
|
||||
}
|
||||
|
||||
// overlap the store with the next iteration's load
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
gmem_out.store(kv_out, result);
|
||||
}
|
||||
|
||||
template <typename Trait, typename InFloat>
|
||||
SGL_DEVICE void c4_write_decode(InFloat* kv_buf, const InFloat* kv_src) {
|
||||
using namespace device;
|
||||
|
||||
using StorageIn = AlignedVector<InFloat, kTileElements>;
|
||||
const auto gmem = tile::Memory<StorageIn>::warp();
|
||||
|
||||
StorageIn data[4];
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < 4; ++i) {
|
||||
data[i] = gmem.load(kv_src + Trait::kHeadDim * i);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < 4; ++i) {
|
||||
gmem.store(kv_buf + Trait::kHeadDim * i, data[i]);
|
||||
}
|
||||
}
|
||||
|
||||
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
|
||||
C4_KERNEL void flash_c4_decode(const __grid_constant__ Compress4DecodeParams params) {
|
||||
using namespace device;
|
||||
using Trait = C4Trait<kHeadDim>;
|
||||
|
||||
const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const uint32_t global_wid = global_tid / kWarpThreads; // warp id
|
||||
const uint32_t global_bid = global_wid / Trait::kNumSplit; // batch id
|
||||
const uint32_t global_sid = global_wid % Trait::kNumSplit; // split id
|
||||
const int64_t split_offset = global_sid * Trait::kTileDim;
|
||||
if (global_bid >= params.batch_size) return;
|
||||
|
||||
const auto plan = params.plan_d[global_bid];
|
||||
const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
|
||||
const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
|
||||
const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
|
||||
const auto score_bias = static_cast<const InFloat*>(params.score_bias) + split_offset;
|
||||
|
||||
const auto kv_src = kv_input + global_bid * Trait::kElementSize;
|
||||
const auto kv_out = kv_output + global_bid * Trait::kHeadDim;
|
||||
const auto kv_buf_0 = kv_buffer + plan.read_page_0 * Trait::kPageElementSize;
|
||||
const auto kv_buf_1 = kv_buffer + plan.read_page_1 * Trait::kPageElementSize;
|
||||
const auto kv_dst = kv_buffer + plan.write_loc * Trait::kElementSize;
|
||||
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
c4_write_decode<Trait>(kv_dst, kv_src);
|
||||
if (plan.seq_len % 4 == 0) {
|
||||
const auto need_overlap = plan.seq_len > 4;
|
||||
c4_forward<Trait, kUsePDL>(kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, 8);
|
||||
}
|
||||
}
|
||||
|
||||
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
|
||||
C4_KERNEL void flash_c4_prefill(const __grid_constant__ Compress4PrefillParams params) {
|
||||
using namespace device;
|
||||
using Trait = C4Trait<kHeadDim>;
|
||||
|
||||
const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const uint32_t global_wid = global_tid / kWarpThreads; // warp id
|
||||
const uint32_t global_pid = global_wid / Trait::kNumSplit; // plan id
|
||||
const uint32_t global_sid = global_wid % Trait::kNumSplit; // split id
|
||||
const int64_t split_offset = global_sid * Trait::kTileDim;
|
||||
if (global_pid >= params.num_compress) return;
|
||||
|
||||
const auto plan = params.plan_c[global_pid];
|
||||
const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
|
||||
const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
|
||||
const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
|
||||
const auto score_bias = static_cast<const InFloat*>(params.score_bias) + split_offset;
|
||||
if (plan.is_invalid()) return;
|
||||
|
||||
const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize;
|
||||
// Compact output: one row per compress plan, indexed by `global_pid`.
|
||||
const auto kv_out = kv_output + global_pid * Trait::kHeadDim;
|
||||
const auto kv_buf_0 = kv_buffer + plan.read_page_0 * Trait::kPageElementSize;
|
||||
const auto kv_buf_1 = kv_buffer + plan.read_page_1 * Trait::kPageElementSize;
|
||||
const bool need_overlap = plan.seq_len > 4;
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
c4_forward<Trait, kUsePDL>(kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, plan.buffer_len);
|
||||
}
|
||||
|
||||
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
|
||||
WRITE_KERNEL void write_c4_prefill(const __grid_constant__ Compress4PrefillParams params) {
|
||||
using namespace device;
|
||||
using Trait = C4Trait<kHeadDim>;
|
||||
using StorageIn = AlignedVector<InFloat, kTileElements>;
|
||||
|
||||
const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const uint32_t global_wid = global_tid / kWarpThreads; // warp id
|
||||
const uint32_t global_pid = global_wid / Trait::kNumSplit; // plan id
|
||||
const uint32_t global_sid = global_wid % Trait::kNumSplit; // split id
|
||||
// split the contiguous `kHeadDim * 4` into `kNumSplit` tiles
|
||||
// each warp handles 1 contiguous tile (in contrast, decode handle the strided head_dim)
|
||||
const int64_t split_offset = global_sid * (Trait::kTileDim * 4);
|
||||
if (global_pid >= params.num_write) return;
|
||||
|
||||
const auto plan = params.plan_w[global_pid];
|
||||
const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
|
||||
const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
|
||||
if (plan.is_invalid()) return;
|
||||
|
||||
// each warp will handle a contiguous region
|
||||
const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize;
|
||||
const auto kv_buf = kv_buffer + plan.write_loc * Trait::kElementSize;
|
||||
const auto gmem = tile::Memory<StorageIn>::warp();
|
||||
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
StorageIn data[4];
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < 4; ++i) {
|
||||
data[i] = gmem.load(kv_src, i);
|
||||
}
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < 4; ++i) {
|
||||
gmem.store(kv_buf, data[i], i);
|
||||
}
|
||||
}
|
||||
|
||||
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
|
||||
struct FlashCompress4Kernel {
|
||||
static constexpr auto decode_kernel = flash_c4_decode<kHeadDim, InFloat, OutFloat, kUsePDL>;
|
||||
static constexpr auto prefill_c_kernel = flash_c4_prefill<kHeadDim, InFloat, OutFloat, kUsePDL>;
|
||||
static constexpr auto prefill_w_kernel = write_c4_prefill<kHeadDim, InFloat, OutFloat, kUsePDL>;
|
||||
static constexpr uint32_t kBlockSize = 128;
|
||||
static constexpr uint32_t kTileDim = kTileElements * device::kWarpThreads;
|
||||
static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
|
||||
static constexpr uint32_t kWarpsPerBlock = kBlockSize / device::kWarpThreads;
|
||||
using Trait = C4Trait<kHeadDim>;
|
||||
|
||||
static void run_decode(
|
||||
const tvm::ffi::TensorView kv_buffer,
|
||||
const tvm::ffi::TensorView kv_input,
|
||||
const tvm::ffi::TensorView kv_output,
|
||||
const tvm::ffi::TensorView ape,
|
||||
const tvm::ffi::TensorView plan_d_) {
|
||||
using namespace host;
|
||||
|
||||
auto N = SymbolicSize{"batch_size"};
|
||||
auto device_ = SymbolicDevice{};
|
||||
device_.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({-1, 4, Trait::kElementSize}) // kv score
|
||||
.with_dtype<InFloat>()
|
||||
.with_device(device_)
|
||||
.verify(kv_buffer);
|
||||
TensorMatcher({N, Trait::kElementSize}) // kv score input
|
||||
.with_dtype<InFloat>()
|
||||
.with_device(device_)
|
||||
.verify(kv_input);
|
||||
TensorMatcher({N, kHeadDim}) // kv compressed output
|
||||
.with_dtype<OutFloat>()
|
||||
.with_device(device_)
|
||||
.verify(kv_output);
|
||||
TensorMatcher({8, kHeadDim}) // ape
|
||||
.with_dtype<InFloat>()
|
||||
.with_device(device_)
|
||||
.verify(ape);
|
||||
|
||||
const auto plan_d = compress::verify_plan_d(plan_d_, N, device_);
|
||||
const auto batch_size = static_cast<uint32_t>(N.unwrap());
|
||||
const auto params = Compress4DecodeParams{
|
||||
.kv_buffer = kv_buffer.data_ptr(),
|
||||
.kv_input = kv_input.data_ptr(),
|
||||
.kv_output = kv_output.data_ptr(),
|
||||
.score_bias = ape.data_ptr(),
|
||||
.plan_d = plan_d,
|
||||
.batch_size = batch_size,
|
||||
};
|
||||
const uint32_t num_blocks = div_ceil(batch_size * kNumSplit, kWarpsPerBlock);
|
||||
LaunchKernel(num_blocks, kBlockSize, device_.unwrap()) //
|
||||
.enable_pdl(kUsePDL)(decode_kernel, params);
|
||||
}
|
||||
|
||||
static void run_prefill(
|
||||
const tvm::ffi::TensorView kv_buffer,
|
||||
const tvm::ffi::TensorView kv_input,
|
||||
const tvm::ffi::TensorView kv_output,
|
||||
const tvm::ffi::TensorView ape,
|
||||
const tvm::ffi::TensorView plan_c_,
|
||||
const tvm::ffi::TensorView plan_w_) {
|
||||
using namespace host;
|
||||
|
||||
auto N = SymbolicSize{"num_q_tokens"};
|
||||
auto C = SymbolicSize{"num_c_plans"};
|
||||
auto W = SymbolicSize{"num_w_plans"};
|
||||
auto device_ = SymbolicDevice{};
|
||||
device_.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({-1, 4, Trait::kElementSize}) // kv score
|
||||
.with_dtype<InFloat>()
|
||||
.with_device(device_)
|
||||
.verify(kv_buffer);
|
||||
TensorMatcher({N, Trait::kElementSize}) // kv score input (ragged)
|
||||
.with_dtype<InFloat>()
|
||||
.with_device(device_)
|
||||
.verify(kv_input);
|
||||
TensorMatcher({C, kHeadDim}) // kv compressed output (compact)
|
||||
.with_dtype<OutFloat>()
|
||||
.with_device(device_)
|
||||
.verify(kv_output);
|
||||
TensorMatcher({8, kHeadDim}) // ape
|
||||
.with_dtype<InFloat>()
|
||||
.with_device(device_)
|
||||
.verify(ape);
|
||||
const auto plan_c = compress::verify_plan_c(plan_c_, C, device_);
|
||||
const auto plan_w = compress::verify_plan_w(plan_w_, W, device_);
|
||||
const auto device = device_.unwrap();
|
||||
const auto num_q_tokens = static_cast<uint32_t>(N.unwrap());
|
||||
const auto num_c = static_cast<uint32_t>(C.unwrap());
|
||||
const auto num_w = static_cast<uint32_t>(W.unwrap());
|
||||
const auto params = Compress4PrefillParams{
|
||||
.kv_buffer = kv_buffer.data_ptr(),
|
||||
.kv_input = kv_input.data_ptr(),
|
||||
.kv_output = kv_output.data_ptr(),
|
||||
.score_bias = ape.data_ptr(),
|
||||
.plan_c = plan_c,
|
||||
.plan_w = plan_w,
|
||||
.num_compress = num_c,
|
||||
.num_write = num_w,
|
||||
};
|
||||
RuntimeCheck(num_q_tokens >= num_w, "invalid prefill plan: num_q < num_w");
|
||||
if (const auto num_c_blocks = div_ceil(num_c * kNumSplit, kWarpsPerBlock)) {
|
||||
LaunchKernel(num_c_blocks, kBlockSize, device) //
|
||||
.enable_pdl(kUsePDL)(prefill_c_kernel, params);
|
||||
}
|
||||
if (const auto num_w_blocks = div_ceil(num_w * kNumSplit, kWarpsPerBlock)) {
|
||||
LaunchKernel(num_w_blocks, kBlockSize, device) //
|
||||
.enable_pdl(kUsePDL)(prefill_w_kernel, params);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
@@ -0,0 +1,827 @@
|
||||
#include <sgl_kernel/ffi.h>
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
#include <sgl_kernel/warp.cuh>
|
||||
|
||||
#include <sgl_kernel/deepseek_v4/compress_v2.cuh>
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/container/tuple.h>
|
||||
|
||||
#include <cstdint>
|
||||
#include <limits>
|
||||
|
||||
namespace host::compress {
|
||||
|
||||
constexpr auto kDLUInt8 = DLDataType{.code = kDLUInt, .bits = 8, .lanes = 1};
|
||||
|
||||
using PlanC = CompressPlan;
|
||||
using PlanW = WritePlan;
|
||||
using PlanD = DecodePlan;
|
||||
|
||||
using RID_T = int64_t;
|
||||
using R2T_T = int32_t;
|
||||
using F2S_T = int64_t;
|
||||
using IDX_T = int64_t;
|
||||
|
||||
/// NOTE: for the internal use, we pack the ragged and batch id, since both not exceed 65536
|
||||
SGL_DEVICE __host__ PlanW pack_w(uint32_t ragged_id, uint32_t batch_id, int32_t seq_len) {
|
||||
return {static_cast<uint32_t>(ragged_id | batch_id << 16), seq_len};
|
||||
}
|
||||
|
||||
/// NOTE: for the internal use, we pack the ragged and batch id, since both not exceed 65536
|
||||
SGL_DEVICE uint2 unpack_w(PlanW plan) {
|
||||
return {static_cast<uint16_t>(plan.ragged_id), static_cast<uint16_t>(plan.ragged_id >> 16)};
|
||||
}
|
||||
|
||||
struct Prefill0Params {
|
||||
PlanC* plan_c;
|
||||
PlanW* plan_w;
|
||||
const IDX_T* seq_lens_ptr; // [batch_size]
|
||||
const IDX_T* extend_lens_ptr; // [batch_size]
|
||||
uint32_t batch_size;
|
||||
uint32_t num_q_tokens;
|
||||
int32_t compress_ratio;
|
||||
int32_t swa_page_size;
|
||||
int32_t mtp_pad;
|
||||
};
|
||||
|
||||
struct Prefill1Params {
|
||||
PlanC* plan_c;
|
||||
PlanW* plan_w;
|
||||
const RID_T* rid_ptr; // [batch_size]
|
||||
const R2T_T* r2t_ptr; // [num_reqs, stride_r2t]
|
||||
const F2S_T* f2s_ptr; // [num_swa_slots]
|
||||
int64_t stride_r2t;
|
||||
uint32_t num_c;
|
||||
uint32_t num_w;
|
||||
uint32_t num_c_padded;
|
||||
uint32_t num_w_padded;
|
||||
uint32_t num_work;
|
||||
int32_t swa_page_size;
|
||||
int32_t ring_size;
|
||||
int32_t compress_ratio;
|
||||
};
|
||||
|
||||
struct DecodeParams {
|
||||
PlanD* plan_d;
|
||||
const RID_T* rid_ptr; // [batch_size]
|
||||
const R2T_T* r2t_ptr; // [num_reqs, stride_r2t]
|
||||
const F2S_T* f2s_ptr; // [num_swa_slots]
|
||||
const IDX_T* seq_ptr; // [batch_size]
|
||||
int64_t stride_r2t;
|
||||
uint32_t batch_size;
|
||||
int32_t swa_page_size;
|
||||
int32_t ring_size;
|
||||
int32_t compress_ratio;
|
||||
};
|
||||
|
||||
struct Prefill1ParamsLegacy {
|
||||
PlanC* plan_c;
|
||||
PlanW* plan_w;
|
||||
const RID_T* rid_ptr; // [batch_size]
|
||||
uint32_t num_c;
|
||||
uint32_t num_w;
|
||||
uint32_t num_c_padded;
|
||||
uint32_t num_w_padded;
|
||||
uint32_t num_work;
|
||||
int32_t compress_ratio;
|
||||
};
|
||||
|
||||
struct DecodeParamsLegacy {
|
||||
PlanD* plan_d;
|
||||
const RID_T* rid_ptr; // [batch_size]
|
||||
const IDX_T* seq_ptr; // [batch_size]
|
||||
uint32_t batch_size;
|
||||
int32_t compress_ratio;
|
||||
};
|
||||
|
||||
inline constexpr uint32_t kMaxPrefillBatchSize = 1024;
|
||||
|
||||
SGL_DEVICE uint32_t warp_inclusive_sum(uint32_t lane_id, uint32_t val) {
|
||||
static_assert(device::kWarpThreads == 32);
|
||||
#pragma unroll
|
||||
for (uint32_t offset = 1; offset < 32; offset *= 2) {
|
||||
uint32_t n = __shfl_up_sync(0xFFFFFFFF, val, offset);
|
||||
if (lane_id >= offset) val += n;
|
||||
}
|
||||
return val;
|
||||
}
|
||||
|
||||
/// Warp-wide max/min for integer types. `device::warp::reduce_max` routes through
|
||||
/// `dtype_trait<T>::max` which is only specialized for FP types.
|
||||
SGL_DEVICE uint32_t warp_reduce_max_u32(uint32_t val) {
|
||||
#pragma unroll
|
||||
for (uint32_t mask = 16; mask > 0; mask >>= 1) {
|
||||
val = max(val, __shfl_xor_sync(0xFFFFFFFF, val, mask, 32));
|
||||
}
|
||||
return val;
|
||||
}
|
||||
|
||||
SGL_DEVICE uint32_t warp_reduce_min_u32(uint32_t val) {
|
||||
#pragma unroll
|
||||
for (uint32_t mask = 16; mask > 0; mask >>= 1) {
|
||||
val = min(val, __shfl_xor_sync(0xFFFFFFFF, val, mask, 32));
|
||||
}
|
||||
return val;
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(1024, 1) //
|
||||
void plan_compress_prefill_kernel0(const Prefill0Params params) {
|
||||
using namespace device;
|
||||
const auto tx = threadIdx.x;
|
||||
const auto block_size = kMaxPrefillBatchSize;
|
||||
constexpr auto kNumWarps = kMaxPrefillBatchSize / kWarpThreads;
|
||||
const auto cr = params.compress_ratio;
|
||||
const auto sps = params.swa_page_size;
|
||||
const bool is_overlap = (cr == 4);
|
||||
const int32_t window_size = cr * (is_overlap ? 2 : 1);
|
||||
|
||||
alignas(128) __shared__ uint32_t counter_c;
|
||||
alignas(128) __shared__ uint32_t counter_w;
|
||||
__shared__ int32_t s_seq_len[kMaxPrefillBatchSize];
|
||||
__shared__ int32_t s_prefix_len[kMaxPrefillBatchSize];
|
||||
__shared__ uint32_t warp_max[kNumWarps];
|
||||
__shared__ uint32_t warp_min[kNumWarps];
|
||||
__shared__ uint32_t s_max_extend;
|
||||
__shared__ uint32_t s_min_extend;
|
||||
|
||||
const auto lane_id = tx % kWarpThreads;
|
||||
const auto warp_id = tx / kWarpThreads;
|
||||
|
||||
// === Stage A: load per-batch fields, init shared scratch ===
|
||||
int32_t seq_len = 0, extend_len = 0, prefix_len = 0;
|
||||
if (tx < params.batch_size) {
|
||||
seq_len = static_cast<int32_t>(params.seq_lens_ptr[tx]);
|
||||
extend_len = static_cast<int32_t>(params.extend_lens_ptr[tx]);
|
||||
prefix_len = seq_len - extend_len;
|
||||
s_seq_len[tx] = seq_len;
|
||||
s_prefix_len[tx] = prefix_len;
|
||||
}
|
||||
if (tx == 0) {
|
||||
counter_c = 0;
|
||||
counter_w = 0;
|
||||
}
|
||||
if (tx < kNumWarps) {
|
||||
warp_max[tx] = 0;
|
||||
warp_min[tx] = 0xFFFFFFFFu;
|
||||
}
|
||||
|
||||
// === Stage B: min/max(extend_len) for MTP-uniform detection ===
|
||||
// For min, treat threads outside `batch_size` as +inf so they don't pull the min down.
|
||||
const uint32_t e_for_max = static_cast<uint32_t>(extend_len);
|
||||
const uint32_t e_for_min = (tx < params.batch_size) ? e_for_max : 0xFFFFFFFFu;
|
||||
warp_max[warp_id] = warp_reduce_max_u32(e_for_max);
|
||||
warp_min[warp_id] = warp_reduce_min_u32(e_for_min);
|
||||
__syncthreads();
|
||||
if (warp_id == 0) {
|
||||
s_max_extend = warp_reduce_max_u32(warp_max[lane_id]);
|
||||
s_min_extend = warp_reduce_min_u32(warp_min[lane_id]);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto num_q = params.num_q_tokens;
|
||||
// MTP-uniform: every batch shares the same small extend_len `E`, so we can decompose
|
||||
// a global token id `k` into (batch_id, j) = (k / E, k % E) and skip the per-batch loop.
|
||||
const bool is_mtp_extend = (s_min_extend == s_max_extend) && (s_max_extend > 0) && (s_max_extend <= 32);
|
||||
|
||||
// === Stage C: emit valid plans, slot allocation via shared-mem atomicAdd ===
|
||||
if (is_mtp_extend) {
|
||||
// Path 1: token-driven. Each global token id maps to exactly one (batch_id, j).
|
||||
const uint32_t E = s_max_extend;
|
||||
for (uint32_t k = tx; k < num_q; k += block_size) {
|
||||
const uint32_t batch_id = k / E;
|
||||
const uint32_t j = k % E;
|
||||
const int32_t pl = s_prefix_len[batch_id];
|
||||
const int32_t sl = s_seq_len[batch_id];
|
||||
const int32_t position = pl + static_cast<int32_t>(j);
|
||||
const uint32_t ragged_id = k;
|
||||
|
||||
if ((position + 1) % cr == 0) {
|
||||
const int32_t buffer_len = window_size - min(static_cast<int32_t>(j) + 1, window_size);
|
||||
const uint32_t out_idx = atomicAdd(&counter_c, 1u);
|
||||
params.plan_c[out_idx] = {
|
||||
.seq_len = static_cast<uint32_t>(position + 1),
|
||||
.ragged_id = static_cast<uint16_t>(ragged_id),
|
||||
.buffer_len = static_cast<uint16_t>(buffer_len),
|
||||
.read_page_0 = -1,
|
||||
.read_page_1 = static_cast<int32_t>(batch_id),
|
||||
};
|
||||
}
|
||||
|
||||
const int32_t last_c_pos = (sl / cr) * cr;
|
||||
const int32_t first_w_pos = min(last_c_pos - (is_overlap ? cr : 0), sl - params.mtp_pad);
|
||||
bool do_write = position >= first_w_pos;
|
||||
if (!do_write && is_overlap) do_write = (position % sps) >= (sps - cr);
|
||||
if (do_write) {
|
||||
const uint32_t out_idx = atomicAdd(&counter_w, 1u);
|
||||
params.plan_w[out_idx] = pack_w(ragged_id, batch_id, position + 1);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Path 2: general prefill (long extend_len). Iterate batches in an outer loop;
|
||||
// the whole block sweeps each batch's tokens in parallel.
|
||||
uint32_t base_e = 0;
|
||||
for (uint32_t batch_id = 0; batch_id < params.batch_size; ++batch_id) {
|
||||
const int32_t pl = s_prefix_len[batch_id];
|
||||
const int32_t sl = s_seq_len[batch_id];
|
||||
const int32_t el = sl - pl;
|
||||
const int32_t last_c_pos = (sl / cr) * cr;
|
||||
const int32_t first_w_pos = min(last_c_pos - (is_overlap ? cr : 0), sl - params.mtp_pad);
|
||||
for (int32_t j = static_cast<int32_t>(tx); j < el; j += static_cast<int32_t>(block_size)) {
|
||||
const int32_t position = pl + j;
|
||||
const uint32_t ragged_id = base_e + static_cast<uint32_t>(j);
|
||||
|
||||
if ((position + 1) % cr == 0) {
|
||||
const int32_t buffer_len = window_size - min(j + 1, window_size);
|
||||
const uint32_t out_idx = atomicAdd(&counter_c, 1u);
|
||||
params.plan_c[out_idx] = {
|
||||
.seq_len = static_cast<uint32_t>(position + 1),
|
||||
.ragged_id = static_cast<uint16_t>(ragged_id),
|
||||
.buffer_len = static_cast<uint16_t>(buffer_len),
|
||||
.read_page_0 = -1,
|
||||
.read_page_1 = static_cast<int32_t>(batch_id),
|
||||
};
|
||||
}
|
||||
|
||||
bool do_write = position >= first_w_pos;
|
||||
if (!do_write && is_overlap) do_write = (position % sps) >= (sps - cr);
|
||||
if (do_write) {
|
||||
const uint32_t out_idx = atomicAdd(&counter_w, 1u);
|
||||
params.plan_w[out_idx] = pack_w(ragged_id, static_cast<uint32_t>(batch_id), position + 1);
|
||||
}
|
||||
}
|
||||
base_e += static_cast<uint32_t>(el);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// === Stage D: pad [counter_c, num_q) / [counter_w, num_q) with invalid ===
|
||||
const auto total_c = counter_c;
|
||||
const auto total_w = counter_w;
|
||||
for (uint32_t k = total_c + tx; k < num_q; k += block_size) {
|
||||
params.plan_c[k] = PlanC::invalid();
|
||||
}
|
||||
for (uint32_t k = total_w + tx; k < num_q; k += block_size) {
|
||||
params.plan_w[k] = PlanW::invalid();
|
||||
}
|
||||
}
|
||||
|
||||
/// NOTE: stage 1
|
||||
__global__ void plan_compress_prefill_kernel_1(const Prefill1Params params) {
|
||||
const auto idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= params.num_work) return;
|
||||
auto plan_c = idx < params.num_c ? params.plan_c[idx] : PlanC::invalid();
|
||||
auto plan_w = idx < params.num_w ? params.plan_w[idx] : PlanW::invalid();
|
||||
|
||||
const auto compute_loc = [&](int32_t swa_loc) {
|
||||
const auto swa_page = swa_loc / params.swa_page_size;
|
||||
const auto ring_offset = swa_loc % params.ring_size;
|
||||
return swa_page * params.ring_size + ring_offset;
|
||||
};
|
||||
|
||||
if (!plan_c.is_invalid()) { // 1. in bound. 2. not masked
|
||||
if (plan_c.buffer_len > 0) {
|
||||
const auto batch_id = plan_c.read_page_1;
|
||||
const auto rid = params.rid_ptr[batch_id];
|
||||
const auto mapping = params.r2t_ptr + rid * params.stride_r2t;
|
||||
// `seq_len` should be ratio-aligned here
|
||||
const auto position_1 = static_cast<int32_t>(plan_c.seq_len - 1);
|
||||
// only used for c4, harmless for c128
|
||||
const auto position_0 = max(position_1 - params.compress_ratio, 0);
|
||||
const auto raw_loc_0 = mapping[position_0];
|
||||
const auto raw_loc_1 = mapping[position_1];
|
||||
const auto swa_loc_0 = params.f2s_ptr[raw_loc_0];
|
||||
const auto swa_loc_1 = params.f2s_ptr[raw_loc_1];
|
||||
plan_c.read_page_0 = compute_loc(swa_loc_0) / params.compress_ratio;
|
||||
plan_c.read_page_1 = compute_loc(swa_loc_1) / params.compress_ratio;
|
||||
params.plan_c[idx] = plan_c;
|
||||
}
|
||||
} else if (idx < params.num_c_padded) {
|
||||
params.plan_c[idx] = PlanC::invalid();
|
||||
}
|
||||
|
||||
if (!plan_w.is_invalid()) { // 1. in bound. 2. not masked
|
||||
const auto [ragged_id, batch_id] = unpack_w(plan_w);
|
||||
const auto rid = params.rid_ptr[batch_id];
|
||||
const auto mapping = params.r2t_ptr + rid * params.stride_r2t;
|
||||
// `seq_len` (`write_loc`) may not be aligned here
|
||||
const auto position = static_cast<int32_t>(plan_w.write_loc - 1);
|
||||
const auto raw_loc = mapping[position];
|
||||
const auto swa_loc = params.f2s_ptr[raw_loc];
|
||||
plan_w.ragged_id = ragged_id;
|
||||
plan_w.write_loc = compute_loc(swa_loc);
|
||||
params.plan_w[idx] = plan_w;
|
||||
} else if (idx < params.num_w_padded) {
|
||||
params.plan_w[idx] = PlanW::invalid();
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void plan_compress_decode_kernel(const DecodeParams params) {
|
||||
const auto idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= params.batch_size) return;
|
||||
const auto rid = params.rid_ptr[idx];
|
||||
const auto mapping = params.r2t_ptr + rid * params.stride_r2t;
|
||||
const auto compute_loc = [&](int32_t swa_loc) {
|
||||
const auto swa_page = swa_loc / params.swa_page_size;
|
||||
const auto ring_offset = swa_loc % params.ring_size;
|
||||
return swa_page * params.ring_size + ring_offset;
|
||||
};
|
||||
const auto seq_len = static_cast<int32_t>(params.seq_ptr[idx]);
|
||||
const auto position_1 = static_cast<int32_t>(seq_len - 1);
|
||||
const auto position_0 = max(position_1 - params.compress_ratio, 0);
|
||||
const auto raw_loc_0 = mapping[position_0];
|
||||
const auto raw_loc_1 = mapping[position_1];
|
||||
const auto swa_loc_0 = params.f2s_ptr[raw_loc_0];
|
||||
const auto swa_loc_1 = params.f2s_ptr[raw_loc_1];
|
||||
const auto write_loc = compute_loc(swa_loc_1);
|
||||
const auto read_page_0 = compute_loc(swa_loc_0) / params.compress_ratio;
|
||||
const auto read_page_1 = write_loc / params.compress_ratio;
|
||||
params.plan_d[idx] = {
|
||||
.seq_len = static_cast<uint32_t>(seq_len),
|
||||
.write_loc = write_loc,
|
||||
.read_page_0 = read_page_0,
|
||||
.read_page_1 = read_page_1,
|
||||
};
|
||||
}
|
||||
|
||||
__global__ void plan_compress_prefill_legacy_kernel(const Prefill1ParamsLegacy params) {
|
||||
const auto idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= params.num_work) return;
|
||||
auto plan_c = idx < params.num_c ? params.plan_c[idx] : PlanC::invalid();
|
||||
auto plan_w = idx < params.num_w ? params.plan_w[idx] : PlanW::invalid();
|
||||
|
||||
/// Per-request ring buffer slot translation:
|
||||
/// - c4: page = rid * 2 + (position / 4) % 2; slot = page * 4 + position % 4
|
||||
/// - c128: page = rid; slot = rid * 128 + position % 128
|
||||
const auto legacy_compute_page = [&](int32_t rid, int32_t position) {
|
||||
if (params.compress_ratio == 4) return rid * 2 + ((position / 4) & 1);
|
||||
return rid; // c128
|
||||
};
|
||||
const auto legacy_compute_loc = [&](int32_t rid, int32_t position) {
|
||||
const auto remainder = position % params.compress_ratio;
|
||||
return legacy_compute_page(rid, position) * params.compress_ratio + remainder;
|
||||
};
|
||||
|
||||
if (!plan_c.is_invalid()) {
|
||||
const auto batch_id = plan_c.read_page_1;
|
||||
const auto rid = static_cast<int32_t>(params.rid_ptr[batch_id]);
|
||||
// `seq_len` is ratio-aligned for compress events
|
||||
const auto position_1 = static_cast<int32_t>(plan_c.seq_len) - 1;
|
||||
const auto position_0 = max(position_1 - params.compress_ratio, 0);
|
||||
plan_c.read_page_0 = legacy_compute_page(rid, position_0);
|
||||
plan_c.read_page_1 = legacy_compute_page(rid, position_1);
|
||||
params.plan_c[idx] = plan_c;
|
||||
} else if (idx < params.num_c_padded) {
|
||||
params.plan_c[idx] = PlanC::invalid();
|
||||
}
|
||||
|
||||
if (!plan_w.is_invalid()) {
|
||||
const auto [ragged_id, batch_id] = unpack_w(plan_w);
|
||||
const auto rid = static_cast<int32_t>(params.rid_ptr[batch_id]);
|
||||
// `write_loc` carries (position + 1) at this stage; may not be ratio-aligned
|
||||
const auto position = static_cast<int32_t>(plan_w.write_loc) - 1;
|
||||
plan_w.ragged_id = ragged_id;
|
||||
plan_w.write_loc = legacy_compute_loc(rid, position);
|
||||
params.plan_w[idx] = plan_w;
|
||||
} else if (idx < params.num_w_padded) {
|
||||
params.plan_w[idx] = PlanW::invalid();
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void plan_compress_decode_legacy_kernel(const DecodeParamsLegacy params) {
|
||||
const auto idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= params.batch_size) return;
|
||||
/// Per-request ring buffer slot translation:
|
||||
/// - c4: page = rid * 2 + (position / 4) % 2; slot = page * 4 + position % 4
|
||||
/// - c128: page = rid; slot = rid * 128 + position % 128
|
||||
const auto legacy_compute_page = [&](int32_t rid, int32_t position) {
|
||||
if (params.compress_ratio == 4) return rid * 2 + ((position / 4) & 1);
|
||||
return rid; // c128
|
||||
};
|
||||
const auto legacy_compute_loc = [&](int32_t rid, int32_t position) {
|
||||
const auto remainder = position % params.compress_ratio;
|
||||
return legacy_compute_page(rid, position) * params.compress_ratio + remainder;
|
||||
};
|
||||
const auto rid = static_cast<int32_t>(params.rid_ptr[idx]);
|
||||
const auto seq_len = static_cast<int32_t>(params.seq_ptr[idx]);
|
||||
const auto position_1 = seq_len - 1;
|
||||
const auto position_0 = max(position_1 - params.compress_ratio, 0);
|
||||
const auto write_loc = legacy_compute_loc(rid, position_1);
|
||||
const auto read_page_0 = legacy_compute_page(rid, position_0);
|
||||
const auto read_page_1 = legacy_compute_page(rid, position_1);
|
||||
params.plan_d[idx] = {
|
||||
.seq_len = static_cast<uint32_t>(seq_len),
|
||||
.write_loc = write_loc,
|
||||
.read_page_0 = read_page_0,
|
||||
.read_page_1 = read_page_1,
|
||||
};
|
||||
}
|
||||
|
||||
using PrefillPlan = tvm::ffi::Tuple<tvm::ffi::Tensor, tvm::ffi::Tensor>;
|
||||
|
||||
/**
|
||||
* \brief Build c4/c128 prefill plan tensors. CPU-resident.
|
||||
* Inputs (all CPU-resident):
|
||||
* @param req_pool_indices `[batch_size]` int64_t
|
||||
* @param req_to_token `[num_reqs, max_tokens_per_req]` int64_t
|
||||
* @param full_to_swa `[num_swa_slots]` int64_t
|
||||
* @param seq_lens `[batch_size]` int64
|
||||
* @param extend_lens `[batch_size]` int64
|
||||
* @param compress_plan `[num_q_tokens, 16]` uint8 (output)
|
||||
* @param write_plan `[num_q_tokens, 8]` uint8 (output)
|
||||
* @param compress_ratio 4 for c4, 128 for c128
|
||||
* @param use_cuda_graph Whether the plans will be used with cuda graph (affects padding)
|
||||
* @return (compress plan tensor, write plan tensor)
|
||||
*/
|
||||
inline PrefillPlan plan_compress_prefill(
|
||||
const tvm::ffi::TensorView req_pool_indices, // GPU
|
||||
const tvm::ffi::TensorView req_to_token, // GPU
|
||||
const tvm::ffi::TensorView full_to_swa, // GPU
|
||||
const tvm::ffi::TensorView seq_lens, // CPU/GPU
|
||||
const tvm::ffi::TensorView extend_lens, // CPU/GPU
|
||||
const tvm::ffi::TensorView pin_buffer, // CPU
|
||||
const uint32_t num_q_tokens,
|
||||
const int32_t compress_ratio,
|
||||
const int32_t swa_page_size,
|
||||
const int32_t ring_size,
|
||||
const bool use_cuda_graph) {
|
||||
auto B = SymbolicSize{"batch_size"};
|
||||
auto N = SymbolicSize{"num_q_tokens"};
|
||||
auto cpu_or_gpu = SymbolicDevice{};
|
||||
auto device_ = SymbolicDevice{};
|
||||
cpu_or_gpu.set_options<kDLCPU, kDLCUDA>();
|
||||
device_.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({B}) //
|
||||
.with_dtype<RID_T>()
|
||||
.with_device(device_)
|
||||
.verify(req_pool_indices);
|
||||
TensorMatcher({-1, -1}) //
|
||||
.with_dtype<R2T_T>()
|
||||
.with_device(device_)
|
||||
.verify(req_to_token);
|
||||
TensorMatcher({-1}) //
|
||||
.with_dtype<F2S_T>()
|
||||
.with_device(device_)
|
||||
.verify(full_to_swa);
|
||||
TensorMatcher({B}) //
|
||||
.with_dtype<IDX_T>()
|
||||
.with_device(cpu_or_gpu)
|
||||
.verify(seq_lens)
|
||||
.verify(extend_lens);
|
||||
TensorMatcher({-1}) //
|
||||
.with_dtype<uint8_t>()
|
||||
.with_device<kDLCPU>()
|
||||
.verify(pin_buffer);
|
||||
|
||||
const bool is_overlap = (compress_ratio == 4);
|
||||
const int32_t window_size = compress_ratio * (is_overlap ? 2 : 1);
|
||||
|
||||
const auto seq_ptr = static_cast<const IDX_T*>(seq_lens.data_ptr());
|
||||
const auto ext_ptr = static_cast<const IDX_T*>(extend_lens.data_ptr());
|
||||
const auto rid_ptr = static_cast<const RID_T*>(req_pool_indices.data_ptr());
|
||||
const auto r2t_ptr = static_cast<const R2T_T*>(req_to_token.data_ptr());
|
||||
const auto f2s_ptr = static_cast<const F2S_T*>(full_to_swa.data_ptr());
|
||||
|
||||
const auto batch_size = static_cast<uint32_t>(B.unwrap());
|
||||
constexpr auto kMaxTokens = static_cast<uint32_t>(std::numeric_limits<uint16_t>::max());
|
||||
RuntimeCheck(compress_ratio == 4 || compress_ratio == 128);
|
||||
RuntimeCheck(batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens);
|
||||
// `swa_page_size` >= `ring_size` >= `compress_ratio`
|
||||
RuntimeCheck(swa_page_size % ring_size == 0 && ring_size % compress_ratio == 0);
|
||||
|
||||
const auto device = device_.unwrap();
|
||||
const auto stream = LaunchKernel::resolve_device(device);
|
||||
|
||||
constexpr int32_t kMaxMTPDraftTokens = 4;
|
||||
const auto mtp_pad = std::min(ring_size - compress_ratio, kMaxMTPDraftTokens);
|
||||
|
||||
if (cpu_or_gpu.unwrap().device_type == kDLCUDA) {
|
||||
// GPU input path: kernel0 builds the (CPU-loop-equivalent) plan metadata directly
|
||||
// on device, padding to num_q_tokens with invalid; kernel_1 then finalizes the
|
||||
// SWA-translated read/write locations. Used for MTP / cuda-graph capture where
|
||||
// a host sync would be expensive.
|
||||
RuntimeCheck(batch_size <= kMaxPrefillBatchSize, "GPU plan only support batch size up to ", kMaxPrefillBatchSize);
|
||||
auto C = ffi::empty({num_q_tokens, sizeof(PlanC)}, kDLUInt8, device);
|
||||
auto W = ffi::empty({num_q_tokens, sizeof(PlanW)}, kDLUInt8, device);
|
||||
const auto params0 = Prefill0Params{
|
||||
.plan_c = static_cast<PlanC*>(C.data_ptr()),
|
||||
.plan_w = static_cast<PlanW*>(W.data_ptr()),
|
||||
.seq_lens_ptr = seq_ptr,
|
||||
.extend_lens_ptr = ext_ptr,
|
||||
.batch_size = batch_size,
|
||||
.num_q_tokens = num_q_tokens,
|
||||
.compress_ratio = compress_ratio,
|
||||
.swa_page_size = swa_page_size,
|
||||
.mtp_pad = mtp_pad,
|
||||
};
|
||||
LaunchKernel(1, kMaxPrefillBatchSize, device)(plan_compress_prefill_kernel0, params0);
|
||||
// kernel_1 sees the already-padded buffers, so num_c == num_w == num_padded == num_q_tokens.
|
||||
const auto params1 = Prefill1Params{
|
||||
.plan_c = static_cast<PlanC*>(C.data_ptr()),
|
||||
.plan_w = static_cast<PlanW*>(W.data_ptr()),
|
||||
.rid_ptr = rid_ptr,
|
||||
.r2t_ptr = r2t_ptr,
|
||||
.f2s_ptr = f2s_ptr,
|
||||
.stride_r2t = req_to_token.stride(0),
|
||||
.num_c = num_q_tokens,
|
||||
.num_w = num_q_tokens,
|
||||
.num_c_padded = num_q_tokens,
|
||||
.num_w_padded = num_q_tokens,
|
||||
.num_work = num_q_tokens,
|
||||
.swa_page_size = swa_page_size,
|
||||
.ring_size = ring_size,
|
||||
.compress_ratio = compress_ratio,
|
||||
};
|
||||
const auto block_size_1 = 256;
|
||||
const auto num_blocks_1 = div_ceil(params1.num_work, block_size_1);
|
||||
LaunchKernel(num_blocks_1, block_size_1, device)(plan_compress_prefill_kernel_1, params1);
|
||||
return PrefillPlan{std::move(C), std::move(W)};
|
||||
}
|
||||
|
||||
// CPU input path: only here do we need the pinned scratch buffer.
|
||||
const auto pin_buffer_bytes = static_cast<size_t>(pin_buffer.numel()) * sizeof(uint8_t);
|
||||
RuntimeCheck(pin_buffer_bytes >= num_q_tokens * (sizeof(PlanC) + sizeof(PlanW)));
|
||||
const auto plan_c_ptr = reinterpret_cast<PlanC*>(pin_buffer.data_ptr());
|
||||
const auto plan_w_ptr = reinterpret_cast<PlanW*>(plan_c_ptr + num_q_tokens);
|
||||
|
||||
uint32_t counter = 0;
|
||||
uint32_t counter_c = 0;
|
||||
uint32_t counter_w = 0;
|
||||
|
||||
const auto should_compress = [=](int32_t position) { return (position + 1) % compress_ratio == 0; };
|
||||
for (const auto i : irange(batch_size)) {
|
||||
const int32_t seq_len = seq_ptr[i];
|
||||
const int32_t extend_len = ext_ptr[i];
|
||||
const int32_t prefix_len = seq_len - extend_len;
|
||||
const int32_t last_c_pos = seq_len / compress_ratio * compress_ratio;
|
||||
const int32_t first_w_pos = last_c_pos - (is_overlap ? compress_ratio : 0);
|
||||
RuntimeCheck(0 < extend_len && extend_len <= seq_len);
|
||||
const auto should_write = [=](int32_t position) {
|
||||
if (position >= first_w_pos) return true;
|
||||
return is_overlap && position % swa_page_size >= (swa_page_size - compress_ratio);
|
||||
};
|
||||
for (const auto j : irange(extend_len)) {
|
||||
const int32_t position = prefix_len + j;
|
||||
const int32_t ragged_id = counter + j;
|
||||
if (should_compress(position)) {
|
||||
const auto buffer_len = window_size - std::min(j + 1, window_size);
|
||||
plan_c_ptr[counter_c++] = {
|
||||
.seq_len = static_cast<uint32_t>(position + 1),
|
||||
.ragged_id = static_cast<uint16_t>(ragged_id),
|
||||
.buffer_len = static_cast<uint16_t>(buffer_len),
|
||||
// to be filled by kernel
|
||||
.read_page_0 = -1,
|
||||
.read_page_1 = static_cast<int32_t>(i),
|
||||
};
|
||||
}
|
||||
if (should_write(position)) {
|
||||
plan_w_ptr[counter_w++] = pack_w(ragged_id, i, position + 1);
|
||||
}
|
||||
}
|
||||
counter += extend_len;
|
||||
}
|
||||
RuntimeCheck(counter == num_q_tokens);
|
||||
|
||||
const auto copy_to_device = [stream](void* cuda_ptr, auto* host_ptr, size_t count) {
|
||||
const auto size_bytes = count * sizeof(*host_ptr);
|
||||
RuntimeDeviceCheck(cudaMemcpyAsync(cuda_ptr, host_ptr, size_bytes, cudaMemcpyHostToDevice, stream));
|
||||
};
|
||||
const auto num_c_padded = use_cuda_graph ? num_q_tokens : counter_c;
|
||||
const auto num_w_padded = use_cuda_graph ? num_q_tokens : counter_w;
|
||||
auto C = ffi::empty({num_c_padded, sizeof(PlanC)}, kDLUInt8, device);
|
||||
auto W = ffi::empty({num_w_padded, sizeof(PlanW)}, kDLUInt8, device);
|
||||
copy_to_device(C.data_ptr(), plan_c_ptr, counter_c);
|
||||
copy_to_device(W.data_ptr(), plan_w_ptr, counter_w);
|
||||
const auto params = Prefill1Params{
|
||||
.plan_c = static_cast<PlanC*>(C.data_ptr()),
|
||||
.plan_w = static_cast<PlanW*>(W.data_ptr()),
|
||||
.rid_ptr = rid_ptr,
|
||||
.r2t_ptr = r2t_ptr,
|
||||
.f2s_ptr = f2s_ptr,
|
||||
.stride_r2t = req_to_token.size(1),
|
||||
.num_c = counter_c,
|
||||
.num_w = counter_w,
|
||||
.num_c_padded = num_c_padded,
|
||||
.num_w_padded = num_w_padded,
|
||||
.num_work = std::max(num_c_padded, num_w_padded),
|
||||
.swa_page_size = swa_page_size,
|
||||
.ring_size = ring_size,
|
||||
.compress_ratio = compress_ratio,
|
||||
};
|
||||
const auto block_size = 256;
|
||||
const auto num_blocks = div_ceil(params.num_work, block_size);
|
||||
LaunchKernel(num_blocks, block_size, device)(plan_compress_prefill_kernel_1, params);
|
||||
return PrefillPlan{std::move(C), std::move(W)};
|
||||
}
|
||||
|
||||
inline tvm::ffi::Tensor plan_compress_decode(
|
||||
const tvm::ffi::TensorView req_pool_indices, // GPU
|
||||
const tvm::ffi::TensorView req_to_token, // GPU
|
||||
const tvm::ffi::TensorView full_to_swa, // GPU
|
||||
const tvm::ffi::TensorView seq_lens, // CPU/GPU
|
||||
const int32_t compress_ratio,
|
||||
const int32_t swa_page_size,
|
||||
const int32_t ring_size) {
|
||||
auto B = SymbolicSize{"batch_size"};
|
||||
auto device_ = SymbolicDevice{};
|
||||
device_.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({B}) //
|
||||
.with_dtype<RID_T>()
|
||||
.with_device(device_)
|
||||
.verify(req_pool_indices);
|
||||
TensorMatcher({-1, -1}) //
|
||||
.with_dtype<R2T_T>()
|
||||
.with_device(device_)
|
||||
.verify(req_to_token);
|
||||
TensorMatcher({-1}) //
|
||||
.with_dtype<F2S_T>()
|
||||
.with_device(device_)
|
||||
.verify(full_to_swa);
|
||||
TensorMatcher({B}) //
|
||||
.with_dtype<IDX_T>()
|
||||
.with_device(device_)
|
||||
.verify(seq_lens);
|
||||
|
||||
const auto batch_size = static_cast<uint32_t>(B.unwrap());
|
||||
const auto device = device_.unwrap();
|
||||
auto D = ffi::empty({batch_size, sizeof(PlanD)}, kDLUInt8, device);
|
||||
const auto params = DecodeParams{
|
||||
.plan_d = static_cast<PlanD*>(D.data_ptr()),
|
||||
.rid_ptr = static_cast<const RID_T*>(req_pool_indices.data_ptr()),
|
||||
.r2t_ptr = static_cast<const R2T_T*>(req_to_token.data_ptr()),
|
||||
.f2s_ptr = static_cast<const F2S_T*>(full_to_swa.data_ptr()),
|
||||
.seq_ptr = static_cast<const IDX_T*>(seq_lens.data_ptr()),
|
||||
.stride_r2t = req_to_token.size(1),
|
||||
.batch_size = batch_size,
|
||||
.swa_page_size = swa_page_size,
|
||||
.ring_size = ring_size,
|
||||
.compress_ratio = compress_ratio,
|
||||
};
|
||||
const auto block_size = 256;
|
||||
const auto num_blocks = div_ceil(batch_size, block_size);
|
||||
LaunchKernel(num_blocks, block_size, device)(plan_compress_decode_kernel, params);
|
||||
return D;
|
||||
}
|
||||
|
||||
/**
|
||||
* \brief Build c4/c128 prefill plan tensors for the legacy non-paged ring
|
||||
* buffer. Uses only `req_pool_indices` to derive ring slots:
|
||||
* - c4 (overlap): each request occupies 2 contiguous pages (8 token slots)
|
||||
* - c128: each request occupies 1 page (128 token slots)
|
||||
*
|
||||
* Inputs:
|
||||
* @param req_pool_indices `[batch_size]` int64 (GPU)
|
||||
* @param seq_lens `[batch_size]` int64 (CPU)
|
||||
* @param extend_lens `[batch_size]` int64 (CPU)
|
||||
* @param pin_buffer pinned scratch (CPU uint8)
|
||||
* @return (compress plan tensor, write plan tensor)
|
||||
*/
|
||||
inline PrefillPlan plan_compress_prefill_legacy(
|
||||
const tvm::ffi::TensorView req_pool_indices, // GPU
|
||||
const tvm::ffi::TensorView seq_lens, // CPU
|
||||
const tvm::ffi::TensorView extend_lens, // CPU
|
||||
const tvm::ffi::TensorView pin_buffer, // CPU
|
||||
const uint32_t num_q_tokens,
|
||||
const int32_t compress_ratio,
|
||||
const bool use_cuda_graph) {
|
||||
auto B = SymbolicSize{"batch_size"};
|
||||
auto device_ = SymbolicDevice{};
|
||||
device_.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({B}) //
|
||||
.with_dtype<RID_T>()
|
||||
.with_device(device_)
|
||||
.verify(req_pool_indices);
|
||||
TensorMatcher({B}) //
|
||||
.with_dtype<IDX_T>()
|
||||
.with_device<kDLCPU>()
|
||||
.verify(seq_lens)
|
||||
.verify(extend_lens);
|
||||
TensorMatcher({-1}) //
|
||||
.with_dtype<uint8_t>()
|
||||
.with_device<kDLCPU>()
|
||||
.verify(pin_buffer);
|
||||
|
||||
const auto pin_buffer_bytes = static_cast<size_t>(pin_buffer.numel()) * sizeof(uint8_t);
|
||||
RuntimeCheck(pin_buffer_bytes >= num_q_tokens * (sizeof(PlanC) + sizeof(PlanW)));
|
||||
const auto plan_c_ptr = reinterpret_cast<PlanC*>(pin_buffer.data_ptr());
|
||||
const auto plan_w_ptr = reinterpret_cast<PlanW*>(plan_c_ptr + num_q_tokens);
|
||||
|
||||
const bool is_overlap = (compress_ratio == 4);
|
||||
const auto seq_ptr = static_cast<const IDX_T*>(seq_lens.data_ptr());
|
||||
const auto ext_ptr = static_cast<const IDX_T*>(extend_lens.data_ptr());
|
||||
const auto rid_ptr = static_cast<const RID_T*>(req_pool_indices.data_ptr());
|
||||
|
||||
const auto window_size = compress_ratio * (is_overlap ? 2 : 1);
|
||||
const auto batch_size = static_cast<uint32_t>(B.unwrap());
|
||||
constexpr auto kMaxTokens = static_cast<uint32_t>(std::numeric_limits<uint16_t>::max());
|
||||
RuntimeCheck(compress_ratio == 4 || compress_ratio == 128);
|
||||
RuntimeCheck(batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens);
|
||||
|
||||
uint32_t counter = 0;
|
||||
uint32_t counter_c = 0;
|
||||
uint32_t counter_w = 0;
|
||||
const auto should_compress = [=](int32_t position) { return (position + 1) % compress_ratio == 0; };
|
||||
for (const auto i : irange(batch_size)) {
|
||||
const int32_t seq_len = seq_ptr[i];
|
||||
const int32_t extend_len = ext_ptr[i];
|
||||
const int32_t prefix_len = seq_len - extend_len;
|
||||
const int32_t last_c_pos = seq_len / compress_ratio * compress_ratio;
|
||||
const int32_t first_w_pos = last_c_pos - (is_overlap ? compress_ratio : 0);
|
||||
RuntimeCheck(0 < extend_len && extend_len <= seq_len);
|
||||
const auto should_write = [=](int32_t position) { return position >= first_w_pos; };
|
||||
for (const auto j : irange(extend_len)) {
|
||||
const int32_t position = prefix_len + j;
|
||||
const int32_t ragged_id = counter + j;
|
||||
if (should_compress(position)) {
|
||||
const auto buffer_len = window_size - std::min(j + 1, window_size);
|
||||
plan_c_ptr[counter_c++] = {
|
||||
.seq_len = static_cast<uint32_t>(position + 1),
|
||||
.ragged_id = static_cast<uint16_t>(ragged_id),
|
||||
.buffer_len = static_cast<uint16_t>(buffer_len),
|
||||
// to be filled by kernel
|
||||
.read_page_0 = -1,
|
||||
.read_page_1 = static_cast<int32_t>(i),
|
||||
};
|
||||
}
|
||||
if (should_write(position)) {
|
||||
plan_w_ptr[counter_w++] = pack_w(ragged_id, i, position + 1);
|
||||
}
|
||||
}
|
||||
counter += extend_len;
|
||||
}
|
||||
RuntimeCheck(counter == num_q_tokens);
|
||||
|
||||
const auto device = device_.unwrap();
|
||||
const auto stream = LaunchKernel::resolve_device(device);
|
||||
const auto copy_to_device = [stream](void* cuda_ptr, auto* host_ptr, size_t count) {
|
||||
const auto size_bytes = count * sizeof(*host_ptr);
|
||||
RuntimeDeviceCheck(cudaMemcpyAsync(cuda_ptr, host_ptr, size_bytes, cudaMemcpyHostToDevice, stream));
|
||||
};
|
||||
const auto num_c_padded = use_cuda_graph ? num_q_tokens : counter_c;
|
||||
const auto num_w_padded = use_cuda_graph ? num_q_tokens : counter_w;
|
||||
auto C = ffi::empty({num_c_padded, sizeof(PlanC)}, kDLUInt8, device);
|
||||
auto W = ffi::empty({num_w_padded, sizeof(PlanW)}, kDLUInt8, device);
|
||||
copy_to_device(C.data_ptr(), plan_c_ptr, counter_c);
|
||||
copy_to_device(W.data_ptr(), plan_w_ptr, counter_w);
|
||||
const auto params = Prefill1ParamsLegacy{
|
||||
.plan_c = static_cast<PlanC*>(C.data_ptr()),
|
||||
.plan_w = static_cast<PlanW*>(W.data_ptr()),
|
||||
.rid_ptr = rid_ptr,
|
||||
.num_c = counter_c,
|
||||
.num_w = counter_w,
|
||||
.num_c_padded = num_c_padded,
|
||||
.num_w_padded = num_w_padded,
|
||||
.num_work = std::max(num_c_padded, num_w_padded),
|
||||
.compress_ratio = compress_ratio,
|
||||
};
|
||||
const auto block_size = 256;
|
||||
const auto num_blocks = div_ceil(params.num_work, block_size);
|
||||
if (num_blocks > 0) {
|
||||
LaunchKernel(num_blocks, block_size, device)(plan_compress_prefill_legacy_kernel, params);
|
||||
}
|
||||
return PrefillPlan{std::move(C), std::move(W)};
|
||||
}
|
||||
|
||||
inline tvm::ffi::Tensor plan_compress_decode_legacy(
|
||||
const tvm::ffi::TensorView req_pool_indices, // GPU
|
||||
const tvm::ffi::TensorView seq_lens, // GPU
|
||||
const int32_t compress_ratio) {
|
||||
auto B = SymbolicSize{"batch_size"};
|
||||
auto device_ = SymbolicDevice{};
|
||||
device_.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({B}) //
|
||||
.with_dtype<RID_T>()
|
||||
.with_device(device_)
|
||||
.verify(req_pool_indices);
|
||||
TensorMatcher({B}) //
|
||||
.with_dtype<IDX_T>()
|
||||
.with_device(device_)
|
||||
.verify(seq_lens);
|
||||
RuntimeCheck(compress_ratio == 4 || compress_ratio == 128);
|
||||
|
||||
const auto batch_size = static_cast<uint32_t>(B.unwrap());
|
||||
const auto device = device_.unwrap();
|
||||
auto D = ffi::empty({batch_size, sizeof(PlanD)}, kDLUInt8, device);
|
||||
const auto params = DecodeParamsLegacy{
|
||||
.plan_d = static_cast<PlanD*>(D.data_ptr()),
|
||||
.rid_ptr = static_cast<const RID_T*>(req_pool_indices.data_ptr()),
|
||||
.seq_ptr = static_cast<const IDX_T*>(seq_lens.data_ptr()),
|
||||
.batch_size = batch_size,
|
||||
.compress_ratio = compress_ratio,
|
||||
};
|
||||
const auto block_size = 256;
|
||||
const auto num_blocks = div_ceil(batch_size, block_size);
|
||||
LaunchKernel(num_blocks, block_size, device)(plan_compress_decode_legacy_kernel, params);
|
||||
return D;
|
||||
}
|
||||
|
||||
} // namespace host::compress
|
||||
|
||||
using namespace host::compress; // expose binding
|
||||
@@ -0,0 +1,419 @@
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/tile.cuh>
|
||||
#include <sgl_kernel/type.cuh>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
#include <sgl_kernel/warp.cuh>
|
||||
|
||||
#include <sgl_kernel/deepseek_v4/compress_v2.cuh>
|
||||
#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
|
||||
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
#include <cstdint>
|
||||
|
||||
namespace {
|
||||
|
||||
using PlanC = device::compress::CompressPlan;
|
||||
using PlanD = device::compress::DecodePlan;
|
||||
using deepseek_v4::fp8::cast_to_ue8m0;
|
||||
using deepseek_v4::fp8::inv_scale_ue8m0;
|
||||
using deepseek_v4::fp8::pack_fp8;
|
||||
|
||||
constexpr uint32_t kBlockSize = 256;
|
||||
constexpr uint32_t kNumWarps = kBlockSize / device::kWarpThreads;
|
||||
|
||||
struct FusedNormRopeStoreParams {
|
||||
void* __restrict__ input;
|
||||
const void* __restrict__ handle; // plan decode / compress
|
||||
const void* __restrict__ weight;
|
||||
const float* __restrict__ freqs_cis;
|
||||
const int32_t* __restrict__ out_loc;
|
||||
uint8_t* __restrict__ kvcache;
|
||||
float eps;
|
||||
uint32_t compress_ratio;
|
||||
uint32_t num_tokens;
|
||||
};
|
||||
|
||||
enum class ForwardMode : bool {
|
||||
CompressExtend = 0,
|
||||
CompressDecode = 1,
|
||||
};
|
||||
|
||||
#define INDEXER_KERNEL __global__ __launch_bounds__(kBlockSize, 8)
|
||||
#define FLASHMLA_KERNEL __global__ __launch_bounds__(kBlockSize, 8)
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// Indexer variant: kHeadDim = 128, 1 token per *warp* (8 tokens per block).
|
||||
// Each warp's 32 lanes cover the full 128-elem head_dim (kVecSize = 4 each).
|
||||
// Cache layout: 132 bytes/token (128 fp8 nope + 4 fp32 scale).
|
||||
// ----------------------------------------------------------------------------
|
||||
template <typename DType, ForwardMode kMode, int32_t kPageBits, bool kUsePDL>
|
||||
INDEXER_KERNEL void fused_norm_rope_indexer(const __grid_constant__ FusedNormRopeStoreParams params) {
|
||||
using namespace device;
|
||||
using enum ForwardMode;
|
||||
|
||||
constexpr int64_t kHeadDim = 128;
|
||||
constexpr int64_t kRopeDim = 64;
|
||||
constexpr int64_t kVecSize = 4;
|
||||
constexpr uint32_t kRopeSize = kRopeDim / kVecSize;
|
||||
constexpr int64_t kPageBytes = 132ll << kPageBits;
|
||||
static_assert(kHeadDim == kWarpThreads * kVecSize);
|
||||
static_assert(kRopeDim == kWarpThreads * 2);
|
||||
static_assert(kRopeSize <= kWarpThreads);
|
||||
using Storage = AlignedVector<DType, kVecSize>;
|
||||
using Float4 = AlignedVector<float, kVecSize>;
|
||||
|
||||
const auto warp_id = threadIdx.x / kWarpThreads;
|
||||
const auto lane_id = threadIdx.x % kWarpThreads;
|
||||
const auto work_id = blockIdx.x * kNumWarps + warp_id;
|
||||
// Lanes whose 4-elem pack lies in the rope tail (= last `kRopeSize` packs).
|
||||
const bool is_rope_lane = lane_id >= kWarpThreads - kRopeSize;
|
||||
|
||||
if (work_id >= params.num_tokens) return;
|
||||
|
||||
const auto input = static_cast<DType*>(params.input) + work_id * kHeadDim;
|
||||
int32_t position;
|
||||
int32_t out_loc;
|
||||
if constexpr (kMode == CompressExtend) {
|
||||
const auto plan = static_cast<const PlanC*>(params.handle)[work_id];
|
||||
if (plan.is_invalid()) return;
|
||||
position = plan.seq_len - params.compress_ratio;
|
||||
out_loc = params.out_loc[plan.ragged_id];
|
||||
} else if constexpr (kMode == CompressDecode) {
|
||||
const auto plan = static_cast<const PlanD*>(params.handle)[work_id];
|
||||
if (plan.seq_len % params.compress_ratio != 0) return;
|
||||
position = plan.seq_len - params.compress_ratio;
|
||||
out_loc = params.out_loc[work_id];
|
||||
} else {
|
||||
static_assert(host::dependent_false_v<DType>, "Unsupported Mode");
|
||||
}
|
||||
const auto freqs_cis = params.freqs_cis + position * kRopeDim;
|
||||
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
Float4 data, freq;
|
||||
|
||||
// part 1: norm
|
||||
{
|
||||
Storage input_vec, weight_vec;
|
||||
input_vec.load(input, lane_id);
|
||||
weight_vec.load(params.weight, lane_id);
|
||||
if (is_rope_lane) freq.load(freqs_cis, lane_id - (kWarpThreads - kRopeSize));
|
||||
|
||||
float sum_of_squares = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecSize; ++i) {
|
||||
const auto fp32_input = cast<float>(input_vec[i]);
|
||||
sum_of_squares += fp32_input * fp32_input;
|
||||
}
|
||||
|
||||
sum_of_squares = warp::reduce_sum(sum_of_squares);
|
||||
const auto norm_factor = math::rsqrt(sum_of_squares / kHeadDim + params.eps);
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecSize; ++i) {
|
||||
const auto fp32_input = cast<float>(input_vec[i]);
|
||||
const auto fp32_weight = cast<float>(weight_vec[i]);
|
||||
data[i] = fp32_input * norm_factor * fp32_weight;
|
||||
}
|
||||
}
|
||||
|
||||
// part 2: rope (rope-lane only, 4 elems per lane = 2 (real, imag) pairs)
|
||||
if (is_rope_lane) {
|
||||
const auto x_real = data[0];
|
||||
const auto x_imag = data[1];
|
||||
const auto y_real = data[2];
|
||||
const auto y_imag = data[3];
|
||||
const auto freq_x_real = freq[0];
|
||||
const auto freq_x_imag = freq[1];
|
||||
const auto freq_y_real = freq[2];
|
||||
const auto freq_y_imag = freq[3];
|
||||
data[0] = x_real * freq_x_real - x_imag * freq_x_imag;
|
||||
data[1] = x_real * freq_x_imag + x_imag * freq_x_real;
|
||||
data[2] = y_real * freq_y_real - y_imag * freq_y_imag;
|
||||
data[3] = y_real * freq_y_imag + y_imag * freq_y_real;
|
||||
}
|
||||
|
||||
// part 3: hadamard transform
|
||||
{
|
||||
// Stage 1: butterfly (data[0], data[1]) and (data[2], data[3]).
|
||||
{
|
||||
const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
|
||||
data[0] = a0 + a1;
|
||||
data[1] = a0 - a1;
|
||||
data[2] = a2 + a3;
|
||||
data[3] = a2 - a3;
|
||||
}
|
||||
// Stage 2: butterfly (data[0], data[2]) and (data[1], data[3]).
|
||||
{
|
||||
const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
|
||||
data[0] = a0 + a2;
|
||||
data[1] = a1 + a3;
|
||||
data[2] = a0 - a2;
|
||||
data[3] = a1 - a3;
|
||||
}
|
||||
// Stages 3..7: cross-lane butterflies. Lower-lane (mask bit clear) keeps
|
||||
// the sum, upper-lane (mask bit set) keeps the difference. shfl_xor is
|
||||
// unsynchronized across early-returned lanes, but invalid-plan returns
|
||||
// happen above for *all* lanes of a warp (work_id is warp-uniform), so
|
||||
// the warp is intact here.
|
||||
#pragma unroll
|
||||
for (uint32_t mask = 1; mask < kWarpThreads; mask <<= 1) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecSize; ++i) {
|
||||
const float other = __shfl_xor_sync(0xFFFFFFFFu, data[i], mask, kWarpThreads);
|
||||
data[i] = (lane_id & mask) ? (other - data[i]) : (data[i] + other);
|
||||
}
|
||||
}
|
||||
const float kHadamardScale = math::rsqrt(static_cast<float>(kHeadDim));
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecSize; ++i)
|
||||
data[i] *= kHadamardScale;
|
||||
}
|
||||
|
||||
// part 4: per-warp UE8M0 quant + store. The whole warp emits one fp8 group
|
||||
// (= 128 elements) plus a single fp32 scale, matching the indexer cache
|
||||
// layout (`fused_store_indexer_cache`).
|
||||
{
|
||||
using OutStorage = AlignedVector<fp8x2_e4m3_t, 2>;
|
||||
float local_max = math::abs(data[0]);
|
||||
#pragma unroll
|
||||
for (int i = 1; i < kVecSize; ++i) {
|
||||
local_max = math::max(local_max, math::abs(data[i]));
|
||||
}
|
||||
const auto abs_max = warp::reduce_max(local_max);
|
||||
const auto scale = fmaxf(1e-4f, abs_max) / math::FP8_E4M3_MAX;
|
||||
const auto inv_scale = 1.0f / scale;
|
||||
const int32_t page = out_loc >> kPageBits;
|
||||
const int32_t offset = out_loc & ((1 << kPageBits) - 1);
|
||||
const auto page_ptr = params.kvcache + page * kPageBytes;
|
||||
const auto value_ptr = page_ptr + offset * 128;
|
||||
const auto scale_ptr = page_ptr + (128 << kPageBits) + offset * 4;
|
||||
OutStorage result;
|
||||
result[0] = pack_fp8(data[0] * inv_scale, data[1] * inv_scale);
|
||||
result[1] = pack_fp8(data[2] * inv_scale, data[3] * inv_scale);
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
result.store(value_ptr, lane_id);
|
||||
// The single fp32 scale is identical across all lanes -- write from any lane.
|
||||
if (lane_id == 0) reinterpret_cast<float*>(scale_ptr)[0] = scale;
|
||||
}
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// FlashMLA variant: kHeadDim = 512, 1 token per *block* (256 threads).
|
||||
// Each thread loads kVecSize=2 BF16, so 256 threads cover the full 512 elems.
|
||||
// Cache layout: 584 bytes/token = 448 fp8 nope + 64 (=32 bf16x2) rope + 8 scale.
|
||||
// ----------------------------------------------------------------------------
|
||||
template <typename DType, ForwardMode kMode, int32_t kPageBits, bool kUsePDL>
|
||||
FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormRopeStoreParams params) {
|
||||
using namespace device;
|
||||
using enum ForwardMode;
|
||||
|
||||
constexpr int64_t kHeadDim = 512;
|
||||
constexpr int64_t kRopeDim = 64;
|
||||
constexpr int64_t kVecSize = 2;
|
||||
// Last warp owns the rope tail. The remaining 7 warps each emit one
|
||||
// 64-element fp8 group (own UE8M0 scale).
|
||||
constexpr uint32_t kRopeWarp = kNumWarps - 1;
|
||||
constexpr int64_t kPageBytes = host::div_ceil(584ll << kPageBits, 576) * 576;
|
||||
static_assert(kHeadDim == kBlockSize * kVecSize);
|
||||
static_assert(kRopeDim == kWarpThreads * kVecSize);
|
||||
static_assert(kHeadDim - kRopeDim == kRopeWarp * kWarpThreads * kVecSize);
|
||||
using Storage = AlignedVector<DType, kVecSize>;
|
||||
using Float2 = AlignedVector<float, kVecSize>;
|
||||
|
||||
const auto tx = threadIdx.x;
|
||||
const auto warp_id = tx / kWarpThreads;
|
||||
const auto lane_id = tx % kWarpThreads;
|
||||
const auto work_id = blockIdx.x;
|
||||
|
||||
if (work_id >= params.num_tokens) return;
|
||||
|
||||
const auto input = static_cast<DType*>(params.input) + work_id * kHeadDim;
|
||||
int32_t position;
|
||||
int32_t out_loc;
|
||||
if constexpr (kMode == CompressExtend) {
|
||||
const auto plan = static_cast<const PlanC*>(params.handle)[work_id];
|
||||
if (plan.is_invalid()) return;
|
||||
position = plan.seq_len - params.compress_ratio;
|
||||
out_loc = params.out_loc[plan.ragged_id];
|
||||
} else if constexpr (kMode == CompressDecode) {
|
||||
const auto plan = static_cast<const PlanD*>(params.handle)[work_id];
|
||||
if (plan.seq_len % params.compress_ratio != 0) return;
|
||||
position = plan.seq_len - params.compress_ratio;
|
||||
out_loc = params.out_loc[work_id];
|
||||
} else {
|
||||
static_assert(host::dependent_false_v<DType>, "Unsupported Mode");
|
||||
}
|
||||
const auto freqs_cis = params.freqs_cis + position * kRopeDim;
|
||||
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
Float2 data, freq;
|
||||
|
||||
// part 1: norm. Each thread owns one 2-elem pack (`tx`-th pack of input).
|
||||
// Sum of squares is reduced across the whole block via per-warp partials.
|
||||
{
|
||||
__shared__ float partial_sums[kNumWarps];
|
||||
|
||||
Storage input_vec, weight_vec;
|
||||
input_vec.load(input, tx);
|
||||
weight_vec.load(params.weight, tx);
|
||||
if (warp_id == kRopeWarp) freq.load(freqs_cis, lane_id);
|
||||
|
||||
float sum_of_squares = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecSize; ++i) {
|
||||
const auto fp32_input = cast<float>(input_vec[i]);
|
||||
sum_of_squares += fp32_input * fp32_input;
|
||||
}
|
||||
|
||||
const auto warp_sum = warp::reduce_sum(sum_of_squares);
|
||||
if (lane_id == 0) partial_sums[warp_id] = warp_sum;
|
||||
__syncthreads();
|
||||
// Replicate the per-warp partial sums to a full warp and reduce. Every
|
||||
// lane-group of `kNumWarps` lanes ends up with the global sum.
|
||||
sum_of_squares = warp::reduce_sum<kNumWarps>(partial_sums[lane_id % kNumWarps]);
|
||||
const auto norm_factor = math::rsqrt(sum_of_squares / kHeadDim + params.eps);
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecSize; ++i) {
|
||||
const auto fp32_input = cast<float>(input_vec[i]);
|
||||
const auto fp32_weight = cast<float>(weight_vec[i]);
|
||||
data[i] = fp32_input * norm_factor * fp32_weight;
|
||||
}
|
||||
}
|
||||
|
||||
const int32_t page = out_loc >> kPageBits;
|
||||
const int32_t offset = out_loc & ((1 << kPageBits) - 1);
|
||||
const auto page_ptr = params.kvcache + page * kPageBytes;
|
||||
const auto value_ptr = page_ptr + offset * 576;
|
||||
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
|
||||
// part 2: rope on the rope warp (BF16 store), or per-warp FP8 quant + store.
|
||||
if (warp_id == kRopeWarp) {
|
||||
// Each rope-warp lane owns exactly one (real, imag) pair within the rope
|
||||
// tail. Apply rotation, downcast to BF16, write to the slot's rope region.
|
||||
const auto x_real = data[0];
|
||||
const auto x_imag = data[1];
|
||||
const auto freq_real = freq[0];
|
||||
const auto freq_imag = freq[1];
|
||||
data[0] = x_real * freq_real - x_imag * freq_imag;
|
||||
data[1] = x_real * freq_imag + x_imag * freq_real;
|
||||
const auto result = cast<bf16x2_t>(fp32x2_t{data[0], data[1]});
|
||||
const auto rope_ptr = value_ptr + 448;
|
||||
reinterpret_cast<bf16x2_t*>(rope_ptr)[lane_id] = result;
|
||||
} else {
|
||||
// Non-rope warp: per-warp UE8M0 group (64 elems -> 64 fp8 + 1 scale byte).
|
||||
const auto x = data[0];
|
||||
const auto y = data[1];
|
||||
const auto abs_max = warp::reduce_max(fmaxf(fabs(x), fabs(y)));
|
||||
const auto scale_raw = fmaxf(1e-4f, abs_max) / math::FP8_E4M3_MAX;
|
||||
const auto scale_ue8m0 = cast_to_ue8m0(scale_raw);
|
||||
const auto inv_scale = inv_scale_ue8m0(scale_ue8m0);
|
||||
const auto result = pack_fp8(x * inv_scale, y * inv_scale);
|
||||
const auto scale_ptr = page_ptr + (576 << kPageBits) + offset * 8;
|
||||
reinterpret_cast<fp8x2_e4m3_t*>(value_ptr)[tx] = result;
|
||||
// All lanes in this warp produce the same scale byte; let lane 0 publish.
|
||||
if (lane_id == 0) static_cast<uint8_t*>(scale_ptr)[warp_id] = scale_ue8m0;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename DType, int64_t kHeadDim, int64_t kRopeDim, uint32_t kPageSize, bool kUsePDL>
|
||||
struct FusedNormRopeKernel {
|
||||
static constexpr int32_t kLogPageSize = std::countr_zero(kPageSize);
|
||||
static constexpr bool kIsIndexer = (kHeadDim == 128);
|
||||
static constexpr int64_t kIndexerBytes = 132 * kPageSize;
|
||||
static constexpr int64_t kFlashMLABytes = host::div_ceil(584 * kPageSize, 576) * 576;
|
||||
static constexpr int64_t kPageBytes = kIsIndexer ? kIndexerBytes : kFlashMLABytes;
|
||||
|
||||
/// TODO: Let's fix the config for now.
|
||||
static_assert(kRopeDim == 64 && (kHeadDim == 128 || kHeadDim == 512));
|
||||
static_assert(std::has_single_bit(kPageSize), "kPageSize must be a power of 2");
|
||||
|
||||
template <ForwardMode kMode>
|
||||
static constexpr auto select_kernel() {
|
||||
if constexpr (kIsIndexer) {
|
||||
return fused_norm_rope_indexer<DType, kMode, kLogPageSize, kUsePDL>;
|
||||
} else {
|
||||
return fused_norm_rope_flashmla<DType, kMode, kLogPageSize, kUsePDL>;
|
||||
}
|
||||
}
|
||||
|
||||
static void forward(
|
||||
const tvm::ffi::TensorView input,
|
||||
const tvm::ffi::TensorView plan,
|
||||
const tvm::ffi::TensorView weight,
|
||||
const float eps,
|
||||
const tvm::ffi::TensorView freqs_cis,
|
||||
const tvm::ffi::TensorView out_loc,
|
||||
const tvm::ffi::TensorView kvcache,
|
||||
const bool is_decode,
|
||||
const uint32_t compress_ratio) {
|
||||
using namespace host;
|
||||
using enum ForwardMode;
|
||||
|
||||
const auto mode = static_cast<ForwardMode>(is_decode);
|
||||
|
||||
auto N = SymbolicSize{"num_tokens"};
|
||||
auto device_ = SymbolicDevice{};
|
||||
device_.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({N, kHeadDim}) // input
|
||||
.with_dtype<DType>()
|
||||
.with_device(device_)
|
||||
.verify(input);
|
||||
TensorMatcher({kHeadDim}) // weight
|
||||
.with_dtype<DType>()
|
||||
.with_device(device_)
|
||||
.verify(weight);
|
||||
TensorMatcher({-1, kRopeDim}) // freqs_cis
|
||||
.with_dtype<float>()
|
||||
.with_device(device_)
|
||||
.verify(freqs_cis);
|
||||
TensorMatcher({-1}) // out_loc
|
||||
.with_dtype<int32_t>()
|
||||
.with_device(device_)
|
||||
.verify(out_loc);
|
||||
TensorMatcher({-1, -1}) // cache
|
||||
.with_strides({kPageBytes, 1})
|
||||
.with_dtype<uint8_t>()
|
||||
.with_device(device_)
|
||||
.verify(kvcache);
|
||||
|
||||
switch (mode) {
|
||||
case CompressExtend:
|
||||
compress::verify_plan_c(plan, N, device_);
|
||||
RuntimeCheck(out_loc.size(0) >= N.unwrap());
|
||||
break;
|
||||
case CompressDecode:
|
||||
compress::verify_plan_d(plan, N, device_);
|
||||
RuntimeCheck(out_loc.size(0) == N.unwrap());
|
||||
break;
|
||||
}
|
||||
|
||||
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
|
||||
if (num_tokens == 0) return;
|
||||
const auto params = FusedNormRopeStoreParams{
|
||||
.input = input.data_ptr(),
|
||||
.handle = plan.data_ptr(),
|
||||
.weight = weight.data_ptr(),
|
||||
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
|
||||
.out_loc = static_cast<const int32_t*>(out_loc.data_ptr()),
|
||||
.kvcache = static_cast<uint8_t*>(kvcache.data_ptr()),
|
||||
.eps = eps,
|
||||
.compress_ratio = compress_ratio,
|
||||
.num_tokens = num_tokens,
|
||||
};
|
||||
// Indexer packs `kNumWarps` tokens per block (warp-major); FlashMLA uses
|
||||
// a whole block per token (cta-major sum-reduce over head_dim=512).
|
||||
const uint32_t num_blocks = kIsIndexer ? div_ceil(num_tokens, kNumWarps) : num_tokens;
|
||||
const auto device = device_.unwrap();
|
||||
const auto kernel = mode == CompressExtend ? select_kernel<CompressExtend>() : select_kernel<CompressDecode>();
|
||||
LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
@@ -0,0 +1,629 @@
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/math.cuh>
|
||||
#include <sgl_kernel/tile.cuh>
|
||||
#include <sgl_kernel/type.cuh>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
#include <sgl_kernel/warp.cuh>
|
||||
|
||||
#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
|
||||
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
#include <bit>
|
||||
#include <cstdint>
|
||||
|
||||
namespace {
|
||||
|
||||
using deepseek_v4::fp8::cast_to_ue8m0;
|
||||
using deepseek_v4::fp8::inv_scale_ue8m0;
|
||||
using deepseek_v4::fp8::pack_fp8;
|
||||
|
||||
// 4 warps per block: warp-per-(token, head) work-item dispatch (Q kernel).
|
||||
constexpr uint32_t kFusedQBlockSize = 128;
|
||||
constexpr uint32_t kFusedQNumWarps = kFusedQBlockSize / device::kWarpThreads;
|
||||
|
||||
// 8 warps per block: block-per-token work-item dispatch (K kernel).
|
||||
constexpr uint32_t kFusedKBlockSize = 256;
|
||||
constexpr uint32_t kFusedKNumWarps = kFusedKBlockSize / device::kWarpThreads;
|
||||
|
||||
#define Q_KERNEL __global__ __launch_bounds__(kFusedQBlockSize, 16)
|
||||
#define K_KERNEL __global__ __launch_bounds__(kFusedKBlockSize, 8)
|
||||
|
||||
// ============================================================================
|
||||
// Q kernel: warp-per-(token, head) rmsnorm-self + RoPE + write to q_out.
|
||||
// ============================================================================
|
||||
|
||||
struct FusedQNormRopeParams {
|
||||
const void* __restrict__ q_input; // (B, num_q_heads, kHeadDim) DType
|
||||
void* __restrict__ q_output; // (B, num_q_heads, kHeadDim) DType
|
||||
const float* __restrict__ freqs_cis; // (max_pos, kRopeDim) fp32 (re/im interleaved)
|
||||
const void* __restrict__ positions; // (B,) PosT
|
||||
int64_t q_input_stride_batch;
|
||||
int64_t q_output_stride_batch;
|
||||
uint32_t batch_size;
|
||||
uint32_t num_q_heads;
|
||||
float eps;
|
||||
};
|
||||
|
||||
template <typename DType, int64_t kHeadDim, int64_t kRopeDim, typename PosT, bool kUsePDL>
|
||||
Q_KERNEL void fused_q_norm_rope(const __grid_constant__ FusedQNormRopeParams params) {
|
||||
using namespace device;
|
||||
|
||||
constexpr int64_t kMaxVecSize = 16 / sizeof(DType);
|
||||
constexpr int64_t kVecSize = std::min(kMaxVecSize, kHeadDim / kWarpThreads);
|
||||
constexpr int64_t kLocalSize = kHeadDim / (kWarpThreads * kVecSize);
|
||||
constexpr uint32_t kRopeSize = kRopeDim / kVecSize;
|
||||
static_assert(kHeadDim % (kWarpThreads * kVecSize) == 0);
|
||||
static_assert(kLocalSize * kVecSize * kWarpThreads == kHeadDim);
|
||||
static_assert(kRopeDim % kVecSize == 0);
|
||||
static_assert(kRopeSize <= kWarpThreads);
|
||||
static_assert(kRopeDim == kWarpThreads * 2, "1 (real, imag) pair per lane");
|
||||
|
||||
using Storage = AlignedVector<DType, kVecSize>;
|
||||
|
||||
const auto warp_id = threadIdx.x / kWarpThreads;
|
||||
const auto lane_id = threadIdx.x % kWarpThreads;
|
||||
const auto work_id = blockIdx.x * kFusedQNumWarps + warp_id;
|
||||
|
||||
const uint32_t total_works = params.batch_size * params.num_q_heads;
|
||||
if (work_id >= total_works) return;
|
||||
|
||||
const uint32_t batch_id = work_id / params.num_q_heads;
|
||||
const uint32_t head_id = work_id % params.num_q_heads;
|
||||
const auto input_ptr =
|
||||
static_cast<const DType*>(params.q_input) + batch_id * params.q_input_stride_batch + head_id * kHeadDim;
|
||||
const auto output_ptr =
|
||||
static_cast<DType*>(params.q_output) + batch_id * params.q_output_stride_batch + head_id * kHeadDim;
|
||||
const auto position = static_cast<int32_t>(static_cast<const PosT*>(params.positions)[batch_id]);
|
||||
|
||||
__shared__ Storage s_rope[kFusedQNumWarps][kRopeSize];
|
||||
|
||||
// Prefetch this lane's freq pair before the PDL gate so the wait happens
|
||||
// outside the dependency chain on `position`.
|
||||
const auto mem_freq = tile::Memory<fp32x2_t>{lane_id, kWarpThreads};
|
||||
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
|
||||
// part 1: rmsnorm-self (no weight).
|
||||
const auto gmem = tile::Memory<Storage>{lane_id, kWarpThreads};
|
||||
Storage input_vec[kLocalSize];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kLocalSize; ++i) {
|
||||
input_vec[i] = gmem.load(input_ptr, i);
|
||||
}
|
||||
|
||||
const auto freq = mem_freq.load(params.freqs_cis + position * kRopeDim);
|
||||
|
||||
float sum_of_squares = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kLocalSize; ++i) {
|
||||
#pragma unroll
|
||||
for (int j = 0; j < kVecSize; ++j) {
|
||||
const auto x = cast<float>(input_vec[i][j]);
|
||||
sum_of_squares += x * x;
|
||||
}
|
||||
}
|
||||
sum_of_squares = warp::reduce_sum(sum_of_squares);
|
||||
const auto norm_factor = math::rsqrt(sum_of_squares / kHeadDim + params.eps);
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kLocalSize; ++i) {
|
||||
#pragma unroll
|
||||
for (int j = 0; j < kVecSize; ++j) {
|
||||
const auto x = cast<float>(input_vec[i][j]);
|
||||
input_vec[i][j] = cast<DType>(x * norm_factor);
|
||||
}
|
||||
}
|
||||
|
||||
// Stash the rope tail (last kRopeSize lanes' last tile) into shared memory;
|
||||
// write nope tiles to gmem directly.
|
||||
const bool is_rope_lane = lane_id >= kWarpThreads - kRopeSize;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kLocalSize; ++i) {
|
||||
if (i == kLocalSize - 1 && is_rope_lane) {
|
||||
const auto rope_id = lane_id - (kWarpThreads - kRopeSize);
|
||||
s_rope[warp_id][rope_id] = input_vec[i];
|
||||
} else {
|
||||
gmem.store(output_ptr, input_vec[i], i);
|
||||
}
|
||||
}
|
||||
__syncwarp();
|
||||
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
|
||||
// part 2: RoPE on all 32 lanes -- one (real, imag) bf16x2 pair per lane.
|
||||
using DType2 = packed_t<DType>;
|
||||
const auto mem_elem = tile::Memory<DType2>{lane_id, kWarpThreads};
|
||||
const auto elem = mem_elem.load(s_rope[warp_id]);
|
||||
const auto [x_real, x_imag] = cast<fp32x2_t>(elem);
|
||||
const auto [freq_real, freq_imag] = freq;
|
||||
const fp32x2_t rotated = {
|
||||
x_real * freq_real - x_imag * freq_imag,
|
||||
x_real * freq_imag + x_imag * freq_real,
|
||||
};
|
||||
mem_elem.store(output_ptr + (kHeadDim - kRopeDim), cast<DType2>(rotated));
|
||||
}
|
||||
|
||||
template <typename DType, int64_t kHeadDim, int64_t kRopeDim, bool kUsePDL>
|
||||
struct FusedQNormRopeKernel {
|
||||
template <typename PosT>
|
||||
static constexpr auto kernel = fused_q_norm_rope<DType, kHeadDim, kRopeDim, PosT, kUsePDL>;
|
||||
|
||||
static void forward(
|
||||
const tvm::ffi::TensorView q_input,
|
||||
const tvm::ffi::TensorView q_output,
|
||||
const tvm::ffi::TensorView freqs_cis,
|
||||
const tvm::ffi::TensorView positions,
|
||||
float eps) {
|
||||
using namespace host;
|
||||
|
||||
auto B = SymbolicSize{"batch_size"};
|
||||
auto H = SymbolicSize{"num_q_heads"};
|
||||
auto device_ = SymbolicDevice{};
|
||||
device_.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({B, H, kHeadDim}) //
|
||||
.with_strides({-1, kHeadDim, 1})
|
||||
.with_dtype<DType>()
|
||||
.with_device(device_)
|
||||
.verify(q_input);
|
||||
TensorMatcher({B, H, kHeadDim}) //
|
||||
.with_strides({-1, kHeadDim, 1})
|
||||
.with_dtype<DType>()
|
||||
.with_device(device_)
|
||||
.verify(q_output);
|
||||
TensorMatcher({-1, kRopeDim}) //
|
||||
.with_dtype<float>()
|
||||
.with_device(device_)
|
||||
.verify(freqs_cis);
|
||||
auto pos_dtype = SymbolicDType{};
|
||||
TensorMatcher({B}) //
|
||||
.with_dtype<int32_t, int64_t>(pos_dtype)
|
||||
.with_device(device_)
|
||||
.verify(positions);
|
||||
|
||||
const auto batch_size = static_cast<uint32_t>(B.unwrap());
|
||||
const auto num_q_heads = static_cast<uint32_t>(H.unwrap());
|
||||
if (batch_size == 0) return;
|
||||
|
||||
const auto params = FusedQNormRopeParams{
|
||||
.q_input = q_input.data_ptr(),
|
||||
.q_output = q_output.data_ptr(),
|
||||
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
|
||||
.positions = positions.data_ptr(),
|
||||
.q_input_stride_batch = q_input.stride(0),
|
||||
.q_output_stride_batch = q_output.stride(0),
|
||||
.batch_size = batch_size,
|
||||
.num_q_heads = num_q_heads,
|
||||
.eps = eps,
|
||||
};
|
||||
const auto total_works = batch_size * num_q_heads;
|
||||
const auto num_blocks = div_ceil(total_works, kFusedQNumWarps);
|
||||
const auto k_int32 = kernel<int32_t>;
|
||||
const auto k_int64 = kernel<int64_t>;
|
||||
const auto k = pos_dtype.is_type<int32_t>() ? k_int32 : k_int64;
|
||||
LaunchKernel(num_blocks, kFusedQBlockSize, device_.unwrap()) //
|
||||
.enable_pdl(kUsePDL)(k, params);
|
||||
}
|
||||
};
|
||||
|
||||
// ============================================================================
|
||||
// K kernel: block-per-token rmsnorm (with kv_weight) + RoPE + FlashMLA store.
|
||||
// ============================================================================
|
||||
|
||||
struct FusedKNormRopeFlashMLAParams {
|
||||
const void* __restrict__ kv; // (B, kHeadDim) DType
|
||||
const void* __restrict__ kv_weight; // (kHeadDim,) DType
|
||||
const float* __restrict__ freqs_cis; // (max_pos, kRopeDim) fp32
|
||||
const void* __restrict__ positions; // (B,) PosT
|
||||
const int32_t* __restrict__ out_loc; // (B,) int32 -> cache slot id
|
||||
uint8_t* __restrict__ kvcache; // (npages, kPageBytes) uint8
|
||||
// Row stride for `kv` in elements. Required because the upstream caller often
|
||||
// passes `qkv_a[..., q_lora_rank:]`, a non-contiguous slice whose stride[0]
|
||||
// equals `q_lora_rank + kHeadDim` rather than `kHeadDim`.
|
||||
int64_t kv_stride_batch;
|
||||
uint32_t batch_size;
|
||||
float eps;
|
||||
};
|
||||
|
||||
template <typename DType, int64_t kHeadDim, int64_t kRopeDim, typename PosT, int32_t kPageBits, bool kUsePDL>
|
||||
K_KERNEL void fused_k_norm_rope_flashmla(const __grid_constant__ FusedKNormRopeFlashMLAParams params) {
|
||||
using namespace device;
|
||||
|
||||
constexpr int64_t kVecSize = 2;
|
||||
constexpr uint32_t kRopeWarp = kFusedKNumWarps - 1;
|
||||
constexpr int64_t kPageBytes = host::div_ceil(584ll << kPageBits, 576) * 576;
|
||||
static_assert(kHeadDim == kFusedKBlockSize * kVecSize);
|
||||
static_assert(kRopeDim == kWarpThreads * kVecSize);
|
||||
static_assert(kHeadDim - kRopeDim == kRopeWarp * kWarpThreads * kVecSize);
|
||||
using Storage = AlignedVector<DType, kVecSize>;
|
||||
using Float2 = AlignedVector<float, kVecSize>;
|
||||
|
||||
const auto tx = threadIdx.x;
|
||||
const auto warp_id = tx / kWarpThreads;
|
||||
const auto lane_id = tx % kWarpThreads;
|
||||
const auto work_id = blockIdx.x;
|
||||
if (work_id >= params.batch_size) return;
|
||||
|
||||
const auto input_ptr = static_cast<const DType*>(params.kv) + work_id * params.kv_stride_batch;
|
||||
const auto position = static_cast<int32_t>(static_cast<const PosT*>(params.positions)[work_id]);
|
||||
const auto out_loc = params.out_loc[work_id];
|
||||
const auto freqs_cis = params.freqs_cis + position * kRopeDim;
|
||||
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
Float2 data, freq;
|
||||
|
||||
// part 1: norm. Each thread owns one 2-elem pack (the `tx`-th).
|
||||
// Sum-of-squares is reduced block-wide via per-warp partials.
|
||||
{
|
||||
__shared__ float partial_sums[kFusedKNumWarps];
|
||||
|
||||
Storage input_vec, weight_vec;
|
||||
input_vec.load(input_ptr, tx);
|
||||
weight_vec.load(params.kv_weight, tx);
|
||||
if (warp_id == kRopeWarp) freq.load(freqs_cis, lane_id);
|
||||
|
||||
float sum_of_squares = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecSize; ++i) {
|
||||
const auto x = cast<float>(input_vec[i]);
|
||||
sum_of_squares += x * x;
|
||||
}
|
||||
const auto warp_sum = warp::reduce_sum(sum_of_squares);
|
||||
if (lane_id == 0) partial_sums[warp_id] = warp_sum;
|
||||
__syncthreads();
|
||||
// Replicate the per-warp partial sums onto all lanes of one warp and
|
||||
// reduce. Every group of `kBlockItemNumWarps` lanes ends up with the
|
||||
// global sum.
|
||||
sum_of_squares = warp::reduce_sum<kFusedKNumWarps>(partial_sums[lane_id % kFusedKNumWarps]);
|
||||
const auto norm_factor = math::rsqrt(sum_of_squares / kHeadDim + params.eps);
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecSize; ++i) {
|
||||
const auto x = cast<float>(input_vec[i]);
|
||||
const auto w = cast<float>(weight_vec[i]);
|
||||
data[i] = x * norm_factor * w;
|
||||
}
|
||||
}
|
||||
|
||||
const int32_t page = out_loc >> kPageBits;
|
||||
const int32_t offset = out_loc & ((1 << kPageBits) - 1);
|
||||
const auto page_ptr = params.kvcache + page * kPageBytes;
|
||||
const auto value_ptr = page_ptr + offset * 576;
|
||||
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
|
||||
// part 2: rope on warp 7 (BF16 store), per-warp UE8M0 quant + store on warps 0..6.
|
||||
if (warp_id == kRopeWarp) {
|
||||
const auto x_real = data[0];
|
||||
const auto x_imag = data[1];
|
||||
const auto freq_real = freq[0];
|
||||
const auto freq_imag = freq[1];
|
||||
data[0] = x_real * freq_real - x_imag * freq_imag;
|
||||
data[1] = x_real * freq_imag + x_imag * freq_real;
|
||||
const auto result = cast<bf16x2_t>(fp32x2_t{data[0], data[1]});
|
||||
const auto rope_ptr = value_ptr + 448;
|
||||
reinterpret_cast<bf16x2_t*>(rope_ptr)[lane_id] = result;
|
||||
} else {
|
||||
const auto x = data[0];
|
||||
const auto y = data[1];
|
||||
const auto abs_max = warp::reduce_max(fmaxf(fabs(x), fabs(y)));
|
||||
const auto scale_raw = fmaxf(1e-4f, abs_max) / math::FP8_E4M3_MAX;
|
||||
const auto scale_ue8m0 = cast_to_ue8m0(scale_raw);
|
||||
const auto inv_scale = inv_scale_ue8m0(scale_ue8m0);
|
||||
const auto result = pack_fp8(x * inv_scale, y * inv_scale);
|
||||
const auto scale_ptr = page_ptr + (576 << kPageBits) + offset * 8;
|
||||
reinterpret_cast<fp8x2_e4m3_t*>(value_ptr)[tx] = result;
|
||||
if (lane_id == 0) static_cast<uint8_t*>(scale_ptr)[warp_id] = scale_ue8m0;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename DType, int64_t kHeadDim, int64_t kRopeDim, uint32_t kPageSize, bool kUsePDL>
|
||||
struct FusedKNormRopeFlashMLAKernel {
|
||||
static constexpr int32_t kLogPageSize = std::countr_zero(kPageSize);
|
||||
static constexpr int64_t kPageBytes = host::div_ceil(584 * kPageSize, 576) * 576;
|
||||
static_assert(std::has_single_bit(kPageSize), "kPageSize must be a power of 2");
|
||||
static_assert(1 << kLogPageSize == kPageSize);
|
||||
static_assert(kHeadDim == 512 && kRopeDim == 64, "FlashMLA layout requires (512, 64)");
|
||||
|
||||
template <typename PosT>
|
||||
static constexpr auto kernel = fused_k_norm_rope_flashmla<DType, kHeadDim, kRopeDim, PosT, kLogPageSize, kUsePDL>;
|
||||
|
||||
static void forward(
|
||||
const tvm::ffi::TensorView kv,
|
||||
const tvm::ffi::TensorView kv_weight,
|
||||
const tvm::ffi::TensorView freqs_cis,
|
||||
const tvm::ffi::TensorView positions,
|
||||
const tvm::ffi::TensorView out_loc,
|
||||
const tvm::ffi::TensorView kvcache,
|
||||
float eps) {
|
||||
using namespace host;
|
||||
|
||||
auto B = SymbolicSize{"batch_size"};
|
||||
auto device_ = SymbolicDevice{};
|
||||
device_.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({B, kHeadDim}) //
|
||||
.with_strides({-1, 1})
|
||||
.with_dtype<DType>()
|
||||
.with_device(device_)
|
||||
.verify(kv);
|
||||
TensorMatcher({kHeadDim}) //
|
||||
.with_dtype<DType>()
|
||||
.with_device(device_)
|
||||
.verify(kv_weight);
|
||||
TensorMatcher({-1, kRopeDim}) //
|
||||
.with_dtype<float>()
|
||||
.with_device(device_)
|
||||
.verify(freqs_cis);
|
||||
auto pos_dtype = SymbolicDType{};
|
||||
TensorMatcher({B}) //
|
||||
.with_dtype<int32_t, int64_t>(pos_dtype)
|
||||
.with_device(device_)
|
||||
.verify(positions);
|
||||
TensorMatcher({B}) //
|
||||
.with_dtype<int32_t>()
|
||||
.with_device(device_)
|
||||
.verify(out_loc);
|
||||
TensorMatcher({-1, -1}) //
|
||||
.with_strides({kPageBytes, 1})
|
||||
.with_dtype<uint8_t>()
|
||||
.with_device(device_)
|
||||
.verify(kvcache);
|
||||
|
||||
const auto batch_size = static_cast<uint32_t>(B.unwrap());
|
||||
if (batch_size == 0) return;
|
||||
|
||||
const auto params = FusedKNormRopeFlashMLAParams{
|
||||
.kv = kv.data_ptr(),
|
||||
.kv_weight = kv_weight.data_ptr(),
|
||||
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
|
||||
.positions = positions.data_ptr(),
|
||||
.out_loc = static_cast<const int32_t*>(out_loc.data_ptr()),
|
||||
.kvcache = static_cast<uint8_t*>(kvcache.data_ptr()),
|
||||
.kv_stride_batch = kv.stride(0),
|
||||
.batch_size = batch_size,
|
||||
.eps = eps,
|
||||
};
|
||||
const auto k_int32 = kernel<int32_t>;
|
||||
const auto k_int64 = kernel<int64_t>;
|
||||
const auto k = pos_dtype.is_type<int32_t>() ? k_int32 : k_int64;
|
||||
LaunchKernel(batch_size, kFusedKBlockSize, device_.unwrap()) //
|
||||
.enable_pdl(kUsePDL)(k, params);
|
||||
}
|
||||
};
|
||||
|
||||
// ============================================================================
|
||||
// Indexer Q kernel: warp-per-(token, head) RoPE + Hadamard + fp8 act-quant.
|
||||
// ============================================================================
|
||||
|
||||
struct FusedQIndexerRopeHadamardQuantParams {
|
||||
const void* __restrict__ q_input; // (B, num_heads, 128) DType
|
||||
void* __restrict__ q_fp8; // (B, num_heads, 128) fp8_e4m3
|
||||
// weights_out[b, h] = weight[b, h] * weight_scale * q_scale[b, h].
|
||||
// q_scale is computed internally and not exposed -- the only consumer of
|
||||
// it is `weights_out`.
|
||||
const void* __restrict__ weight; // (B, num_heads) DType
|
||||
float* __restrict__ weights_out; // (B, num_heads) fp32 (== (B, H, 1) flat)
|
||||
float weight_scale; // scalar c4_indexer.weight_scale
|
||||
const float* __restrict__ freqs_cis; // (max_pos, 64) fp32
|
||||
const void* __restrict__ positions; // (B,) PosT
|
||||
uint32_t batch_size;
|
||||
uint32_t num_heads;
|
||||
};
|
||||
|
||||
template <typename DType, typename PosT, bool kUsePDL>
|
||||
Q_KERNEL void fused_q_indexer_rope_hadamard_quant(const __grid_constant__ FusedQIndexerRopeHadamardQuantParams params) {
|
||||
using namespace device;
|
||||
|
||||
constexpr int64_t kHeadDim = 128;
|
||||
constexpr int64_t kRopeDim = 64;
|
||||
constexpr int64_t kVecSize = 4;
|
||||
constexpr uint32_t kRopeSize = kRopeDim / kVecSize; // = 16
|
||||
static_assert(kHeadDim == kWarpThreads * kVecSize);
|
||||
static_assert(kRopeDim == kWarpThreads * 2);
|
||||
static_assert(kRopeSize <= kWarpThreads);
|
||||
|
||||
using Storage = AlignedVector<DType, kVecSize>;
|
||||
using Float4 = AlignedVector<float, kVecSize>;
|
||||
using OutStorage = AlignedVector<fp8x2_e4m3_t, 2>; // 4 fp8 / lane
|
||||
|
||||
const auto warp_id = threadIdx.x / kWarpThreads;
|
||||
const auto lane_id = threadIdx.x % kWarpThreads;
|
||||
const auto work_id = blockIdx.x * kFusedQNumWarps + warp_id;
|
||||
// Last `kRopeSize` lanes own the rope tail; their 4-elem packs cover the
|
||||
// trailing kRopeDim elements.
|
||||
const bool is_rope_lane = lane_id >= kWarpThreads - kRopeSize;
|
||||
|
||||
const uint32_t total_works = params.batch_size * params.num_heads;
|
||||
if (work_id >= total_works) return;
|
||||
|
||||
const uint32_t batch_id = work_id / params.num_heads;
|
||||
const auto input_ptr = static_cast<const DType*>(params.q_input) + work_id * kHeadDim;
|
||||
const auto position = static_cast<int32_t>(static_cast<const PosT*>(params.positions)[batch_id]);
|
||||
const auto freqs_cis = params.freqs_cis + position * kRopeDim;
|
||||
|
||||
// Lane 0 prefetches the weight scalar for this (token, head) work item.
|
||||
// Weight is (B, num_heads) DType; we need one scalar per warp -- offload
|
||||
// the load to lane 0 only. The multiply + store happens once the q_scale
|
||||
// is known (part 4).
|
||||
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
Float4 data, freq;
|
||||
const auto weight_val = cast<float>(static_cast<const DType*>(params.weight)[work_id]);
|
||||
|
||||
// part 1: load (no norm). Each lane owns a 4-elem pack.
|
||||
{
|
||||
Storage input_vec;
|
||||
input_vec.load(input_ptr, lane_id);
|
||||
if (is_rope_lane) freq.load(freqs_cis, lane_id - (kWarpThreads - kRopeSize));
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecSize; ++i) {
|
||||
data[i] = cast<float>(input_vec[i]);
|
||||
}
|
||||
}
|
||||
|
||||
// part 2: rope on rope lanes only (4 elems / lane = 2 (real, imag) pairs).
|
||||
if (is_rope_lane) {
|
||||
const auto x_real = data[0];
|
||||
const auto x_imag = data[1];
|
||||
const auto y_real = data[2];
|
||||
const auto y_imag = data[3];
|
||||
const auto fxr = freq[0];
|
||||
const auto fxi = freq[1];
|
||||
const auto fyr = freq[2];
|
||||
const auto fyi = freq[3];
|
||||
data[0] = x_real * fxr - x_imag * fxi;
|
||||
data[1] = x_real * fxi + x_imag * fxr;
|
||||
data[2] = y_real * fyr - y_imag * fyi;
|
||||
data[3] = y_real * fyi + y_imag * fyr;
|
||||
}
|
||||
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
|
||||
// part 3: 128-point Hadamard (2 local stages + 5 cross-lane shfl_xor stages).
|
||||
// Same recipe as `fused_norm_rope_indexer`; see comments there for the
|
||||
// butterfly invariants and the early-return safety argument.
|
||||
{
|
||||
{
|
||||
const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
|
||||
data[0] = a0 + a1;
|
||||
data[1] = a0 - a1;
|
||||
data[2] = a2 + a3;
|
||||
data[3] = a2 - a3;
|
||||
}
|
||||
{
|
||||
const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
|
||||
data[0] = a0 + a2;
|
||||
data[1] = a1 + a3;
|
||||
data[2] = a0 - a2;
|
||||
data[3] = a1 - a3;
|
||||
}
|
||||
#pragma unroll
|
||||
for (uint32_t mask = 1; mask < kWarpThreads; mask <<= 1) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecSize; ++i) {
|
||||
const float other = __shfl_xor_sync(0xFFFFFFFFu, data[i], mask, kWarpThreads);
|
||||
data[i] = (lane_id & mask) ? (other - data[i]) : (data[i] + other);
|
||||
}
|
||||
}
|
||||
const float kHadamardScale = math::rsqrt(static_cast<float>(kHeadDim));
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecSize; ++i)
|
||||
data[i] *= kHadamardScale;
|
||||
}
|
||||
|
||||
{
|
||||
float local_max = math::abs(data[0]);
|
||||
#pragma unroll
|
||||
for (int i = 1; i < kVecSize; ++i) {
|
||||
local_max = math::max(local_max, math::abs(data[i]));
|
||||
}
|
||||
const auto abs_max = warp::reduce_max(local_max);
|
||||
const auto scale = fmaxf(1e-4f, abs_max) / math::FP8_E4M3_MAX;
|
||||
const auto inv_scale = 1.0f / scale;
|
||||
OutStorage result;
|
||||
result[0] = pack_fp8(data[0] * inv_scale, data[1] * inv_scale);
|
||||
result[1] = pack_fp8(data[2] * inv_scale, data[3] * inv_scale);
|
||||
|
||||
// q_fp8 row pointer: 128 fp8 / row = 32 OutStorage / row, one per lane.
|
||||
auto out_row = static_cast<uint8_t*>(params.q_fp8) + work_id * kHeadDim;
|
||||
result.store(out_row, lane_id);
|
||||
params.weights_out[work_id] = weight_val * params.weight_scale * scale;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename DType, bool kUsePDL>
|
||||
struct FusedQIndexerRopeHadamardQuantKernel {
|
||||
template <typename PosT>
|
||||
static constexpr auto kernel = fused_q_indexer_rope_hadamard_quant<DType, PosT, kUsePDL>;
|
||||
|
||||
static void forward(
|
||||
const tvm::ffi::TensorView q_input,
|
||||
const tvm::ffi::TensorView q_fp8,
|
||||
const tvm::ffi::TensorView weight,
|
||||
const tvm::ffi::TensorView weights_out,
|
||||
double weight_scale,
|
||||
const tvm::ffi::TensorView freqs_cis,
|
||||
const tvm::ffi::TensorView positions) {
|
||||
using namespace host;
|
||||
constexpr int64_t kHeadDim = 128;
|
||||
constexpr int64_t kRopeDim = 64;
|
||||
|
||||
auto B = SymbolicSize{"batch_size"};
|
||||
auto H = SymbolicSize{"num_heads"};
|
||||
auto device_ = SymbolicDevice{};
|
||||
device_.set_options<kDLCUDA>();
|
||||
|
||||
// Caller path is `wq_b(q_lora).view(-1, H, D)` -> contiguous; the kernel
|
||||
// assumes a flat `(B*H, kHeadDim)` layout for both q_input and q_fp8.
|
||||
// Pin the head/innermost strides; assert the batch stride below.
|
||||
TensorMatcher({B, H, kHeadDim}) //
|
||||
.with_strides({-1, kHeadDim, 1})
|
||||
.with_dtype<DType>()
|
||||
.with_device(device_)
|
||||
.verify(q_input);
|
||||
TensorMatcher({B, H, kHeadDim}) //
|
||||
.with_strides({-1, kHeadDim, 1})
|
||||
.with_dtype<fp8_e4m3_t>()
|
||||
.with_device(device_)
|
||||
.verify(q_fp8);
|
||||
TensorMatcher({B, H}) //
|
||||
.with_dtype<DType>()
|
||||
.with_device(device_)
|
||||
.verify(weight);
|
||||
TensorMatcher({B, H, 1}) //
|
||||
.with_dtype<float>()
|
||||
.with_device(device_)
|
||||
.verify(weights_out);
|
||||
TensorMatcher({-1, kRopeDim}) //
|
||||
.with_dtype<float>()
|
||||
.with_device(device_)
|
||||
.verify(freqs_cis);
|
||||
auto pos_dtype = SymbolicDType{};
|
||||
TensorMatcher({B}) //
|
||||
.with_dtype<int32_t, int64_t>(pos_dtype)
|
||||
.with_device(device_)
|
||||
.verify(positions);
|
||||
|
||||
const auto batch_size = static_cast<uint32_t>(B.unwrap());
|
||||
const auto num_heads = static_cast<uint32_t>(H.unwrap());
|
||||
if (batch_size == 0) return;
|
||||
|
||||
// The kernel computes row pointers as `base + work_id * kHeadDim`, so
|
||||
// both inputs must be contiguous in (batch, head, elem) order.
|
||||
const int64_t expected_batch_stride = static_cast<int64_t>(num_heads) * kHeadDim;
|
||||
RuntimeCheck(
|
||||
q_input.stride(0) == expected_batch_stride,
|
||||
"q_input must be contiguous (B, H, kHeadDim); got stride[0]=",
|
||||
q_input.stride(0));
|
||||
RuntimeCheck(
|
||||
q_fp8.stride(0) == expected_batch_stride,
|
||||
"q_fp8 must be contiguous (B, H, kHeadDim); got stride[0]=",
|
||||
q_fp8.stride(0));
|
||||
|
||||
const auto params = FusedQIndexerRopeHadamardQuantParams{
|
||||
.q_input = q_input.data_ptr(),
|
||||
.q_fp8 = q_fp8.data_ptr(),
|
||||
.weight = weight.data_ptr(),
|
||||
.weights_out = static_cast<float*>(weights_out.data_ptr()),
|
||||
.weight_scale = static_cast<float>(weight_scale),
|
||||
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
|
||||
.positions = positions.data_ptr(),
|
||||
.batch_size = batch_size,
|
||||
.num_heads = num_heads,
|
||||
};
|
||||
const auto total_works = batch_size * num_heads;
|
||||
const auto num_blocks = div_ceil(total_works, kFusedQNumWarps);
|
||||
const auto k_int32 = kernel<int32_t>;
|
||||
const auto k_int64 = kernel<int64_t>;
|
||||
const auto k = pos_dtype.is_type<int32_t>() ? k_int32 : k_int64;
|
||||
LaunchKernel(num_blocks, kFusedQBlockSize, device_.unwrap()) //
|
||||
.enable_pdl(kUsePDL)(k, params);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
@@ -18,8 +18,10 @@ constexpr uint32_t kNumWarps = kBlockSize / device::kWarpThreads;
|
||||
struct RMSNormSelfParams {
|
||||
const void* __restrict__ input;
|
||||
void* __restrict__ output;
|
||||
int64_t stride_batch_bytes;
|
||||
int64_t stride_head_bytes;
|
||||
int64_t stride_batch_bytes_0;
|
||||
int64_t stride_head_bytes_0;
|
||||
int64_t stride_batch_bytes_1;
|
||||
int64_t stride_head_bytes_1;
|
||||
uint32_t batch_size;
|
||||
uint32_t num_head;
|
||||
float eps;
|
||||
@@ -42,12 +44,12 @@ __global__ __launch_bounds__(kBlockSize, 20) //
|
||||
if (batch_id >= params.batch_size) return;
|
||||
const auto input_ptr = pointer::offset( //
|
||||
params.input,
|
||||
batch_id * params.stride_batch_bytes,
|
||||
head_id * params.stride_head_bytes);
|
||||
// use contiguous layout
|
||||
batch_id * params.stride_batch_bytes_0,
|
||||
head_id * params.stride_head_bytes_0);
|
||||
const auto output_ptr = pointer::offset( //
|
||||
params.output,
|
||||
warp_id * kHeadDim * sizeof(DType));
|
||||
batch_id * params.stride_batch_bytes_1,
|
||||
head_id * params.stride_head_bytes_1);
|
||||
PDLWaitPrimary<kUsePDL>(); // wait for primary kernel
|
||||
|
||||
Vec inputs[kNumLoop];
|
||||
@@ -93,31 +95,30 @@ struct RMSNormKernel {
|
||||
|
||||
auto N = SymbolicSize{"batch_size"};
|
||||
auto H = SymbolicSize{"num_heads"};
|
||||
auto Dn = SymbolicSize{"stride_head"};
|
||||
auto Dh = SymbolicSize{"stride_batch"};
|
||||
constexpr auto D = kHeadDim;
|
||||
auto device = SymbolicDevice{};
|
||||
device.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({N, H, D}) // input
|
||||
.with_strides({Dh, Dn, 1})
|
||||
.with_strides({-1, -1, 1})
|
||||
.with_dtype<DType>()
|
||||
.with_device(device)
|
||||
.verify(input);
|
||||
TensorMatcher({N, H, D}) // output, must be contiguous
|
||||
TensorMatcher({N, H, D}) // output
|
||||
.with_strides({-1, -1, 1})
|
||||
.with_dtype<DType>()
|
||||
.with_device(device)
|
||||
.verify(output);
|
||||
|
||||
const auto batch_size = static_cast<uint32_t>(N.unwrap());
|
||||
const auto num_head = static_cast<uint32_t>(H.unwrap());
|
||||
const auto stride_head_bytes = static_cast<int64_t>(Dn.unwrap() * sizeof(DType));
|
||||
const auto stride_batch_bytes = static_cast<int64_t>(Dh.unwrap() * sizeof(DType));
|
||||
const auto params = RMSNormSelfParams{
|
||||
.input = input.data_ptr(),
|
||||
.output = output.data_ptr(),
|
||||
.stride_batch_bytes = stride_batch_bytes,
|
||||
.stride_head_bytes = stride_head_bytes,
|
||||
.stride_batch_bytes_0 = static_cast<int64_t>(input.stride(0) * sizeof(DType)),
|
||||
.stride_head_bytes_0 = static_cast<int64_t>(input.stride(1) * sizeof(DType)),
|
||||
.stride_batch_bytes_1 = static_cast<int64_t>(output.stride(0) * sizeof(DType)),
|
||||
.stride_head_bytes_1 = static_cast<int64_t>(output.stride(1) * sizeof(DType)),
|
||||
.batch_size = batch_size,
|
||||
.num_head = num_head,
|
||||
.eps = eps,
|
||||
|
||||
@@ -195,6 +195,52 @@ def _jit_fused_store_module(
|
||||
)
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_main_q_norm_rope_module(
|
||||
dtype: torch.dtype, head_dim: int, rope_dim: int
|
||||
) -> Module:
|
||||
"""Main MLA path Q kernel: rmsnorm-self + RoPE, warp per (token, head)."""
|
||||
args = make_cpp_args(dtype, head_dim, rope_dim, is_arch_support_pdl())
|
||||
return load_jit(
|
||||
make_name("main_q_norm_rope"),
|
||||
*args,
|
||||
cuda_files=["deepseek_v4/main_norm_rope.cuh"],
|
||||
cuda_wrappers=[
|
||||
("forward", f"FusedQNormRopeKernel<{args}>::forward"),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_main_k_norm_rope_flashmla_module(
|
||||
dtype: torch.dtype, head_dim: int, rope_dim: int, page_size: int
|
||||
) -> Module:
|
||||
"""Main MLA path K kernel: rmsnorm + RoPE + write to FlashMLA paged cache."""
|
||||
args = make_cpp_args(dtype, head_dim, rope_dim, page_size, is_arch_support_pdl())
|
||||
return load_jit(
|
||||
make_name("main_k_norm_rope_flashmla"),
|
||||
*args,
|
||||
cuda_files=["deepseek_v4/main_norm_rope.cuh"],
|
||||
cuda_wrappers=[
|
||||
("forward", f"FusedKNormRopeFlashMLAKernel<{args}>::forward"),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_main_q_indexer_rope_hadamard_quant_module(dtype: torch.dtype) -> Module:
|
||||
"""C4 indexer Q kernel: RoPE + 128-pt Hadamard + fp8 act-quant (no norm)."""
|
||||
args = make_cpp_args(dtype, is_arch_support_pdl())
|
||||
return load_jit(
|
||||
make_name("main_q_indexer_rope_hadamard_quant"),
|
||||
*args,
|
||||
cuda_files=["deepseek_v4/main_norm_rope.cuh"],
|
||||
cuda_wrappers=[
|
||||
("forward", f"FusedQIndexerRopeHadamardQuantKernel<{args}>::forward"),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_metadata_module():
|
||||
return load_jit(
|
||||
@@ -571,6 +617,26 @@ def compress_fused_norm_rope_inplace(
|
||||
)
|
||||
|
||||
|
||||
def fused_norm_rope_inplace(
|
||||
kv: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
eps: float,
|
||||
freq_cis: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
) -> None:
|
||||
freq_cis = torch.view_as_real(freq_cis).flatten(-2)
|
||||
module = _jit_norm_rope_module(kv.dtype, kv.shape[-1], freq_cis.shape[-1])
|
||||
module.forward(
|
||||
kv,
|
||||
weight,
|
||||
positions,
|
||||
freq_cis,
|
||||
2,
|
||||
eps,
|
||||
0,
|
||||
)
|
||||
|
||||
|
||||
def fused_rope(
|
||||
q: torch.Tensor,
|
||||
k: Optional[torch.Tensor],
|
||||
@@ -583,6 +649,62 @@ def fused_rope(
|
||||
module.forward(q, k, freqs_real, positions, inverse)
|
||||
|
||||
|
||||
# Alias for V2 code paths
|
||||
fused_rope_inplace = fused_rope
|
||||
|
||||
|
||||
def fused_q_norm_rope(
|
||||
q_input: torch.Tensor,
|
||||
q_output: torch.Tensor,
|
||||
eps: float,
|
||||
freqs_cis: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
) -> None:
|
||||
freqs_real = torch.view_as_real(freqs_cis).flatten(-2)
|
||||
head_dim = q_input.shape[-1]
|
||||
rope_dim = freqs_real.shape[-1]
|
||||
module = _jit_main_q_norm_rope_module(q_input.dtype, head_dim, rope_dim)
|
||||
module.forward(q_input, q_output, freqs_real, positions, eps)
|
||||
|
||||
|
||||
def fused_q_indexer_rope_hadamard_quant(
|
||||
q_input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
weight_scale: float,
|
||||
freqs_cis: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
freqs_real = torch.view_as_real(freqs_cis).flatten(-2)
|
||||
q_fp8 = torch.empty(q_input.shape, dtype=torch.float8_e4m3fn, device=q_input.device)
|
||||
weights_out = torch.empty(
|
||||
(*q_input.shape[:-1], 1), dtype=torch.float32, device=q_input.device
|
||||
)
|
||||
module = _jit_main_q_indexer_rope_hadamard_quant_module(q_input.dtype)
|
||||
module.forward(
|
||||
q_input, q_fp8, weight, weights_out, float(weight_scale), freqs_real, positions
|
||||
)
|
||||
return q_fp8, weights_out
|
||||
|
||||
|
||||
def fused_k_norm_rope_flashmla(
|
||||
kv: torch.Tensor,
|
||||
kv_weight: torch.Tensor,
|
||||
eps: float,
|
||||
freqs_cis: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
out_loc: torch.Tensor,
|
||||
kvcache: torch.Tensor,
|
||||
page_size: int,
|
||||
) -> None:
|
||||
freqs_real = torch.view_as_real(freqs_cis).flatten(-2)
|
||||
head_dim = kv.shape[-1]
|
||||
rope_dim = freqs_real.shape[-1]
|
||||
module = _jit_main_k_norm_rope_flashmla_module(
|
||||
kv.dtype, head_dim, rope_dim, page_size
|
||||
)
|
||||
module.forward(kv, kv_weight, freqs_real, positions, out_loc, kvcache, eps)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def create_paged_compress_data_kernel(
|
||||
req_pool_indices_ptr,
|
||||
@@ -818,9 +940,12 @@ def get_paged_mqa_logits_metadata(seq_lens: torch.Tensor, page_size: int, num_sm
|
||||
return metadata
|
||||
|
||||
|
||||
def rmsnorm_self(q: torch.Tensor, eps: float) -> torch.Tensor:
|
||||
def rmsnorm_self(
|
||||
q: torch.Tensor, eps: float, out: Optional[torch.Tensor] = None
|
||||
) -> torch.Tensor:
|
||||
module = _jit_rmsnorm_head_module(q.shape[-1], q.dtype)
|
||||
out = q.new_empty(q.shape)
|
||||
if out is None:
|
||||
out = q.new_empty(q.shape)
|
||||
module.run_self(q, out, eps)
|
||||
return out
|
||||
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
from .compress import *
|
||||
from .utils import make_name
|
||||
|
||||
__all__ = [
|
||||
"CompressorDecodePlan",
|
||||
"CompressorPrefillPlan",
|
||||
"compress_forward",
|
||||
"compress_norm_rope_store",
|
||||
"make_name",
|
||||
]
|
||||
@@ -0,0 +1,349 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Literal, NamedTuple, Optional, Union
|
||||
|
||||
import torch
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
from sglang.jit_kernel.utils import (
|
||||
cache_once,
|
||||
is_arch_support_pdl,
|
||||
load_jit,
|
||||
make_cpp_args,
|
||||
)
|
||||
|
||||
from .utils import make_name
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_compress_norm_rope_module(
|
||||
dtype: torch.dtype,
|
||||
head_dim: int,
|
||||
rope_dim: int,
|
||||
page_size: int,
|
||||
) -> Module:
|
||||
args = make_cpp_args(dtype, head_dim, rope_dim, page_size, is_arch_support_pdl())
|
||||
return load_jit(
|
||||
make_name(f"fused_norm_rope_v2"),
|
||||
*args,
|
||||
cuda_files=[f"deepseek_v4/fused_norm_rope_v2.cuh"],
|
||||
cuda_wrappers=[("forward", f"FusedNormRopeKernel<{args}>::forward")],
|
||||
)
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_compress_module(
|
||||
head_dim: int,
|
||||
dtype_in: torch.dtype,
|
||||
dtype_out: torch.dtype,
|
||||
ratio: Literal[4, 128],
|
||||
) -> Module:
|
||||
args = make_cpp_args(head_dim, dtype_in, dtype_out, is_arch_support_pdl())
|
||||
kernel_class = f"FlashCompress{ratio}Kernel<{args}>"
|
||||
return load_jit(
|
||||
make_name(f"compress_{ratio}_v2"),
|
||||
*args,
|
||||
cuda_files=[f"deepseek_v4/c{ratio}_v2.cuh"],
|
||||
cuda_wrappers=[
|
||||
("decode", f"{kernel_class}::run_decode"),
|
||||
("prefill", f"{kernel_class}::run_prefill"),
|
||||
],
|
||||
extra_cuda_cflags=["-use_fast_math"],
|
||||
)
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_compress_128_online_module(head_dim: int) -> Module:
|
||||
assert head_dim == 512
|
||||
args = make_cpp_args(head_dim, is_arch_support_pdl())
|
||||
kernel_class = f"FlashCompress128OnlineKernel<{args}>"
|
||||
return load_jit(
|
||||
make_name(f"compress_128_online_v2"),
|
||||
*args,
|
||||
cuda_files=["deepseek_v4/c128_online_v2.cuh"],
|
||||
cuda_wrappers=[
|
||||
("decode", f"{kernel_class}::run_decode"),
|
||||
("prefill", f"{kernel_class}::run_prefill"),
|
||||
("plan_decode", "plan_compress_128_online_decode"),
|
||||
("plan_prefill", "plan_compress_128_online_prefill"),
|
||||
],
|
||||
extra_cuda_cflags=["-use_fast_math"],
|
||||
)
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_compress_plan_module() -> Module:
|
||||
return load_jit(
|
||||
make_name(f"compress_plan"),
|
||||
cuda_files=[f"deepseek_v4/c_plan.cuh"],
|
||||
cuda_wrappers=[
|
||||
("plan_prefill", "plan_compress_prefill"),
|
||||
("plan_decode", "plan_compress_decode"),
|
||||
("plan_prefill_legacy", "plan_compress_prefill_legacy"),
|
||||
("plan_decode_legacy", "plan_compress_decode_legacy"),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Plan tensor sizes (must match the C++ structs in compress.cuh).
|
||||
# ----------------------------------------------------------------------------
|
||||
_PREFILL_PLAN_BYTES = 24
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Plan dataclasses. The element at index 1 is the consumer for
|
||||
# `compress_fused_norm_rope_inplace` (which reads ragged_id / seq_len from a
|
||||
# 16-byte plan tensor --- both DecodePlan and CompressPlan satisfy that layout).
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
class CompressorDecodePlan(NamedTuple):
|
||||
compress_ratio: int
|
||||
plan_d: torch.Tensor # [batch_size, 16] uint8 --- DecodePlan
|
||||
|
||||
def copy_(self, other) -> None:
|
||||
assert isinstance(other, CompressorDecodePlan)
|
||||
assert self.compress_ratio == other.compress_ratio
|
||||
self.plan_d.copy_(other.plan_d)
|
||||
|
||||
@staticmethod
|
||||
def generate(
|
||||
compress_ratio: Literal[4, 128],
|
||||
req_pool_indices: torch.Tensor,
|
||||
req_to_token: torch.Tensor,
|
||||
full_to_swa: torch.Tensor,
|
||||
seq_lens: torch.Tensor,
|
||||
swa_page_size: int,
|
||||
ring_size: int,
|
||||
) -> CompressorDecodePlan:
|
||||
module = _jit_compress_plan_module()
|
||||
plan_d = module.plan_decode(
|
||||
req_pool_indices,
|
||||
req_to_token,
|
||||
full_to_swa,
|
||||
seq_lens,
|
||||
int(compress_ratio),
|
||||
int(swa_page_size),
|
||||
int(ring_size),
|
||||
)
|
||||
return CompressorDecodePlan(compress_ratio, torch.from_dlpack(plan_d))
|
||||
|
||||
@staticmethod
|
||||
def generate_legacy(
|
||||
compress_ratio: Literal[4, 128],
|
||||
req_pool_indices: torch.Tensor,
|
||||
seq_lens: torch.Tensor,
|
||||
) -> CompressorDecodePlan:
|
||||
module = _jit_compress_plan_module()
|
||||
plan_d = module.plan_decode_legacy(req_pool_indices, seq_lens, compress_ratio)
|
||||
return CompressorDecodePlan(compress_ratio, torch.from_dlpack(plan_d))
|
||||
|
||||
@staticmethod
|
||||
def generate_online(
|
||||
seq_lens: torch.Tensor,
|
||||
req_pool_indices: torch.Tensor,
|
||||
req_to_token: torch.Tensor,
|
||||
full_to_swa: torch.Tensor,
|
||||
swa_page_size: int,
|
||||
) -> CompressorDecodePlan:
|
||||
batch_size = int(seq_lens.shape[0])
|
||||
module = _jit_compress_128_online_module(512)
|
||||
plan_d = torch.empty(
|
||||
(batch_size, 16),
|
||||
dtype=torch.uint8,
|
||||
device=req_pool_indices.device,
|
||||
)
|
||||
module.plan_decode(
|
||||
seq_lens, req_pool_indices, req_to_token, full_to_swa, plan_d, swa_page_size
|
||||
)
|
||||
return CompressorDecodePlan(128, plan_d)
|
||||
|
||||
@property
|
||||
def is_decode(self) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
class CompressorPrefillPlan(NamedTuple):
|
||||
compress_ratio: int
|
||||
plan_c: torch.Tensor # [num_q_tokens, 16] uint8 --- CompressPlan
|
||||
plan_w: torch.Tensor # [num_q_tokens, 8] uint8 --- WritePlan
|
||||
pin_buffer: Optional[torch.Tensor] = None # keep alive
|
||||
|
||||
def copy_(self, other) -> None:
|
||||
assert isinstance(other, CompressorPrefillPlan)
|
||||
assert self.compress_ratio == other.compress_ratio
|
||||
self.plan_c.copy_(other.plan_c)
|
||||
self.plan_w.copy_(other.plan_w)
|
||||
|
||||
@staticmethod
|
||||
def generate(
|
||||
compress_ratio: Literal[4, 128],
|
||||
req_pool_indices: torch.Tensor,
|
||||
seq_lens: torch.Tensor,
|
||||
extend_lens: torch.Tensor,
|
||||
req_to_token: torch.Tensor,
|
||||
full_to_swa: torch.Tensor,
|
||||
swa_page_size: int,
|
||||
ring_size: int,
|
||||
num_q_tokens: int,
|
||||
use_cuda_graph: bool = False,
|
||||
) -> CompressorPrefillPlan:
|
||||
is_gpu_input = seq_lens.device.type == "cuda"
|
||||
pin_buffer = torch.empty(
|
||||
0 if is_gpu_input else num_q_tokens * _PREFILL_PLAN_BYTES,
|
||||
dtype=torch.uint8,
|
||||
pin_memory=not is_gpu_input,
|
||||
)
|
||||
module = _jit_compress_plan_module()
|
||||
plan_c, plan_w = module.plan_prefill(
|
||||
req_pool_indices,
|
||||
req_to_token,
|
||||
full_to_swa,
|
||||
seq_lens,
|
||||
extend_lens,
|
||||
pin_buffer,
|
||||
int(num_q_tokens),
|
||||
int(compress_ratio),
|
||||
int(swa_page_size),
|
||||
int(ring_size),
|
||||
bool(use_cuda_graph),
|
||||
)
|
||||
return CompressorPrefillPlan(
|
||||
compress_ratio,
|
||||
torch.from_dlpack(plan_c),
|
||||
torch.from_dlpack(plan_w),
|
||||
pin_buffer,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def generate_legacy(
|
||||
compress_ratio: Literal[4, 128],
|
||||
req_pool_indices: torch.Tensor,
|
||||
seq_lens: torch.Tensor,
|
||||
extend_lens: torch.Tensor,
|
||||
num_q_tokens: int,
|
||||
device: torch.device,
|
||||
use_cuda_graph: bool = False,
|
||||
) -> CompressorPrefillPlan:
|
||||
pin_buffer = torch.empty(
|
||||
num_q_tokens * _PREFILL_PLAN_BYTES,
|
||||
dtype=torch.uint8,
|
||||
pin_memory=True,
|
||||
)
|
||||
module = _jit_compress_plan_module()
|
||||
plan_c, plan_w = module.plan_prefill_legacy(
|
||||
req_pool_indices,
|
||||
seq_lens,
|
||||
extend_lens,
|
||||
pin_buffer,
|
||||
int(num_q_tokens),
|
||||
int(compress_ratio),
|
||||
bool(use_cuda_graph),
|
||||
)
|
||||
return CompressorPrefillPlan(
|
||||
compress_ratio,
|
||||
torch.from_dlpack(plan_c),
|
||||
torch.from_dlpack(plan_w),
|
||||
pin_buffer,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def generate_online(
|
||||
seq_lens: torch.Tensor,
|
||||
extend_lens: torch.Tensor,
|
||||
req_pool_indices: torch.Tensor,
|
||||
req_to_token: torch.Tensor,
|
||||
full_to_swa: torch.Tensor,
|
||||
num_q_tokens: int,
|
||||
swa_page_size: int,
|
||||
) -> CompressorPrefillPlan:
|
||||
seq_lens_cpu = seq_lens.to(torch.int64)
|
||||
extend_lens_cpu = extend_lens.to(torch.int64)
|
||||
rid_i64 = req_pool_indices.to(torch.int64)
|
||||
r2t_i32 = req_to_token.to(torch.int32)
|
||||
f2s_i64 = full_to_swa.to(torch.int64)
|
||||
pin_buffer = torch.empty(
|
||||
(2, num_q_tokens, 16), dtype=torch.uint8, pin_memory=True
|
||||
)
|
||||
plan_c_pin, plan_w_pin = pin_buffer[0], pin_buffer[1]
|
||||
device = req_pool_indices.device
|
||||
plan_c_dev = torch.empty((num_q_tokens, 16), dtype=torch.uint8, device=device)
|
||||
plan_w_dev = torch.empty((num_q_tokens, 16), dtype=torch.uint8, device=device)
|
||||
module = _jit_compress_128_online_module(512) # NOTE: only support dim=512
|
||||
num_c, num_w = module.plan_prefill(
|
||||
seq_lens_cpu,
|
||||
extend_lens_cpu,
|
||||
rid_i64,
|
||||
r2t_i32,
|
||||
f2s_i64,
|
||||
plan_c_pin,
|
||||
plan_w_pin,
|
||||
plan_c_dev,
|
||||
plan_w_dev,
|
||||
int(swa_page_size),
|
||||
)
|
||||
return CompressorPrefillPlan(
|
||||
128,
|
||||
plan_c_dev[: int(num_c)],
|
||||
plan_w_dev[: int(num_w)],
|
||||
pin_buffer,
|
||||
)
|
||||
|
||||
@property
|
||||
def is_decode(self) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def compress_forward(
|
||||
kv_score_buffer: torch.Tensor,
|
||||
kv_score_input: torch.Tensor,
|
||||
ape: torch.Tensor,
|
||||
plan: Union[CompressorDecodePlan, CompressorPrefillPlan],
|
||||
*,
|
||||
head_dim: int,
|
||||
compress_ratio: Literal[4, 128],
|
||||
out: Optional[torch.Tensor] = None,
|
||||
is_online: bool = False,
|
||||
) -> torch.Tensor:
|
||||
if out is None:
|
||||
num_q_tokens = plan[1].shape[0] # NOTE: decode = bs, prefill = dynamic
|
||||
out = kv_score_input.new_empty((num_q_tokens, head_dim))
|
||||
assert plan.compress_ratio == compress_ratio
|
||||
if is_online:
|
||||
assert compress_ratio == 128 and head_dim == 512
|
||||
module = _jit_compress_128_online_module(512)
|
||||
else:
|
||||
dtype_in, dtype_out = kv_score_input.dtype, out.dtype
|
||||
module = _jit_compress_module(head_dim, dtype_in, dtype_out, compress_ratio)
|
||||
fn = module.decode if plan.is_decode else module.prefill
|
||||
fn(kv_score_buffer, kv_score_input, out, ape, *plan[1:3])
|
||||
return out
|
||||
|
||||
|
||||
def compress_norm_rope_store(
|
||||
kv: torch.Tensor,
|
||||
plan: Union[CompressorDecodePlan, CompressorPrefillPlan],
|
||||
*,
|
||||
norm_weight: torch.Tensor,
|
||||
norm_eps: float,
|
||||
freq_cis: torch.Tensor,
|
||||
out_loc: torch.Tensor,
|
||||
kvcache: torch.Tensor,
|
||||
page_size: int,
|
||||
) -> None:
|
||||
freq_cis = torch.view_as_real(freq_cis).flatten(-2)
|
||||
module = _jit_compress_norm_rope_module(
|
||||
kv.dtype, kv.shape[-1], freq_cis.shape[-1], page_size
|
||||
)
|
||||
module.forward(
|
||||
kv,
|
||||
plan[1],
|
||||
norm_weight,
|
||||
norm_eps,
|
||||
freq_cis,
|
||||
out_loc,
|
||||
kvcache,
|
||||
plan.is_decode,
|
||||
plan.compress_ratio,
|
||||
)
|
||||
@@ -0,0 +1,2 @@
|
||||
def make_name(name: str) -> str:
|
||||
return f"dpsk_v4_{name}"
|
||||
@@ -0,0 +1,99 @@
|
||||
#pragma once
|
||||
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
#include <cstdint>
|
||||
|
||||
namespace device::compress {
|
||||
|
||||
/// \brief Per-batch decode plan. Layout: 16 bytes.
|
||||
struct alignas(16) DecodePlan {
|
||||
uint32_t seq_len;
|
||||
int32_t write_loc;
|
||||
int32_t read_page_0;
|
||||
int32_t read_page_1;
|
||||
};
|
||||
|
||||
/// \brief Per-token compress plan (used by c4/c128 prefill). Layout: 16 bytes.
|
||||
struct alignas(16) CompressPlan {
|
||||
uint32_t seq_len;
|
||||
uint16_t ragged_id;
|
||||
uint16_t buffer_len;
|
||||
int32_t read_page_0;
|
||||
/// \brief Stage 0 (CPU): batch_id (used to look up page table).
|
||||
/// \brief Stage 1 (GPU): final state-pool write location.
|
||||
int32_t read_page_1;
|
||||
|
||||
static SGL_DEVICE __host__ CompressPlan invalid() {
|
||||
return CompressPlan{-1u, 0, 0, -1, -1};
|
||||
}
|
||||
|
||||
SGL_DEVICE __host__ bool is_invalid() const {
|
||||
return seq_len == -1u;
|
||||
}
|
||||
};
|
||||
|
||||
/// \brief Per-token write plan (used by c4/c128 prefill). Layout: 8 bytes.
|
||||
struct alignas(8) WritePlan {
|
||||
/// \brief Stage 0 (CPU): packed `(batch_id << 16) | ragged_id`.
|
||||
/// \brief Stage 1 (GPU): just `ragged_id`.
|
||||
uint32_t ragged_id;
|
||||
/// \brief Stage 0 (CPU): position + 1 (used to look up state slot).
|
||||
/// \brief Stage 1 (GPU): final state-pool write location.
|
||||
int32_t write_loc;
|
||||
|
||||
static SGL_DEVICE __host__ WritePlan invalid() {
|
||||
return WritePlan{-1u, -1};
|
||||
}
|
||||
|
||||
SGL_DEVICE __host__ bool is_invalid() const {
|
||||
return ragged_id == -1u;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace device::compress
|
||||
|
||||
namespace host::compress {
|
||||
|
||||
using device::compress::CompressPlan;
|
||||
using device::compress::DecodePlan;
|
||||
using device::compress::WritePlan;
|
||||
|
||||
static_assert(alignof(DecodePlan) == sizeof(DecodePlan));
|
||||
static_assert(sizeof(DecodePlan) == 16);
|
||||
static_assert(alignof(CompressPlan) == sizeof(CompressPlan));
|
||||
static_assert(sizeof(CompressPlan) == 16);
|
||||
static_assert(alignof(WritePlan) == sizeof(WritePlan));
|
||||
static_assert(sizeof(WritePlan) == 8);
|
||||
|
||||
inline auto verify_plan_d(tvm::ffi::TensorView t, SymbolicSize& N, SymbolicDevice& device) -> const DecodePlan* {
|
||||
TensorMatcher({N, sizeof(DecodePlan)}) //
|
||||
.with_dtype<uint8_t>()
|
||||
.with_device(device)
|
||||
.verify(t);
|
||||
return static_cast<const DecodePlan*>(t.data_ptr());
|
||||
}
|
||||
|
||||
inline auto verify_plan_c(tvm::ffi::TensorView t, SymbolicSize& N, SymbolicDevice& device) -> const CompressPlan* {
|
||||
TensorMatcher({N, sizeof(CompressPlan)}) //
|
||||
.with_dtype<uint8_t>()
|
||||
.with_device(device)
|
||||
.verify(t);
|
||||
return static_cast<const CompressPlan*>(t.data_ptr());
|
||||
}
|
||||
|
||||
inline auto verify_plan_w(tvm::ffi::TensorView t, SymbolicSize& N, SymbolicDevice& device) -> const WritePlan* {
|
||||
TensorMatcher({N, sizeof(WritePlan)}) //
|
||||
.with_dtype<uint8_t>()
|
||||
.with_device(device)
|
||||
.verify(t);
|
||||
return static_cast<const WritePlan*>(t.data_ptr());
|
||||
}
|
||||
|
||||
} // namespace host::compress
|
||||
@@ -0,0 +1,195 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Literal, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.dsv4 import CompressorDecodePlan, CompressorPrefillPlan
|
||||
|
||||
|
||||
@dataclass
|
||||
class LegacyContext:
|
||||
"""Per-request ring buffer (no req_to_token / full_to_swa).
|
||||
|
||||
`req_pool_indices[i]` directly maps to the request's ring base slot.
|
||||
"""
|
||||
|
||||
bs: int
|
||||
head_dim: int
|
||||
compress_ratio: int
|
||||
req_pool_indices: torch.Tensor # int64 [bs] on cuda
|
||||
pages_per_req: int
|
||||
|
||||
@property
|
||||
def num_pages(self) -> int:
|
||||
# Reserve enough pages to hold all batched requests' rings.
|
||||
return int(self.req_pool_indices.max().item() + 1) * self.pages_per_req
|
||||
|
||||
def state_loc(self, b: int, position: int) -> int:
|
||||
rid = int(self.req_pool_indices[b].item())
|
||||
if self.compress_ratio == 4:
|
||||
page = rid * 2 + (position // 4) % 2
|
||||
else:
|
||||
page = rid
|
||||
return page * self.compress_ratio + position % self.compress_ratio
|
||||
|
||||
def make_prefill_plan(
|
||||
self,
|
||||
seq_lens_cpu: torch.Tensor,
|
||||
extend_lens_cpu: torch.Tensor,
|
||||
num_q_tokens: int,
|
||||
) -> CompressorPrefillPlan:
|
||||
return CompressorPrefillPlan.generate_legacy(
|
||||
compress_ratio=self.compress_ratio, # type: ignore
|
||||
req_pool_indices=self.req_pool_indices,
|
||||
seq_lens=seq_lens_cpu,
|
||||
extend_lens=extend_lens_cpu,
|
||||
num_q_tokens=num_q_tokens,
|
||||
device=torch.device("cuda"),
|
||||
)
|
||||
|
||||
def make_decode_plan(self, seq_lens_gpu: torch.Tensor) -> CompressorDecodePlan:
|
||||
return CompressorDecodePlan.generate_legacy(
|
||||
compress_ratio=self.compress_ratio, # type: ignore
|
||||
req_pool_indices=self.req_pool_indices,
|
||||
seq_lens=seq_lens_gpu,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class PagedContext:
|
||||
"""SWA paged layout with identity req_to_token + identity full_to_swa.
|
||||
|
||||
Each request occupies `num_swa_pages_per_req` contiguous swa_pages, so
|
||||
`req_to_token[r, p] = r * (num_swa_pages_per_req * swa_page_size) + p`.
|
||||
"""
|
||||
|
||||
bs: int
|
||||
head_dim: int
|
||||
compress_ratio: int
|
||||
swa_page_size: int
|
||||
ring_size: int
|
||||
num_swa_pages_per_req: int
|
||||
req_pool_indices: torch.Tensor # int64 [bs] on cuda
|
||||
req_to_token: torch.Tensor # int64 [num_reqs_capacity, max_tokens_per_req] on cuda
|
||||
full_to_swa: torch.Tensor # int64 [num_swa_slots] on cuda
|
||||
|
||||
@property
|
||||
def num_pages(self) -> int:
|
||||
# Upper bound: every (request, position) state slot fits.
|
||||
max_state_loc = (
|
||||
self.bs * self.num_swa_pages_per_req * self.ring_size
|
||||
+ self.swa_page_size # slack for the largest tail
|
||||
)
|
||||
return max_state_loc // self.compress_ratio + 1
|
||||
|
||||
def state_loc(self, b: int, position: int) -> int:
|
||||
rid = int(self.req_pool_indices[b].item())
|
||||
loc = int(self.req_to_token[rid, position].item())
|
||||
swa_loc = int(self.full_to_swa[loc].item())
|
||||
swa_page = swa_loc // self.swa_page_size
|
||||
return swa_page * self.ring_size + swa_loc % self.ring_size
|
||||
|
||||
def make_prefill_plan(
|
||||
self,
|
||||
seq_lens_cpu: torch.Tensor,
|
||||
extend_lens_cpu: torch.Tensor,
|
||||
num_q_tokens: int,
|
||||
) -> CompressorPrefillPlan:
|
||||
return CompressorPrefillPlan.generate(
|
||||
compress_ratio=self.compress_ratio, # type: ignore
|
||||
req_pool_indices=self.req_pool_indices,
|
||||
seq_lens=seq_lens_cpu,
|
||||
extend_lens=extend_lens_cpu,
|
||||
req_to_token=self.req_to_token,
|
||||
full_to_swa=self.full_to_swa,
|
||||
swa_page_size=self.swa_page_size,
|
||||
ring_size=self.ring_size,
|
||||
num_q_tokens=num_q_tokens,
|
||||
)
|
||||
|
||||
def make_decode_plan(self, seq_lens_gpu: torch.Tensor) -> CompressorDecodePlan:
|
||||
return CompressorDecodePlan.generate(
|
||||
compress_ratio=self.compress_ratio, # type: ignore
|
||||
req_pool_indices=self.req_pool_indices,
|
||||
req_to_token=self.req_to_token,
|
||||
full_to_swa=self.full_to_swa,
|
||||
seq_lens=seq_lens_gpu,
|
||||
swa_page_size=self.swa_page_size,
|
||||
ring_size=self.ring_size,
|
||||
)
|
||||
|
||||
|
||||
def make_legacy_context(
|
||||
bs: int,
|
||||
compress_ratio: Literal[4, 128],
|
||||
head_dim: int = 512,
|
||||
) -> LegacyContext:
|
||||
pages_per_req = 2 if compress_ratio == 4 else 1
|
||||
req_pool_indices = torch.arange(bs, dtype=torch.int64, device="cuda")
|
||||
return LegacyContext(
|
||||
bs=bs,
|
||||
head_dim=head_dim,
|
||||
compress_ratio=compress_ratio,
|
||||
req_pool_indices=req_pool_indices,
|
||||
pages_per_req=pages_per_req,
|
||||
)
|
||||
|
||||
|
||||
def make_paged_context(
|
||||
bs: int,
|
||||
compress_ratio: Literal[4, 128],
|
||||
head_dim: int = 512,
|
||||
swa_page_size: int = 256,
|
||||
ring_size: Optional[int] = None,
|
||||
num_swa_pages_per_req: int = 8,
|
||||
max_tokens_per_req: int = 8192,
|
||||
num_reqs_capacity: int = 16,
|
||||
) -> PagedContext:
|
||||
if ring_size is None:
|
||||
ring_size = 8 if compress_ratio == 4 else 128
|
||||
assert swa_page_size % ring_size == 0
|
||||
assert ring_size % compress_ratio == 0
|
||||
assert num_swa_pages_per_req * swa_page_size <= max_tokens_per_req
|
||||
|
||||
stride = num_swa_pages_per_req * swa_page_size
|
||||
req_to_token = torch.zeros(
|
||||
(num_reqs_capacity, max_tokens_per_req), dtype=torch.int32
|
||||
)
|
||||
for r in range(bs):
|
||||
req_to_token[r, :stride] = torch.arange(
|
||||
r * stride, (r + 1) * stride, dtype=torch.int32
|
||||
)
|
||||
total_swa_slots = num_reqs_capacity * stride
|
||||
full_to_swa = torch.arange(total_swa_slots, dtype=torch.int64)
|
||||
req_pool_indices = torch.arange(bs, dtype=torch.int64)
|
||||
return PagedContext(
|
||||
bs=bs,
|
||||
head_dim=head_dim,
|
||||
compress_ratio=compress_ratio,
|
||||
swa_page_size=swa_page_size,
|
||||
ring_size=ring_size,
|
||||
num_swa_pages_per_req=num_swa_pages_per_req,
|
||||
req_pool_indices=req_pool_indices.cuda(),
|
||||
req_to_token=req_to_token.cuda(),
|
||||
full_to_swa=full_to_swa.cuda(),
|
||||
)
|
||||
|
||||
|
||||
def make_state_pool(num_pages: int, compress_ratio: int, head_dim: int) -> torch.Tensor:
|
||||
last_dim = head_dim * (4 if compress_ratio == 4 else 2)
|
||||
return torch.zeros(
|
||||
(num_pages, compress_ratio, last_dim),
|
||||
dtype=torch.float32,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
|
||||
def to_seq_extend(
|
||||
seq_extend_pairs: List[Tuple[int, int]],
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, int]:
|
||||
seq_lens = torch.tensor([s for s, _ in seq_extend_pairs], dtype=torch.int64)
|
||||
extend_lens = torch.tensor([e for _, e in seq_extend_pairs], dtype=torch.int64)
|
||||
num_q = int(extend_lens.sum().item())
|
||||
return seq_lens, extend_lens, num_q
|
||||
@@ -0,0 +1,260 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from typing import Tuple, Union
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.jit_kernel.benchmark.bench_activation import register_cuda_ci
|
||||
from sglang.jit_kernel.dsv4 import compress_forward
|
||||
from sglang.jit_kernel.tests.deepseek_v4.common import (
|
||||
LegacyContext,
|
||||
PagedContext,
|
||||
make_legacy_context,
|
||||
make_paged_context,
|
||||
make_state_pool,
|
||||
to_seq_extend,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=30, suite="stage-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=30, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
Context = Union[LegacyContext, PagedContext]
|
||||
|
||||
# c128 input row layout: | kv | score | each [head_dim]
|
||||
HEAD_DIM = 512
|
||||
RATIO = 128
|
||||
ATOL = 5e-3
|
||||
RTOL = 5e-3
|
||||
|
||||
|
||||
def _gt_compress(
|
||||
kv_score_input_cpu: torch.Tensor, # [num_q, head_dim*2]
|
||||
ape_cpu: torch.Tensor, # [128, head_dim]
|
||||
P: int,
|
||||
head_dim: int,
|
||||
) -> torch.Tensor:
|
||||
"""fp64 reference for compress event at ragged position ``P`` (P % 128 == 127)."""
|
||||
lo = P - (RATIO - 1)
|
||||
kv = kv_score_input_cpu[lo : P + 1, :head_dim].double()
|
||||
sc = kv_score_input_cpu[lo : P + 1, head_dim:].double()
|
||||
return ((kv * (sc + ape_cpu.double()).softmax(dim=0)).sum(dim=0)).float()
|
||||
|
||||
|
||||
def _make_inputs(
|
||||
num_q: int, head_dim: int, seed: int
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
g = torch.Generator(device="cpu").manual_seed(seed)
|
||||
kv_score_input_cpu = torch.randn(
|
||||
num_q, head_dim * 2, generator=g, dtype=torch.float32
|
||||
)
|
||||
ape_cpu = torch.randn(RATIO, head_dim, generator=g, dtype=torch.float32)
|
||||
return kv_score_input_cpu, ape_cpu
|
||||
|
||||
|
||||
def _run_prefill(
|
||||
ctx: Context,
|
||||
pool: torch.Tensor,
|
||||
kv_score_input: torch.Tensor,
|
||||
ape: torch.Tensor,
|
||||
seq_lens_cpu: torch.Tensor,
|
||||
extend_lens_cpu: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
num_q = int(extend_lens_cpu.sum().item())
|
||||
plan = ctx.make_prefill_plan(seq_lens_cpu, extend_lens_cpu, num_q)
|
||||
return compress_forward(
|
||||
pool,
|
||||
kv_score_input,
|
||||
ape,
|
||||
plan,
|
||||
head_dim=ctx.head_dim,
|
||||
compress_ratio=RATIO,
|
||||
)
|
||||
|
||||
|
||||
def _run_decode(
|
||||
ctx: Context,
|
||||
pool: torch.Tensor,
|
||||
kv_score_input: torch.Tensor,
|
||||
ape: torch.Tensor,
|
||||
seq_lens_gpu: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
plan = ctx.make_decode_plan(seq_lens_gpu)
|
||||
return compress_forward(
|
||||
pool,
|
||||
kv_score_input,
|
||||
ape,
|
||||
plan,
|
||||
head_dim=ctx.head_dim,
|
||||
compress_ratio=RATIO,
|
||||
)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Tests
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["legacy", "paged"])
|
||||
@pytest.mark.parametrize("seq_len", [128, 256, 512])
|
||||
def test_prefill_no_context(mode: str, seq_len: int) -> None:
|
||||
"""Single-shot prefill, no prefix. Every compress event must match fp64 GT."""
|
||||
if mode == "legacy":
|
||||
ctx: Context = make_legacy_context(
|
||||
bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM
|
||||
)
|
||||
else:
|
||||
ctx = make_paged_context(bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM)
|
||||
|
||||
seq_lens_cpu, extend_lens_cpu, num_q = to_seq_extend([(seq_len, seq_len)])
|
||||
kv_in_cpu, ape_cpu = _make_inputs(num_q, ctx.head_dim, seed=seq_len)
|
||||
|
||||
pool = make_state_pool(ctx.num_pages, RATIO, ctx.head_dim)
|
||||
out = _run_prefill(
|
||||
ctx, pool, kv_in_cpu.cuda(), ape_cpu.cuda(), seq_lens_cpu, extend_lens_cpu
|
||||
)
|
||||
|
||||
# Compact prefill output: row per compress plan, in CPU-planner order.
|
||||
for plan_id, P in enumerate(range(RATIO - 1, seq_len, RATIO)):
|
||||
gt = _gt_compress(kv_in_cpu, ape_cpu, P=P, head_dim=ctx.head_dim)
|
||||
triton.testing.assert_close(out[plan_id].cpu(), gt, atol=ATOL, rtol=RTOL)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["legacy", "paged"])
|
||||
@pytest.mark.parametrize("prefix_len", [0, 128, 256])
|
||||
def test_prefill_then_decode(mode: str, prefix_len: int) -> None:
|
||||
"""Prefill ``prefix_len`` tokens, then decode through to the next 128 boundary."""
|
||||
seq_len = prefix_len + RATIO # one full compress chunk after prefix
|
||||
|
||||
if mode == "legacy":
|
||||
ctx: Context = make_legacy_context(
|
||||
bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM
|
||||
)
|
||||
else:
|
||||
ctx = make_paged_context(bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM)
|
||||
|
||||
kv_full_cpu, ape_cpu = _make_inputs(
|
||||
seq_len, ctx.head_dim, seed=seq_len + prefix_len
|
||||
)
|
||||
pool = make_state_pool(ctx.num_pages, RATIO, ctx.head_dim)
|
||||
|
||||
if prefix_len > 0:
|
||||
seq_lens_cpu, extend_lens_cpu, _ = to_seq_extend([(prefix_len, prefix_len)])
|
||||
_run_prefill(
|
||||
ctx,
|
||||
pool,
|
||||
kv_full_cpu[:prefix_len].cuda(),
|
||||
ape_cpu.cuda(),
|
||||
seq_lens_cpu,
|
||||
extend_lens_cpu,
|
||||
)
|
||||
|
||||
final_out = None
|
||||
for k in range(RATIO):
|
||||
cur_seq_len = prefix_len + k + 1
|
||||
seq_lens_gpu = torch.tensor([cur_seq_len], dtype=torch.int64, device="cuda")
|
||||
kv_step = kv_full_cpu[prefix_len + k : prefix_len + k + 1].cuda()
|
||||
out = _run_decode(ctx, pool, kv_step, ape_cpu.cuda(), seq_lens_gpu)
|
||||
if cur_seq_len % RATIO == 0:
|
||||
final_out = out
|
||||
|
||||
P = seq_len - 1
|
||||
gt = _gt_compress(kv_full_cpu, ape_cpu, P=P, head_dim=ctx.head_dim)
|
||||
assert final_out is not None
|
||||
triton.testing.assert_close(final_out[0].cpu(), gt, atol=ATOL, rtol=RTOL)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["legacy", "paged"])
|
||||
@pytest.mark.parametrize("prefix_len", [128, 256])
|
||||
def test_prefill_then_extend(mode: str, prefix_len: int) -> None:
|
||||
"""Prefill once, then a second prefill that extends across one compress event.
|
||||
|
||||
First prefill ends at a 128-boundary so the second prefill starts fresh.
|
||||
"""
|
||||
extend_len = RATIO
|
||||
seq_len = prefix_len + extend_len
|
||||
|
||||
if mode == "legacy":
|
||||
ctx: Context = make_legacy_context(
|
||||
bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM
|
||||
)
|
||||
else:
|
||||
ctx = make_paged_context(bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM)
|
||||
|
||||
kv_full_cpu, ape_cpu = _make_inputs(seq_len, ctx.head_dim, seed=prefix_len)
|
||||
pool = make_state_pool(ctx.num_pages, RATIO, ctx.head_dim)
|
||||
|
||||
seq_lens_cpu, extend_lens_cpu, _ = to_seq_extend([(prefix_len, prefix_len)])
|
||||
_run_prefill(
|
||||
ctx,
|
||||
pool,
|
||||
kv_full_cpu[:prefix_len].cuda(),
|
||||
ape_cpu.cuda(),
|
||||
seq_lens_cpu,
|
||||
extend_lens_cpu,
|
||||
)
|
||||
|
||||
seq_lens_cpu, extend_lens_cpu, _ = to_seq_extend([(seq_len, extend_len)])
|
||||
out = _run_prefill(
|
||||
ctx,
|
||||
pool,
|
||||
kv_full_cpu[prefix_len:].cuda(),
|
||||
ape_cpu.cuda(),
|
||||
seq_lens_cpu,
|
||||
extend_lens_cpu,
|
||||
)
|
||||
|
||||
P = seq_len - 1
|
||||
gt = _gt_compress(kv_full_cpu, ape_cpu, P=P, head_dim=ctx.head_dim)
|
||||
# Single compress event in this extend; compact plan_id 0.
|
||||
triton.testing.assert_close(out[0].cpu(), gt, atol=ATOL, rtol=RTOL)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["legacy", "paged"])
|
||||
def test_prefill_multibatch(mode: str) -> None:
|
||||
"""Multi-batch prefill, each batch ending at a different chunk count."""
|
||||
seq_extend = [(128, 128), (256, 256), (384, 384)]
|
||||
bs = len(seq_extend)
|
||||
if mode == "legacy":
|
||||
ctx: Context = make_legacy_context(
|
||||
bs=bs, compress_ratio=RATIO, head_dim=HEAD_DIM
|
||||
)
|
||||
else:
|
||||
ctx = make_paged_context(bs=bs, compress_ratio=RATIO, head_dim=HEAD_DIM)
|
||||
|
||||
seq_lens_cpu, extend_lens_cpu, num_q = to_seq_extend(seq_extend)
|
||||
kv_in_cpu, ape_cpu = _make_inputs(num_q, ctx.head_dim, seed=99)
|
||||
pool = make_state_pool(ctx.num_pages, RATIO, ctx.head_dim)
|
||||
out = _run_prefill(
|
||||
ctx, pool, kv_in_cpu.cuda(), ape_cpu.cuda(), seq_lens_cpu, extend_lens_cpu
|
||||
)
|
||||
|
||||
# Compact: walk batches in order, then positions in order; matches the
|
||||
# CPU planner's emit order for plan_c.
|
||||
base = 0
|
||||
plan_id = 0
|
||||
for b, (seq, ext) in enumerate(seq_extend):
|
||||
for j in range(ext):
|
||||
P = j # prefix=0
|
||||
if (P + 1) % RATIO != 0:
|
||||
continue
|
||||
gt = _gt_compress(
|
||||
kv_in_cpu[base : base + ext],
|
||||
ape_cpu,
|
||||
P=P,
|
||||
head_dim=ctx.head_dim,
|
||||
)
|
||||
triton.testing.assert_close(
|
||||
out[plan_id].cpu(),
|
||||
gt,
|
||||
atol=ATOL,
|
||||
rtol=RTOL,
|
||||
)
|
||||
plan_id += 1
|
||||
base += ext
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -0,0 +1,337 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from typing import Tuple, Union
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.jit_kernel.benchmark.bench_activation import register_cuda_ci
|
||||
from sglang.jit_kernel.dsv4 import compress_forward
|
||||
from sglang.jit_kernel.tests.deepseek_v4.common import (
|
||||
LegacyContext,
|
||||
PagedContext,
|
||||
make_legacy_context,
|
||||
make_paged_context,
|
||||
make_state_pool,
|
||||
to_seq_extend,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=30, suite="stage-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=30, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
Context = Union[LegacyContext, PagedContext]
|
||||
|
||||
# c4 input row layout: | kv_overlap | kv | score_overlap | score |
|
||||
HEAD_DIM = 512
|
||||
RATIO = 4
|
||||
WINDOW = 8 # = 2 * RATIO (overlap + current)
|
||||
ATOL = 5e-3
|
||||
RTOL = 5e-3
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# fp64 ground truth (single compress event over a 8-token window).
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _gt_compress(
|
||||
kv_score_input_cpu: torch.Tensor, # [num_q, head_dim*4]
|
||||
ape_cpu: torch.Tensor, # [8, head_dim]
|
||||
P: int,
|
||||
head_dim: int,
|
||||
) -> torch.Tensor:
|
||||
"""fp64 reference for compress event at ragged position ``P``.
|
||||
|
||||
Tokens at positions [P-7..P-4] contribute their *overlap* halves, tokens
|
||||
at [P-3..P] contribute their *fresh* halves. Bias[0..3] for overlap,
|
||||
bias[4..7] for fresh. When P < 7, the overlap is masked (kv=0, score=-inf)
|
||||
so the softmax sees only the 4 fresh tokens.
|
||||
"""
|
||||
if P < 7:
|
||||
kv_ov = torch.zeros(4, head_dim, dtype=torch.float64)
|
||||
sc_ov = torch.full((4, head_dim), float("-inf"), dtype=torch.float64)
|
||||
else:
|
||||
kv_ov = kv_score_input_cpu[P - 7 : P - 3, :head_dim].double()
|
||||
sc_ov = kv_score_input_cpu[P - 7 : P - 3, 2 * head_dim : 3 * head_dim].double()
|
||||
kv_fr = kv_score_input_cpu[P - 3 : P + 1, head_dim : 2 * head_dim].double()
|
||||
sc_fr = kv_score_input_cpu[P - 3 : P + 1, 3 * head_dim :].double()
|
||||
kv = torch.cat([kv_ov, kv_fr], dim=0)
|
||||
sc = torch.cat([sc_ov, sc_fr], dim=0) + ape_cpu.double()
|
||||
return ((kv * sc.softmax(dim=0)).sum(dim=0)).float()
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Driver
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _run_prefill(
|
||||
ctx: Context,
|
||||
pool: torch.Tensor,
|
||||
kv_score_input: torch.Tensor,
|
||||
ape: torch.Tensor,
|
||||
seq_lens_cpu: torch.Tensor,
|
||||
extend_lens_cpu: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
num_q = int(extend_lens_cpu.sum().item())
|
||||
plan = ctx.make_prefill_plan(seq_lens_cpu, extend_lens_cpu, num_q)
|
||||
return compress_forward(
|
||||
pool,
|
||||
kv_score_input,
|
||||
ape,
|
||||
plan,
|
||||
head_dim=ctx.head_dim,
|
||||
compress_ratio=RATIO,
|
||||
)
|
||||
|
||||
|
||||
def _run_decode(
|
||||
ctx: Context,
|
||||
pool: torch.Tensor,
|
||||
kv_score_input: torch.Tensor,
|
||||
ape: torch.Tensor,
|
||||
seq_lens_gpu: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
plan = ctx.make_decode_plan(seq_lens_gpu)
|
||||
return compress_forward(
|
||||
pool,
|
||||
kv_score_input,
|
||||
ape,
|
||||
plan,
|
||||
head_dim=ctx.head_dim,
|
||||
compress_ratio=RATIO,
|
||||
)
|
||||
|
||||
|
||||
def _make_inputs(
|
||||
num_q: int, head_dim: int, seed: int
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
g = torch.Generator(device="cpu").manual_seed(seed)
|
||||
kv_score_input_cpu = torch.randn(
|
||||
num_q, head_dim * 4, generator=g, dtype=torch.float32
|
||||
)
|
||||
ape_cpu = torch.randn(WINDOW, head_dim, generator=g, dtype=torch.float32)
|
||||
return kv_score_input_cpu, ape_cpu
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Tests
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["legacy", "paged"])
|
||||
@pytest.mark.parametrize("seq_len", [4, 8, 32, 256, 1024])
|
||||
def test_prefill_no_context(mode: str, seq_len: int) -> None:
|
||||
"""Prefill once, no prefix. Every compress event must match fp64 GT."""
|
||||
if mode == "legacy":
|
||||
ctx: Context = make_legacy_context(
|
||||
bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM
|
||||
)
|
||||
else:
|
||||
ctx = make_paged_context(bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM)
|
||||
|
||||
seq_lens_cpu, extend_lens_cpu, num_q = to_seq_extend([(seq_len, seq_len)])
|
||||
kv_in_cpu, ape_cpu = _make_inputs(num_q, ctx.head_dim, seed=seq_len)
|
||||
|
||||
pool = make_state_pool(ctx.num_pages, RATIO, ctx.head_dim)
|
||||
out = _run_prefill(
|
||||
ctx, pool, kv_in_cpu.cuda(), ape_cpu.cuda(), seq_lens_cpu, extend_lens_cpu
|
||||
)
|
||||
|
||||
# Compact prefill output: row per compress plan, in CPU-planner order
|
||||
# (batch-major, position-ascending).
|
||||
for plan_id, P in enumerate(range(RATIO - 1, seq_len, RATIO)):
|
||||
gt = _gt_compress(kv_in_cpu, ape_cpu, P=P, head_dim=ctx.head_dim)
|
||||
triton.testing.assert_close(out[plan_id].cpu(), gt, atol=ATOL, rtol=RTOL)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["legacy", "paged"])
|
||||
@pytest.mark.parametrize("prefix_len", [4, 256])
|
||||
def test_prefill_then_decode(mode: str, prefix_len: int) -> None:
|
||||
"""Prefill once, then decode 4 more tokens through one compress boundary."""
|
||||
extend_decode = 4
|
||||
seq_len = prefix_len + extend_decode
|
||||
|
||||
if mode == "legacy":
|
||||
ctx: Context = make_legacy_context(
|
||||
bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM
|
||||
)
|
||||
else:
|
||||
ctx = make_paged_context(bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM)
|
||||
|
||||
kv_full_cpu, ape_cpu = _make_inputs(
|
||||
seq_len, ctx.head_dim, seed=seq_len + prefix_len
|
||||
)
|
||||
pool = make_state_pool(ctx.num_pages, RATIO, ctx.head_dim)
|
||||
|
||||
# Prefill the prefix.
|
||||
seq_lens_cpu, extend_lens_cpu, _ = to_seq_extend([(prefix_len, prefix_len)])
|
||||
_run_prefill(
|
||||
ctx,
|
||||
pool,
|
||||
kv_full_cpu[:prefix_len].cuda(),
|
||||
ape_cpu.cuda(),
|
||||
seq_lens_cpu,
|
||||
extend_lens_cpu,
|
||||
)
|
||||
|
||||
# Decode `extend_decode` tokens one at a time.
|
||||
final_out = None
|
||||
for k in range(extend_decode):
|
||||
cur_seq_len = prefix_len + k + 1
|
||||
seq_lens_gpu = torch.tensor([cur_seq_len], dtype=torch.int64, device="cuda")
|
||||
kv_step = kv_full_cpu[prefix_len + k : prefix_len + k + 1].cuda()
|
||||
out = _run_decode(ctx, pool, kv_step, ape_cpu.cuda(), seq_lens_gpu)
|
||||
if cur_seq_len % RATIO == 0:
|
||||
final_out = out
|
||||
|
||||
# Check the trailing compress: position P = seq_len - 1 = prefix + 3.
|
||||
P = seq_len - 1
|
||||
gt = _gt_compress(kv_full_cpu, ape_cpu, P=P, head_dim=ctx.head_dim)
|
||||
assert final_out is not None
|
||||
triton.testing.assert_close(final_out[0].cpu(), gt, atol=ATOL, rtol=RTOL)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["legacy", "paged"])
|
||||
@pytest.mark.parametrize("prefix_len", [256, 512, 768])
|
||||
def test_prefill_then_extend(mode: str, prefix_len: int) -> None:
|
||||
"""Prefill once, then prefill an extend that crosses one compress event.
|
||||
|
||||
The first prefill ends at a swa_page boundary (only relevant for paged),
|
||||
so the second prefill's overlap must be read out of the buffer.
|
||||
"""
|
||||
extend_len = 4
|
||||
|
||||
if mode == "legacy":
|
||||
ctx: Context = make_legacy_context(
|
||||
bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM
|
||||
)
|
||||
else:
|
||||
ctx = make_paged_context(bs=1, compress_ratio=RATIO, head_dim=HEAD_DIM)
|
||||
|
||||
seq_len = prefix_len + extend_len
|
||||
kv_full_cpu, ape_cpu = _make_inputs(seq_len, ctx.head_dim, seed=prefix_len)
|
||||
pool = make_state_pool(ctx.num_pages, RATIO, ctx.head_dim)
|
||||
|
||||
# First prefill: seq=prefix, ext=prefix.
|
||||
seq_lens_cpu, extend_lens_cpu, _ = to_seq_extend([(prefix_len, prefix_len)])
|
||||
_run_prefill(
|
||||
ctx,
|
||||
pool,
|
||||
kv_full_cpu[:prefix_len].cuda(),
|
||||
ape_cpu.cuda(),
|
||||
seq_lens_cpu,
|
||||
extend_lens_cpu,
|
||||
)
|
||||
|
||||
# Second prefill: seq=prefix+extend, ext=extend, prefix=prefix_len.
|
||||
seq_lens_cpu, extend_lens_cpu, num_q = to_seq_extend([(seq_len, extend_len)])
|
||||
out = _run_prefill(
|
||||
ctx,
|
||||
pool,
|
||||
kv_full_cpu[prefix_len:].cuda(),
|
||||
ape_cpu.cuda(),
|
||||
seq_lens_cpu,
|
||||
extend_lens_cpu,
|
||||
)
|
||||
|
||||
P = seq_len - 1
|
||||
gt = _gt_compress(kv_full_cpu, ape_cpu, P=P, head_dim=ctx.head_dim)
|
||||
# Single compress event in this extend; compact plan_id 0.
|
||||
triton.testing.assert_close(out[0].cpu(), gt, atol=ATOL, rtol=RTOL)
|
||||
|
||||
|
||||
def test_paged_buffer_intermediate() -> None:
|
||||
"""Paged-only: after a multi-page prefill, verify the trailing 4 tokens of
|
||||
every swa_page sit in the correct state-pool slots.
|
||||
|
||||
These slots are what radix-cache resume reads when prefix-matching from a
|
||||
swa_page boundary, so they MUST match the original token data.
|
||||
"""
|
||||
ctx = make_paged_context(
|
||||
bs=1,
|
||||
compress_ratio=RATIO,
|
||||
head_dim=HEAD_DIM,
|
||||
swa_page_size=256,
|
||||
ring_size=8,
|
||||
num_swa_pages_per_req=8,
|
||||
)
|
||||
seq_len = 1024 # 4 swa_pages
|
||||
seq_lens_cpu, extend_lens_cpu, num_q = to_seq_extend([(seq_len, seq_len)])
|
||||
kv_in_cpu, ape_cpu = _make_inputs(num_q, ctx.head_dim, seed=42)
|
||||
|
||||
pool = make_state_pool(ctx.num_pages, RATIO, ctx.head_dim)
|
||||
_run_prefill(
|
||||
ctx, pool, kv_in_cpu.cuda(), ape_cpu.cuda(), seq_lens_cpu, extend_lens_cpu
|
||||
)
|
||||
|
||||
pool_cpu = pool.cpu()
|
||||
# For each swa_page boundary, the trailing `RATIO` tokens must have been
|
||||
# written. The state slot for token at position p is
|
||||
# `state_loc(0, p) = (p // swa_page_size) * ring_size + p % ring_size`.
|
||||
for swa_page_end in range(ctx.swa_page_size, seq_len + 1, ctx.swa_page_size):
|
||||
for offset in range(RATIO):
|
||||
p = swa_page_end - RATIO + offset
|
||||
sl = ctx.state_loc(0, p)
|
||||
page_idx = sl // RATIO
|
||||
slot_idx = sl % RATIO
|
||||
actual = pool_cpu[page_idx, slot_idx]
|
||||
# Token-row layout: the c4 prefill write copies the full
|
||||
# head_dim*4 row from kv_input verbatim into the state pool.
|
||||
expected = kv_in_cpu[p]
|
||||
triton.testing.assert_close(
|
||||
actual,
|
||||
expected,
|
||||
atol=ATOL,
|
||||
rtol=RTOL,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["legacy", "paged"])
|
||||
def test_prefill_multibatch(mode: str) -> None:
|
||||
"""Multi-batch prefill, both modes."""
|
||||
seq_extend = [(8, 8), (256, 256), (260, 260), (1023, 1023)]
|
||||
bs = len(seq_extend)
|
||||
if mode == "legacy":
|
||||
ctx: Context = make_legacy_context(
|
||||
bs=bs, compress_ratio=RATIO, head_dim=HEAD_DIM
|
||||
)
|
||||
else:
|
||||
ctx = make_paged_context(bs=bs, compress_ratio=RATIO, head_dim=HEAD_DIM)
|
||||
|
||||
seq_lens_cpu, extend_lens_cpu, num_q = to_seq_extend(seq_extend)
|
||||
kv_in_cpu, ape_cpu = _make_inputs(num_q, ctx.head_dim, seed=99)
|
||||
pool = make_state_pool(ctx.num_pages, RATIO, ctx.head_dim)
|
||||
out = _run_prefill(
|
||||
ctx, pool, kv_in_cpu.cuda(), ape_cpu.cuda(), seq_lens_cpu, extend_lens_cpu
|
||||
)
|
||||
|
||||
# Compact: walk batches in order, then positions in order; matches the
|
||||
# CPU planner's emit order for plan_c.
|
||||
base = 0
|
||||
plan_id = 0
|
||||
for b, (seq, ext) in enumerate(seq_extend):
|
||||
for j in range(ext):
|
||||
P = j # prefix=0 here
|
||||
if (P + 1) % RATIO != 0:
|
||||
continue
|
||||
gt = _gt_compress(
|
||||
kv_in_cpu[base : base + ext],
|
||||
ape_cpu,
|
||||
P=P,
|
||||
head_dim=ctx.head_dim,
|
||||
)
|
||||
triton.testing.assert_close(
|
||||
out[plan_id].cpu(),
|
||||
gt,
|
||||
atol=ATOL,
|
||||
rtol=RTOL,
|
||||
)
|
||||
plan_id += 1
|
||||
base += ext
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -581,6 +581,7 @@ class Envs:
|
||||
SGLANG_OPT_USE_TILELANG_INDEXER = EnvBool(False)
|
||||
SGLANG_OPT_USE_JIT_INDEXER_METADATA = EnvBool(False)
|
||||
SGLANG_OPT_USE_ONLINE_COMPRESS = EnvBool(False)
|
||||
SGLANG_OPT_USE_COMPRESSOR_V2 = EnvBool(True)
|
||||
SGLANG_FP8_PAGED_MQA_LOGITS_TORCH = EnvBool(False)
|
||||
SGLANG_TOPK_TRANSFORM_512_TORCH = EnvBool(False)
|
||||
|
||||
@@ -611,7 +612,6 @@ class Envs:
|
||||
|
||||
# Cache / overlap
|
||||
SGLANG_OPT_USE_FUSED_STORE_CACHE = EnvBool(True)
|
||||
SGLANG_OPT_USE_OVERLAP_STORE_CACHE = EnvBool(True)
|
||||
SGLANG_OPT_USE_MULTI_STREAM_OVERLAP = EnvBool(True)
|
||||
|
||||
# CUDA graph
|
||||
|
||||
@@ -20,11 +20,21 @@ import torch.nn.functional as F
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
|
||||
from sglang.srt.layers.attention.dsv4.compressor import (
|
||||
CompressorBackendMixin,
|
||||
FusedCompressMetadata,
|
||||
create_paged_compressor_data,
|
||||
)
|
||||
|
||||
if envs.SGLANG_OPT_USE_COMPRESSOR_V2.get():
|
||||
# NOTE: should eventually be the only compressor backend
|
||||
from sglang.srt.layers.attention.dsv4.compressor_v2 import (
|
||||
CompressorBackendMixin,
|
||||
FusedCompressMetadata,
|
||||
create_paged_compressor_data,
|
||||
)
|
||||
else:
|
||||
from sglang.srt.layers.attention.dsv4.compressor import (
|
||||
CompressorBackendMixin,
|
||||
FusedCompressMetadata,
|
||||
create_paged_compressor_data,
|
||||
)
|
||||
|
||||
from sglang.srt.layers.attention.dsv4.indexer import C4IndexerBackendMixin
|
||||
from sglang.srt.layers.attention.dsv4.metadata import (
|
||||
PagedIndexerMetadata,
|
||||
|
||||
@@ -334,7 +334,8 @@ class Compressor(nn.Module):
|
||||
ape = torch.cat([ape[0], ape[1]], dim=0)
|
||||
self.ape.data.copy_(ape.view(self.ratio, -1))
|
||||
|
||||
def _get_state_pool(self, forward_batch: ForwardBatch) -> CompressStatePool:
|
||||
# NOTE: used by v2 compressor backend
|
||||
def get_state_pool(self, forward_batch: ForwardBatch) -> CompressStatePool:
|
||||
token_to_kv_pool = forward_batch.token_to_kv_pool
|
||||
assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool)
|
||||
if self.is_in_indexer:
|
||||
@@ -346,11 +347,8 @@ class Compressor(nn.Module):
|
||||
|
||||
return ret
|
||||
|
||||
def forward(self, x: torch.Tensor, forward_batch: ForwardBatch) -> torch.Tensor:
|
||||
if forward_batch.forward_mode.is_idle():
|
||||
assert x.shape[0] == 0
|
||||
return x.new_empty(0, self.head_dim)
|
||||
|
||||
# NOTE: used by v2 compressor backend
|
||||
def compute_kv_score(self, x: torch.Tensor, forward_batch: ForwardBatch):
|
||||
kv_score = linear_bf16_fp32(x, self.wkv_gate.weight)
|
||||
if nsa_use_prefill_cp(forward_batch):
|
||||
kv_score = cp_all_gather_rerange_output(
|
||||
@@ -359,11 +357,19 @@ class Compressor(nn.Module):
|
||||
forward_batch,
|
||||
torch.cuda.current_stream(),
|
||||
)
|
||||
return kv_score
|
||||
|
||||
def forward(self, x: torch.Tensor, forward_batch: ForwardBatch) -> torch.Tensor:
|
||||
if forward_batch.forward_mode.is_idle():
|
||||
assert x.shape[0] == 0
|
||||
return x.new_empty(0, self.head_dim)
|
||||
|
||||
kv_score = self.compute_kv_score(x, forward_batch)
|
||||
|
||||
backend = forward_batch.attn_backend
|
||||
if TYPE_CHECKING:
|
||||
assert isinstance(backend, DeepseekV4AttnBackend)
|
||||
kv_score_buffer = self._get_state_pool(forward_batch)
|
||||
kv_score_buffer = self.get_state_pool(forward_batch)
|
||||
kv_score_buffer = kv_score_buffer.kv_score_buffer.kv_score
|
||||
return backend.forward_compress(
|
||||
kv_score_buffer=kv_score_buffer,
|
||||
|
||||
@@ -0,0 +1,267 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, List, Literal, Optional, TypeAlias, Union, cast
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.dsv4 import (
|
||||
CompressorDecodePlan,
|
||||
CompressorPrefillPlan,
|
||||
compress_forward,
|
||||
compress_norm_rope_store,
|
||||
)
|
||||
from sglang.srt.environ import envs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.attention.deepseek_v4_backend import DSV4Metadata
|
||||
from sglang.srt.layers.attention.dsv4.compressor import Compressor
|
||||
from sglang.srt.layers.layernorm import RMSNorm
|
||||
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||
|
||||
|
||||
CompressMetadata: TypeAlias = Union[CompressorDecodePlan, CompressorPrefillPlan]
|
||||
# NOTE: alias for backward compatibility
|
||||
FusedCompressMetadata: TypeAlias = CompressMetadata
|
||||
|
||||
|
||||
def _use_online_compress(compress_ratio: int) -> bool:
|
||||
"""Online state-pool path is c128-only."""
|
||||
return compress_ratio == 128 and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get()
|
||||
|
||||
|
||||
class CompressorBackendMixin:
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.forward_metadata: DSV4Metadata
|
||||
|
||||
# NOTE: Will be overridden
|
||||
def _maybe_upgrade_forward_metadata(self): ...
|
||||
|
||||
def _get_paged_compress_metadata(self, compress_ratio: int) -> CompressMetadata:
|
||||
attr_name = f"c{compress_ratio}_compress_metadata"
|
||||
return getattr(self.forward_metadata, attr_name)
|
||||
|
||||
def _get_out_loc(self, compress_ratio: int) -> torch.Tensor:
|
||||
attr_name = f"c{compress_ratio}_out_loc"
|
||||
return getattr(self.forward_metadata.core_metadata, attr_name)
|
||||
|
||||
def _forward_compress_all_in_one(
|
||||
self,
|
||||
*,
|
||||
kv_score_buffer: torch.Tensor,
|
||||
kv_score_input: torch.Tensor,
|
||||
ape: torch.Tensor,
|
||||
head_dim: int,
|
||||
norm: RMSNorm,
|
||||
freqs_cis_cache: torch.Tensor,
|
||||
kv_cache: torch.Tensor,
|
||||
is_indexer: bool,
|
||||
rotate: bool,
|
||||
compress_ratio: int,
|
||||
page_size: int,
|
||||
) -> None:
|
||||
assert compress_ratio == 4 or compress_ratio == 128
|
||||
assert rotate == is_indexer == (head_dim == 128)
|
||||
|
||||
plan = self._get_paged_compress_metadata(compress_ratio)
|
||||
is_online = _use_online_compress(compress_ratio)
|
||||
if is_online:
|
||||
kv_score_buffer = kv_score_buffer.view(-1, 1, head_dim * 3)
|
||||
else:
|
||||
coff = 2 if is_overlap_compress(compress_ratio) else 1
|
||||
last_dim = 2 * head_dim * coff
|
||||
assert kv_score_buffer.shape[-1] == last_dim
|
||||
kv_score_buffer = kv_score_buffer.view(-1, compress_ratio, last_dim)
|
||||
kv_compressed = compress_forward(
|
||||
kv_score_buffer=kv_score_buffer,
|
||||
kv_score_input=kv_score_input,
|
||||
ape=ape.view(-1, head_dim),
|
||||
plan=plan,
|
||||
compress_ratio=compress_ratio,
|
||||
head_dim=head_dim,
|
||||
is_online=is_online,
|
||||
)
|
||||
# NOTE: we use some hack here...
|
||||
compress_norm_rope_store(
|
||||
kv_compressed,
|
||||
plan,
|
||||
norm_weight=norm.weight,
|
||||
norm_eps=norm.variance_epsilon,
|
||||
freq_cis=freqs_cis_cache,
|
||||
out_loc=self._get_out_loc(compress_ratio),
|
||||
kvcache=kv_cache,
|
||||
page_size=page_size,
|
||||
)
|
||||
|
||||
def forward_unified(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
layer_id: int,
|
||||
compressor: Compressor,
|
||||
) -> None:
|
||||
if forward_batch.forward_mode.is_idle():
|
||||
return
|
||||
|
||||
self._maybe_upgrade_forward_metadata()
|
||||
token_to_kv_pool = forward_batch.token_to_kv_pool
|
||||
token_to_kv_pool = cast("DeepSeekV4TokenToKVPool", token_to_kv_pool)
|
||||
kv_score_input = compressor.compute_kv_score(x, forward_batch)
|
||||
state_pool = compressor.get_state_pool(forward_batch)
|
||||
if compressor.is_in_indexer:
|
||||
kv_cache = token_to_kv_pool.get_index_k_with_scale_buffer(layer_id)
|
||||
page_size = token_to_kv_pool.get_index_k_page_size()
|
||||
else:
|
||||
kv_cache = token_to_kv_pool.get_extra_key_buffer(layer_id)
|
||||
page_size = token_to_kv_pool.get_extra_key_page_size(layer_id)
|
||||
self._forward_compress_all_in_one(
|
||||
kv_score_buffer=state_pool.kv_score_buffer.kv_score,
|
||||
kv_score_input=kv_score_input,
|
||||
ape=compressor.ape,
|
||||
head_dim=compressor.head_dim,
|
||||
norm=compressor.norm,
|
||||
freqs_cis_cache=compressor.freqs_cis,
|
||||
kv_cache=kv_cache.view(dtype=torch.uint8),
|
||||
is_indexer=compressor.is_in_indexer,
|
||||
rotate=compressor.rotate,
|
||||
compress_ratio=compressor.ratio,
|
||||
page_size=page_size,
|
||||
)
|
||||
|
||||
# NOTE: alias for backward compatibility
|
||||
forward_indexer_compressor = forward_unified
|
||||
forward_core_compressor = forward_unified
|
||||
|
||||
|
||||
def is_overlap_compress(compress_ratio: int) -> bool:
|
||||
return compress_ratio == 4
|
||||
|
||||
|
||||
def create_paged_compressor_data(
|
||||
compress_ratio: Literal[4, 128],
|
||||
*,
|
||||
is_prefill: bool,
|
||||
token_to_kv_pool: DeepSeekV4TokenToKVPool,
|
||||
req_to_token: torch.Tensor,
|
||||
req_pool_indices: torch.Tensor,
|
||||
seq_lens: torch.Tensor,
|
||||
extend_lens: Optional[torch.Tensor] = None,
|
||||
seq_lens_cpu: Optional[List[int]] = None,
|
||||
extend_lens_cpu: Optional[List[int]] = None,
|
||||
use_prefill_cuda_graph: bool = False,
|
||||
num_q_tokens: Optional[int] = None,
|
||||
) -> CompressMetadata:
|
||||
"""Build the paged compress metadata (= the plan).
|
||||
|
||||
State-pool slot translation is done inside the C++ planner; the
|
||||
Python side just hands the relevant tensors over.
|
||||
"""
|
||||
if _use_online_compress(compress_ratio):
|
||||
return _create_online_paged_compressor_data(
|
||||
is_prefill=is_prefill,
|
||||
token_to_kv_pool=token_to_kv_pool,
|
||||
req_to_token=req_to_token,
|
||||
req_pool_indices=req_pool_indices,
|
||||
seq_lens=seq_lens,
|
||||
extend_lens=extend_lens,
|
||||
seq_lens_cpu=seq_lens_cpu,
|
||||
extend_lens_cpu=extend_lens_cpu,
|
||||
use_prefill_cuda_graph=use_prefill_cuda_graph,
|
||||
num_q_tokens=num_q_tokens,
|
||||
)
|
||||
|
||||
swa_page_size = token_to_kv_pool.swa_page_size
|
||||
ring_size = token_to_kv_pool.get_ring_size(compress_ratio=compress_ratio)
|
||||
# NOTE: This is actually a proxy, which encounter some bug with tvm-ffi.
|
||||
# As a workaround, we use `.detach()` to get the real tensor.
|
||||
full_to_swa = token_to_kv_pool.full_to_swa_index_mapping.detach()
|
||||
req_pool_indices_i64 = req_pool_indices.to(torch.int64)
|
||||
|
||||
if is_prefill:
|
||||
assert extend_lens is not None
|
||||
if seq_lens_cpu is not None:
|
||||
assert extend_lens_cpu is not None
|
||||
seq_lens_planner = torch.tensor(seq_lens_cpu, dtype=torch.int64)
|
||||
extend_lens_planner = torch.tensor(extend_lens_cpu, dtype=torch.int64)
|
||||
num_q_tokens = sum(extend_lens_cpu)
|
||||
else:
|
||||
assert num_q_tokens is not None
|
||||
seq_lens_planner = seq_lens.to(torch.int64)
|
||||
extend_lens_planner = extend_lens.to(torch.int64)
|
||||
|
||||
return CompressorPrefillPlan.generate(
|
||||
compress_ratio=compress_ratio,
|
||||
req_pool_indices=req_pool_indices_i64,
|
||||
seq_lens=seq_lens_planner,
|
||||
extend_lens=extend_lens_planner,
|
||||
req_to_token=req_to_token,
|
||||
full_to_swa=full_to_swa,
|
||||
swa_page_size=swa_page_size,
|
||||
ring_size=ring_size,
|
||||
num_q_tokens=num_q_tokens,
|
||||
use_cuda_graph=use_prefill_cuda_graph,
|
||||
)
|
||||
else:
|
||||
return CompressorDecodePlan.generate(
|
||||
compress_ratio=compress_ratio,
|
||||
req_pool_indices=req_pool_indices_i64,
|
||||
req_to_token=req_to_token,
|
||||
full_to_swa=full_to_swa,
|
||||
seq_lens=seq_lens.to(torch.int64),
|
||||
swa_page_size=swa_page_size,
|
||||
ring_size=ring_size,
|
||||
)
|
||||
|
||||
|
||||
def _create_online_paged_compressor_data(
|
||||
*,
|
||||
is_prefill: bool,
|
||||
token_to_kv_pool: DeepSeekV4TokenToKVPool,
|
||||
req_to_token: torch.Tensor,
|
||||
req_pool_indices: torch.Tensor,
|
||||
seq_lens: torch.Tensor,
|
||||
extend_lens: Optional[torch.Tensor],
|
||||
seq_lens_cpu: Optional[List[int]],
|
||||
extend_lens_cpu: Optional[List[int]],
|
||||
use_prefill_cuda_graph: bool,
|
||||
num_q_tokens: Optional[int],
|
||||
) -> 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)
|
||||
|
||||
if is_prefill:
|
||||
# Sync-on-entry: catch IMA from a prior layer / kernel BEFORE we touch
|
||||
# anything in this builder, so blame doesn't land on us spuriously.
|
||||
assert extend_lens is not None
|
||||
if seq_lens_cpu is not None:
|
||||
assert extend_lens_cpu is not None
|
||||
seq_lens_planner = torch.tensor(seq_lens_cpu, dtype=torch.int64)
|
||||
extend_lens_planner = torch.tensor(extend_lens_cpu, dtype=torch.int64)
|
||||
num_q_tokens_planner = sum(extend_lens_cpu)
|
||||
else:
|
||||
assert num_q_tokens is not None
|
||||
seq_lens_planner = seq_lens.to(torch.int64)
|
||||
extend_lens_planner = extend_lens.to(torch.int64)
|
||||
num_q_tokens_planner = num_q_tokens
|
||||
|
||||
return CompressorPrefillPlan.generate_online(
|
||||
seq_lens=seq_lens_planner,
|
||||
extend_lens=extend_lens_planner,
|
||||
req_pool_indices=req_pool_indices,
|
||||
req_to_token=req_to_token,
|
||||
full_to_swa=full_to_swa,
|
||||
num_q_tokens=int(num_q_tokens_planner),
|
||||
swa_page_size=swa_page_size,
|
||||
)
|
||||
else:
|
||||
return CompressorDecodePlan.generate_online(
|
||||
seq_lens=seq_lens.to(torch.int64),
|
||||
req_pool_indices=req_pool_indices,
|
||||
req_to_token=req_to_token,
|
||||
full_to_swa=full_to_swa,
|
||||
swa_page_size=swa_page_size,
|
||||
)
|
||||
@@ -9,7 +9,7 @@ import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.jit_kernel.deepseek_v4 import (
|
||||
fused_rope,
|
||||
fused_q_indexer_rope_hadamard_quant,
|
||||
topk_transform_512,
|
||||
topk_transform_512_v2,
|
||||
)
|
||||
@@ -17,8 +17,6 @@ from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.attention.dsv4.compressor import Compressor
|
||||
from sglang.srt.layers.attention.dsv4.metadata import PagedIndexerMetadata
|
||||
from sglang.srt.layers.attention.nsa.nsa_indexer import rotate_activation
|
||||
from sglang.srt.layers.attention.nsa.triton_kernel import act_quant
|
||||
from sglang.srt.layers.linear import ReplicatedLinear
|
||||
from sglang.srt.state_capturer.indexer_topk import get_global_indexer_capturer
|
||||
from sglang.srt.utils import add_prefix, is_hip
|
||||
@@ -267,21 +265,21 @@ class C4IndexerBackendMixin:
|
||||
layer_id=c4_indexer.layer_id,
|
||||
)
|
||||
|
||||
# The weight projection is small and fast; compute it on its own
|
||||
# stream, then have the Q stream wait on it before launching the big
|
||||
# fused Q kernel (which folds rope + hadamard + fp8 quant + the
|
||||
# weight*weight_scale*q_scale step into one pass).
|
||||
with torch.cuda.stream(stream_weights):
|
||||
weights = c4_indexer.compute_weights(x, skip_scale=True)
|
||||
weights_ready = stream_weights.record_event()
|
||||
|
||||
with torch.cuda.stream(stream_q):
|
||||
if q_lora_ready is not None:
|
||||
stream_q.wait_event(q_lora_ready)
|
||||
q = c4_indexer.compute_q(q_lora, positions=positions)
|
||||
q_fp8, q_scale = act_quant(q)
|
||||
q_scale_ready = stream_q.record_event()
|
||||
|
||||
with torch.cuda.stream(stream_weights):
|
||||
weights = c4_indexer.compute_weights(x, skip_scale=True)
|
||||
stream_weights.wait_event(q_scale_ready)
|
||||
weights = fused_scale(weights, c4_indexer.weight_scale, q_scale)
|
||||
stream_q.wait_event(weights_ready)
|
||||
q_fp8, weights = c4_indexer.compute_q(q_lora, positions, weights)
|
||||
|
||||
current_stream.wait_stream(stream_q)
|
||||
current_stream.wait_stream(stream_weights)
|
||||
|
||||
return q_fp8, weights, c4_indexer_kv_cache
|
||||
|
||||
def _forward_prepare_normal(
|
||||
@@ -296,10 +294,8 @@ class C4IndexerBackendMixin:
|
||||
if TYPE_CHECKING:
|
||||
assert isinstance(self, CompressorBackendMixin)
|
||||
|
||||
q = c4_indexer.compute_q(q_lora, positions=positions)
|
||||
q_fp8, q_scale = act_quant(q)
|
||||
weights = c4_indexer.compute_weights(x, skip_scale=True)
|
||||
weights = fused_scale(weights, c4_indexer.weight_scale, q_scale)
|
||||
q_fp8, weights = c4_indexer.compute_q(q_lora, positions, weights)
|
||||
self.forward_indexer_compressor(
|
||||
x=x,
|
||||
forward_batch=forward_batch,
|
||||
@@ -523,17 +519,17 @@ class C4Indexer(nn.Module):
|
||||
self.weight_scale: float = self.softmax_scale * self.n_heads**-0.5
|
||||
self.alt_streams = alt_streams
|
||||
|
||||
def compute_q(self, q_lora: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
|
||||
def compute_q(
|
||||
self,
|
||||
q_lora: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
q, _ = self.wq_b(q_lora)
|
||||
q = q.view(-1, self.n_local_heads, self.head_dim)
|
||||
fused_rope(
|
||||
q[..., -self.rope_head_dim :],
|
||||
None,
|
||||
self.freqs_cis,
|
||||
positions=positions,
|
||||
return fused_q_indexer_rope_hadamard_quant(
|
||||
q, weight, self.weight_scale, self.freqs_cis, positions
|
||||
)
|
||||
q = rotate_activation(q)
|
||||
return q
|
||||
|
||||
def compute_weights(self, x: torch.Tensor, skip_scale=False) -> torch.Tensor:
|
||||
out, _ = self.weights_proj(x)
|
||||
|
||||
@@ -6,7 +6,7 @@ from typing import List, Literal, NamedTuple, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.deepseek_v4 import fused_store_cache
|
||||
from sglang.jit_kernel.deepseek_v4 import fused_k_norm_rope_flashmla, fused_store_cache
|
||||
from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.attention.dsv4 import (
|
||||
@@ -630,7 +630,12 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
|
||||
) -> None:
|
||||
self.swa_kv_pool.set_key_buffer(layer_id, loc, cache_nope_fp8_rope_bf16_pack)
|
||||
|
||||
def get_extra_key_buffer(self, layer_id: int) -> torch.Tensor | None:
|
||||
def get_extra_key_page_size(self, layer_id: int) -> int:
|
||||
_, _, compress_kv_pool = self.layer_mapping[layer_id]
|
||||
assert compress_kv_pool is not None
|
||||
return compress_kv_pool.page_size
|
||||
|
||||
def get_extra_key_buffer(self, layer_id: int) -> torch.Tensor:
|
||||
_, compress_layer_id, compress_kv_pool = self.layer_mapping[layer_id]
|
||||
assert compress_kv_pool is not None
|
||||
return compress_kv_pool.get_key_buffer(compress_layer_id)
|
||||
@@ -647,6 +652,9 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
|
||||
compress_layer_id, loc, cache_nope_fp8_rope_bf16_pack
|
||||
)
|
||||
|
||||
def get_index_k_page_size(self) -> int:
|
||||
return self.c4_indexer_kv_pool.page_size
|
||||
|
||||
def get_index_k_with_scale_buffer(self, layer_id: int) -> torch.Tensor:
|
||||
compress_ratio, compress_layer_id, _ = self.layer_mapping[layer_id]
|
||||
assert compress_ratio == 4, f"only c4 has indexer, got {compress_ratio = }"
|
||||
@@ -717,6 +725,33 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
|
||||
swa_loc = self.translate_loc_from_full_to_swa(raw_loc)
|
||||
return self.swa_kv_pool.set_key_buffer_fused(layer_id, swa_loc, cache_k)
|
||||
|
||||
def set_swa_key_buffer_radix_fused_norm_rope(
|
||||
self,
|
||||
layer_id: int,
|
||||
raw_loc: torch.Tensor,
|
||||
kv: torch.Tensor,
|
||||
kv_weight: torch.Tensor,
|
||||
eps: float,
|
||||
freqs_cis: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
) -> None:
|
||||
if self._should_cache_swa:
|
||||
if layer_id == self.start_layer or self.cached_loc is None:
|
||||
self.cached_loc = self.translate_loc_from_full_to_swa(raw_loc)
|
||||
swa_loc = self.cached_loc
|
||||
else:
|
||||
swa_loc = self.translate_loc_from_full_to_swa(raw_loc)
|
||||
fused_k_norm_rope_flashmla(
|
||||
kv=kv,
|
||||
kv_weight=kv_weight,
|
||||
eps=eps,
|
||||
freqs_cis=freqs_cis,
|
||||
positions=positions,
|
||||
out_loc=swa_loc,
|
||||
kvcache=self.swa_kv_pool.kv_buffer[layer_id],
|
||||
page_size=self.swa_kv_pool.page_size,
|
||||
)
|
||||
|
||||
def set_extra_key_buffer_fused(
|
||||
self,
|
||||
layer_id: int,
|
||||
|
||||
@@ -11,7 +11,11 @@ import triton
|
||||
import triton.language as tl
|
||||
|
||||
import sglang.srt.models.deepseek_v2 as deepseek_v2
|
||||
from sglang.jit_kernel.deepseek_v4 import fused_rope, rmsnorm_self
|
||||
from sglang.jit_kernel.deepseek_v4 import (
|
||||
fused_norm_rope_inplace,
|
||||
fused_q_norm_rope,
|
||||
fused_rope_inplace,
|
||||
)
|
||||
from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config
|
||||
from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size
|
||||
from sglang.srt.environ import envs
|
||||
@@ -25,7 +29,6 @@ from sglang.srt.layers.attention.nsa.utils import (
|
||||
nsa_use_prefill_cp,
|
||||
)
|
||||
from sglang.srt.layers.communicator import get_attn_tp_context
|
||||
from sglang.srt.layers.deepseek_v4_rope import apply_rotary_emb_triton
|
||||
from sglang.srt.layers.dp_attention import (
|
||||
_DpGatheredBufferWrapper,
|
||||
attn_tp_all_gather,
|
||||
@@ -82,6 +85,7 @@ if TYPE_CHECKING:
|
||||
DeepseekV4AttnBackend,
|
||||
)
|
||||
from sglang.srt.layers.quantization import QuantizationConfig
|
||||
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
|
||||
from sglang.srt.model_executor.forward_batch_info import (
|
||||
ForwardBatch,
|
||||
PPProxyTensors,
|
||||
@@ -315,8 +319,9 @@ class MQALayer(nn.Module):
|
||||
prefix=add_prefix("attn_mqa", prefix),
|
||||
)
|
||||
|
||||
self.overlap_store_cache = envs.SGLANG_OPT_USE_OVERLAP_STORE_CACHE.get()
|
||||
self.use_jit_norm = envs.SGLANG_OPT_USE_JIT_NORM.get()
|
||||
# KV cache write is always fused into the K kernel
|
||||
# (`_compute_kv_to_cache`), so the legacy "overlap store cache" flag
|
||||
# has no effect here -- the fused path is on by default.
|
||||
|
||||
def _compute_q_a(
|
||||
self,
|
||||
@@ -327,52 +332,69 @@ class MQALayer(nn.Module):
|
||||
q = qkv_a[..., : self.q_lora_rank]
|
||||
else:
|
||||
q, _ = self.wq_a(x)
|
||||
q = self.q_norm(q)
|
||||
q_lora = q
|
||||
return q_lora
|
||||
return self.q_norm(q)
|
||||
|
||||
def _compute_q_b(
|
||||
self,
|
||||
q: torch.Tensor,
|
||||
positions: Optional[torch.Tensor] = None,
|
||||
positions: torch.Tensor,
|
||||
q_out: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
q, _ = self.wq_b(q)
|
||||
q = q.view(-1, self.n_local_heads, self.head_dim)
|
||||
if self.use_jit_norm:
|
||||
q = rmsnorm_self(q, self.eps)
|
||||
else:
|
||||
q = rms_normalize_triton(q, self.eps)
|
||||
if positions is not None:
|
||||
fused_rope(
|
||||
q[..., -self.qk_rope_head_dim :],
|
||||
None,
|
||||
self.freqs_cis,
|
||||
positions=positions,
|
||||
)
|
||||
else:
|
||||
apply_rotary_emb_triton(q[..., -self.qk_rope_head_dim :], self.freqs_cis)
|
||||
return q
|
||||
if q_out is None:
|
||||
q_out = torch.empty_like(q)
|
||||
# Fused warp-per-(token, head) rmsnorm-self + RoPE + write to q_out.
|
||||
fused_q_norm_rope(q, q_out, self.eps, self.freqs_cis, positions)
|
||||
return q_out
|
||||
|
||||
def _compute_kv(
|
||||
def _compute_kv_to_cache(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
positions: Optional[torch.Tensor] = None,
|
||||
positions: torch.Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
qkv_a: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
) -> None:
|
||||
"""Fused: rmsnorm + RoPE + write directly to FlashMLA paged cache.
|
||||
|
||||
Replaces the bf16-kv-intermediate path. Used everywhere except the NSA
|
||||
prefill-CP case (which needs bf16 kv for the cross-rank all-gather).
|
||||
"""
|
||||
if qkv_a is not None:
|
||||
kv = qkv_a[..., self.q_lora_rank :]
|
||||
else:
|
||||
kv, _ = self.wkv(x)
|
||||
kv = self.kv_norm(kv)
|
||||
if positions is not None:
|
||||
fused_rope(
|
||||
kv[..., -self.qk_rope_head_dim :].unsqueeze(1),
|
||||
None,
|
||||
self.freqs_cis,
|
||||
positions=positions,
|
||||
)
|
||||
token_to_kv_pool = forward_batch.token_to_kv_pool
|
||||
if TYPE_CHECKING:
|
||||
assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool)
|
||||
token_to_kv_pool.set_swa_key_buffer_radix_fused_norm_rope(
|
||||
layer_id=self.layer_id,
|
||||
raw_loc=forward_batch.out_cache_loc,
|
||||
kv=kv,
|
||||
kv_weight=self.kv_norm.weight.data,
|
||||
eps=self.eps,
|
||||
freqs_cis=self.freqs_cis,
|
||||
positions=positions,
|
||||
)
|
||||
|
||||
def _compute_kv_bf16(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
qkv_a: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Bf16-kv path used by the NSA prefill-CP case (needs all-gather)."""
|
||||
if qkv_a is not None:
|
||||
kv = qkv_a[..., self.q_lora_rank :]
|
||||
else:
|
||||
apply_rotary_emb_triton(kv[..., -self.qk_rope_head_dim :], self.freqs_cis)
|
||||
kv, _ = self.wkv(x)
|
||||
fused_norm_rope_inplace(
|
||||
kv,
|
||||
self.kv_norm.weight.data,
|
||||
self.eps,
|
||||
self.freqs_cis,
|
||||
positions,
|
||||
)
|
||||
return kv
|
||||
|
||||
def _forward_prepare_multi_stream(
|
||||
@@ -382,7 +404,7 @@ class MQALayer(nn.Module):
|
||||
forward_batch: ForwardBatch,
|
||||
attn_backend: DeepseekV4AttnBackend,
|
||||
q_out: Optional[torch.Tensor] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
) -> torch.Tensor:
|
||||
assert self.alt_streams is not None
|
||||
assert len(self.alt_streams) >= 3
|
||||
|
||||
@@ -417,13 +439,8 @@ class MQALayer(nn.Module):
|
||||
with torch.cuda.stream(stream_kv):
|
||||
if qkv_a_ready is not None:
|
||||
stream_kv.wait_event(qkv_a_ready)
|
||||
kv = self._compute_kv(x, positions, qkv_a=qkv_a)
|
||||
if self.overlap_store_cache:
|
||||
attn_backend.store_cache(
|
||||
layer_id=self.layer_id,
|
||||
swa_k=kv,
|
||||
forward_batch=forward_batch,
|
||||
)
|
||||
# Fused norm + rope + cache write -- no bf16 KV intermediate.
|
||||
self._compute_kv_to_cache(x, positions, forward_batch, qkv_a=qkv_a)
|
||||
|
||||
del qkv_a
|
||||
|
||||
@@ -433,15 +450,12 @@ class MQALayer(nn.Module):
|
||||
x, forward_batch, self.layer_id, self.compressor
|
||||
)
|
||||
|
||||
q = self._compute_q_b(q_lora, positions)
|
||||
if q_out is not None:
|
||||
q_out.copy_(q)
|
||||
|
||||
q = self._compute_q_b(q_lora, positions, q_out)
|
||||
current_stream.wait_stream(stream_kv)
|
||||
current_stream.wait_stream(stream_compressor)
|
||||
current_stream.wait_stream(stream_indexer)
|
||||
|
||||
return q, kv
|
||||
return q
|
||||
|
||||
def _forward_prepare(
|
||||
self,
|
||||
@@ -450,47 +464,38 @@ class MQALayer(nn.Module):
|
||||
forward_batch: ForwardBatch,
|
||||
attn_backend: DeepseekV4AttnBackend,
|
||||
q_out: Optional[torch.Tensor] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
|
||||
if self.fuse_wqa_wkv:
|
||||
qkv_a, _ = self.wqkv_a(x)
|
||||
q = qkv_a[..., : self.q_lora_rank]
|
||||
kv = qkv_a[..., self.q_lora_rank :]
|
||||
del qkv_a
|
||||
q_lora = qkv_a[..., : self.q_lora_rank]
|
||||
else:
|
||||
kv, _ = self.wkv(x)
|
||||
q, _ = self.wq_a(x)
|
||||
q = self.q_norm(q)
|
||||
q_lora = q
|
||||
q, _ = self.wq_b(q)
|
||||
q = q.view(-1, self.n_local_heads, self.head_dim)
|
||||
if self.use_jit_norm:
|
||||
q = rmsnorm_self(q, self.eps)
|
||||
else:
|
||||
q = rms_normalize_triton(q, self.eps)
|
||||
q_lora, _ = self.wq_a(x)
|
||||
qkv_a = None
|
||||
q_lora = self.q_norm(q_lora)
|
||||
q = self._compute_q_b(q_lora, positions, q_out)
|
||||
|
||||
kv = self.kv_norm(kv)
|
||||
|
||||
fused_rope(
|
||||
q[..., -self.qk_rope_head_dim :],
|
||||
kv[..., -self.qk_rope_head_dim :].unsqueeze(1),
|
||||
self.freqs_cis,
|
||||
positions=positions,
|
||||
)
|
||||
|
||||
if self.nsa_enable_prefill_cp and nsa_use_prefill_cp(forward_batch):
|
||||
use_cp = self.nsa_enable_prefill_cp and nsa_use_prefill_cp(forward_batch)
|
||||
kv: Optional[torch.Tensor]
|
||||
if use_cp:
|
||||
# NSA CP: keep bf16 kv around for the cross-rank all-gather, then
|
||||
# write to the FlashMLA cache after gather.
|
||||
kv = self._compute_kv_bf16(x, positions, qkv_a=qkv_a)
|
||||
kv = cp_all_gather_rerange_output(
|
||||
kv.contiguous(),
|
||||
self.cp_size,
|
||||
forward_batch,
|
||||
torch.cuda.current_stream(),
|
||||
)
|
||||
|
||||
if self.overlap_store_cache:
|
||||
attn_backend.store_cache(
|
||||
layer_id=self.layer_id,
|
||||
swa_k=kv,
|
||||
forward_batch=forward_batch,
|
||||
)
|
||||
else:
|
||||
self._compute_kv_to_cache(x, positions, forward_batch, qkv_a=qkv_a)
|
||||
kv = None
|
||||
|
||||
del qkv_a
|
||||
|
||||
if self.indexer is not None:
|
||||
self.indexer(x=x, q_lora=q_lora, forward_batch=forward_batch)
|
||||
@@ -502,8 +507,6 @@ class MQALayer(nn.Module):
|
||||
self.compressor,
|
||||
)
|
||||
|
||||
if q_out is not None:
|
||||
q_out.copy_(q)
|
||||
return q, kv
|
||||
|
||||
def forward(
|
||||
@@ -538,26 +541,34 @@ class MQALayer(nn.Module):
|
||||
q_out = q_padded[:, tp_slice, :]
|
||||
|
||||
if enable_multi_stream:
|
||||
q, kv = self._forward_prepare_multi_stream(
|
||||
# Multi-stream path always fuses cache write into the K kernel,
|
||||
# so the bf16 KV intermediate is gone.
|
||||
q = self._forward_prepare_multi_stream(
|
||||
x, positions, forward_batch, attn_backend, q_out
|
||||
)
|
||||
kv = None
|
||||
else:
|
||||
q, kv = self._forward_prepare(
|
||||
x, positions, forward_batch, attn_backend, q_out
|
||||
)
|
||||
|
||||
# The cache write is always fused / already done by _forward_prepare* --
|
||||
# tell the backend to skip its own store_cache. When `kv is None`
|
||||
# (no NSA-CP), pass `q` as a sentinel for the `k is v` assert; the
|
||||
# attention path doesn't read it once `save_kv_cache=False`.
|
||||
attn_k = kv if kv is not None else q
|
||||
o = attn_backend.forward(
|
||||
q=q_padded if q_padded is not None else q,
|
||||
k=kv,
|
||||
v=kv,
|
||||
k=attn_k,
|
||||
v=attn_k,
|
||||
layer=self.attn_mqa,
|
||||
forward_batch=forward_batch,
|
||||
compress_ratio=self.compress_ratio,
|
||||
attn_sink=self.attn_sink,
|
||||
save_kv_cache=not self.overlap_store_cache,
|
||||
save_kv_cache=False,
|
||||
)
|
||||
o = o[:, tp_slice, :]
|
||||
fused_rope(
|
||||
fused_rope_inplace(
|
||||
o[..., -self.qk_rope_head_dim :],
|
||||
None,
|
||||
self.freqs_cis,
|
||||
|
||||
Reference in New Issue
Block a user