diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c128_online_v2.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/c128_online_v2.cuh new file mode 100644 index 000000000..71e600dc3 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/c128_online_v2.cuh @@ -0,0 +1,875 @@ +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +#include +#include +#include +#include +#include + +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 +__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; + const auto gmem = tile::Memory::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(); + + 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(params.kv_score_buffer); + const auto kv_score_input = static_cast(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(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; + +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; + 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(local_val_max); + const auto rescale = expf(local_val_max - global_val_max); + const auto global_exp_sum = warp::reduce_sum(local_exp_sum * rescale); + const auto final_scale = rescale / global_exp_sum; + const auto global_product = warp::reduce_sum(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 +__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(); + + const auto plan = plan_ptr[global_pid]; + if (plan.is_invalid()) [[unlikely]] + return; + + const auto kv_score_buffer = static_cast(params.kv_score_buffer); + const auto kv_score_input = static_cast(params.kv_score_input); + const auto kv_compressed_output = static_cast(params.kv_compressed_output); + const auto score_bias_base = static_cast(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(plan.ragged_id) - static_cast(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(); + + // --- 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(buf_store + 0 * kHeadDim)[lane_id] = out_max_vec; + reinterpret_cast(buf_store + 1 * kHeadDim)[lane_id] = out_sum_vec; + reinterpret_cast(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(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 +struct FlashCompress128OnlineKernel { + static constexpr auto decode_kernel = flash_c128_online_decode_v2; + template + static constexpr auto prefill_kernel = flash_c128_online_prefill_v2; + 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(); + + TensorMatcher({-1, 1, kHeadDim * 3}) // kv score buffer (max, sum, kv) + .with_dtype() + .with_device(device_) + .verify(kv_score_buffer); + TensorMatcher({B, kHeadDim * 2}) // kv score input + .with_dtype() + .with_device(device_) + .verify(kv_score_input); + TensorMatcher({B, kHeadDim}) // kv compressed output (sparse by batch_id) + .with_dtype() + .with_device(device_) + .verify(kv_compressed_output); + TensorMatcher({128, kHeadDim}) // ape + .with_dtype() + .with_device(device_) + .verify(ape); + + const auto plan_d = compress::verify_plan_d(plan_d_, B, device_); + const auto batch_size = static_cast(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(); + + TensorMatcher({-1, 1, kHeadDim * 3}) // kv score buffer + .with_dtype() + .with_device(device_) + .verify(kv_score_buffer); + TensorMatcher({N, kHeadDim * 2}) // kv score input (ragged) + .with_dtype() + .with_device(device_) + .verify(kv_score_input); + TensorMatcher({C, kHeadDim}) // kv compressed output (compact, by plan_c index) + .with_dtype() + .with_device(device_) + .verify(kv_compressed_output); + TensorMatcher({128, kHeadDim}) // ape + .with_dtype() + .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(N.unwrap()); + const auto num_c = static_cast(C.unwrap()); + const auto num_w = static_cast(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, params); + } + if (const auto num_w_blocks = num_w * kNumSplit) { + LaunchKernel(num_w_blocks, kPrefillBlockSize, device) // + .enable_pdl(kUsePDL)(prefill_kernel, 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(params.seq_lens[idx]); + const auto rid = params.req_pool_indices[idx]; + const int32_t chunk_start = static_cast((seq_len - 1u) / 128u * 128u); + const int32_t full_loc = params.req_to_token[rid * params.stride_r2t + chunk_start]; + const int32_t swa_loc = static_cast(params.full_to_swa[full_loc]); + const int32_t slot = swa_loc / params.swa_page_size; + params.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(); + + auto seq_dtype = SymbolicDType{}; + TensorMatcher({B}) // + .with_dtype(seq_dtype) + .with_device(device_) + .verify(seq_lens); + TensorMatcher({B}) // + .with_dtype() + .with_device(device_) + .verify(req_pool_indices); + TensorMatcher({-1, -1}) // + .with_dtype() + .with_device(device_) + .verify(req_to_token); + TensorMatcher({-1}) // + .with_dtype() + .with_device(device_) + .verify(full_to_swa); + TensorMatcher({B, sizeof(DecodePlan)}) // + .with_dtype() + .with_device(device_) + .verify(plan_d_dev_); + RuntimeCheck(swa_page_size > 0); + + const auto batch_size = static_cast(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(plan_d_dev_.data_ptr()), + .seq_lens = static_cast(seq_lens.data_ptr()), + .req_pool_indices = static_cast(req_pool_indices.data_ptr()), + .req_to_token = static_cast(req_to_token.data_ptr()), + .full_to_swa = static_cast(full_to_swa.data_ptr()), + .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 _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(p.seq_lens[i]); + const uint32_t extend_len = static_cast(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(last_ragged), + .buffer_len = static_cast(seg_len), + .read_page_0 = static_cast(i), // batch_id placeholder + .read_page_1 = -1, // unused, kept so MSB layout is stable + }; + 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{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(plan.seq_len - 1u); + const int32_t chunk_start = (position / 128) * 128; + const int32_t full_loc = params.req_to_token[rid * params.stride_r2t + chunk_start]; + const int32_t swa_loc = static_cast(params.full_to_swa[full_loc]); + plan.read_page_0 = swa_loc / params.swa_page_size; + *plan_ptr = plan; +} + +using OnlinePrefillPlan = tvm::ffi::Tuple; + +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(); + device_.set_options(); + + TensorMatcher({B}) // + .with_dtype() + .with_device(cpu) + .verify(seq_lens) + .verify(extend_lens); + TensorMatcher({B}) // + .with_dtype() + .with_device(device_) + .verify(req_pool_indices); + TensorMatcher({-1, -1}) // + .with_dtype() + .with_device(device_) + .verify(req_to_token); + TensorMatcher({-1}) // + .with_dtype() + .with_device(device_) + .verify(full_to_swa); + TensorMatcher({N, sizeof(CompressPlan)}) // + .with_dtype() + .with_device(cpu) + .verify(plan_c_pin) + .verify(plan_w_pin); + TensorMatcher({N, sizeof(CompressPlan)}) // + .with_dtype() + .with_device(device_) + .verify(plan_c_dev_) + .verify(plan_w_dev_); + + const auto stage0_params = OnlinePrefillStage0Params{ + .plan_c = static_cast(plan_c_pin.data_ptr()), + .plan_w = static_cast(plan_w_pin.data_ptr()), + .seq_lens = static_cast(seq_lens.data_ptr()), + .extend_lens = static_cast(extend_lens.data_ptr()), + .batch_size = static_cast(B.unwrap()), + .num_q_tokens = static_cast(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(plan_c_pin.data_ptr()); + auto* pw = static_cast(plan_w_pin.data_ptr()); + const auto bytes = static_cast(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(plan_c_pin.data_ptr()); + auto* pw = static_cast(plan_w_pin.data_ptr()); + const auto end_c = static_cast(num_c) * sizeof(CompressPlan); + const auto end_w = static_cast(num_w) * sizeof(CompressPlan); + const auto pin_bytes = static_cast(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(plan_c_dev_.data_ptr()); + auto* const plan_w_dev_ptr = static_cast(plan_w_dev_.data_ptr()); + + if (const auto total = num_c + num_w) { + 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(req_pool_indices.data_ptr()), + .req_to_token = static_cast(req_to_token.data_ptr()), + .full_to_swa = static_cast(full_to_swa.data_ptr()), + .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 diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh index 498f2601e..e625919e1 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh @@ -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 #include @@ -8,7 +21,7 @@ #include #include -#include +#include #include #include @@ -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 -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 +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; - const auto element_size = head_dim * 2; - const auto gmem = tile::Memory{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 +template 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; @@ -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(score[j][i]) + cast(bias[j][i]); + } + } #pragma unroll - for (int32_t j = 0; j < kElementsPerWarp; ++j) { - score_fp32[j] = cast(score[j][i]) + cast(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(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(); + #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 +SGL_DEVICE void c128_write_decode(InFloat* kv_buf, const InFloat* kv_src) { + using namespace device; + + using Storage = AlignedVector; + const auto gmem = tile::Memory::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 C128_KERNEL void flash_c128_decode(const __grid_constant__ Compress128DecodeParams params) { using namespace device; + using Trait = C128Trait; - 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(params.kv_input) + split_offset; + const auto kv_output = static_cast(params.kv_output) + split_offset; + const auto kv_buffer = static_cast(params.kv_buffer) + split_offset; + const auto score_bias = static_cast(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(_kv_score_buffer); - const auto kv_buf = kv_score_buffer + index * (kElementSize * 128) + split_offset; - - // kv input - const auto kv_score_input = static_cast(_kv_score_input); - const auto kv_src = kv_score_input + global_bid * kElementSize + split_offset; - - // kv output - const auto kv_compressed_output = static_cast(_kv_compressed_output); - const auto kv_out = kv_compressed_output + global_bid * kHeadDim + split_offset; - - // score bias (ape) - const auto score_bias = static_cast(_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(); - - /// 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(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(kv_buf, kv_src, kv_out, score_bias, 128); } - - PDLTriggerSecondary(); } // compress kernel -template +template C128_KERNEL void flash_c128_prefill(const __grid_constant__ Compress128PrefillParams params) { using namespace device; + using Trait = C128Trait; - 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(params.kv_input) + split_offset; + const auto kv_output = static_cast(params.kv_output) + split_offset; + const auto kv_buffer = static_cast(params.kv_buffer) + split_offset; + const auto score_bias = static_cast(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(); + c128_forward(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 +WRITE_KERNEL void write_c128_prefill(const __grid_constant__ Compress128PrefillParams params) { + using namespace device; + using Trait = C128Trait; + using StorageIn = AlignedVector; - 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(params.kv_input) + split_offset; + const auto kv_buffer = static_cast(params.kv_buffer) + split_offset; + if (plan.is_invalid()) return; - // kv input - const auto kv_score_input = static_cast(_kv_score_input); - const auto kv_src = kv_score_input + ragged_id * kElementSize + split_offset; - - // kv output - const auto kv_compressed_output = static_cast(_kv_compressed_output); - const auto kv_out = kv_compressed_output + ragged_id * kHeadDim + split_offset; - - // score bias (ape) - const auto score_bias = static_cast(_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(index) >= _num_indices) [[unlikely]] - return; - - // kv score - const auto kv_score_buffer = static_cast(_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::warp(); PDLWaitPrimary(); - - // 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(); +#pragma unroll + for (int32_t i = 0; i < 2; ++i) { + gmem.store(kv_buf, data[i], i); + } } template struct FlashCompress128Kernel { static constexpr auto decode_kernel = flash_c128_decode; - template - static constexpr auto prefill_kernel = flash_c128_prefill; - static constexpr auto prefill_c_kernel = prefill_kernel; - static constexpr auto prefill_w_kernel = prefill_kernel; + static constexpr auto prefill_c_kernel = flash_c128_prefill; + static constexpr auto prefill_w_kernel = write_c128_prefill; 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; 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 /* 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(); + auto N = SymbolicSize{"batch_size"}; + auto device_ = SymbolicDevice{}; + device_.set_options(); - TensorMatcher({-1, 128, kHeadDim * 2}) // kv score + TensorMatcher({-1, 128, Trait::kElementSize}) // kv score .with_dtype() - .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() - .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() - .with_device(device) - .verify(kv_compressed_output); + .with_device(device_) + .verify(kv_output); TensorMatcher({128, kHeadDim}) // ape .with_dtype() - .with_device(device) + .with_device(device_) .verify(ape); - TensorMatcher({B}) // indices - .with_dtype() - .with_device(device) - .verify(indices); - TensorMatcher({B}) // seq lens - .with_dtype() - .with_device(device) - .verify(seq_lens); - const auto batch_size = static_cast(B.unwrap()); + const auto plan_d = compress::verify_plan_d(plan_d_, N, device_); + const auto batch_size = static_cast(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(indices.data_ptr()), - .seq_lens = static_cast(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 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(); - TensorMatcher({K, 128, kHeadDim * 2}) // kv score + TensorMatcher({-1, 128, Trait::kElementSize}) // kv score .with_dtype() .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() .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() .with_device(device_) - .verify(kv_compressed_output); + .verify(kv_output); TensorMatcher({128, kHeadDim}) // ape .with_dtype() .with_device(device_) .verify(ape); - TensorMatcher({B}) // indices - .with_dtype() - .with_device(device_) - .verify(indices); - TensorMatcher({X, compress::kPrefillPlanDim}) // compress plan - .with_dtype() - .with_device(device_) - .verify(compress_plan); - TensorMatcher({Y, compress::kPrefillPlanDim}) // write plan - .with_dtype() - .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() - .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(B.unwrap()); const auto num_q_tokens = static_cast(N.unwrap()); - const auto num_c = static_cast(X.unwrap()); - const auto num_w = static_cast(Y.unwrap()); - const auto num_indices = static_cast(K.unwrap()); + const auto num_c = static_cast(C.unwrap()); + const auto num_w = static_cast(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(indices.data_ptr()), - .load_indices = static_cast(load_indices.data_ptr()), - .compress_plan = static_cast(compress_plan.data_ptr()), - .write_plan = static_cast(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); } diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh new file mode 100644 index 000000000..b970e3847 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh @@ -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 +#include + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +#include +#include + +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 +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 +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; + /// NOTE: load one tile_dim (< head_dim) at at time + const auto gmem_in = tile::Memory::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(0.0f)); + score[i].fill(cast(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; + const auto gmem_out = tile::Memory::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(score[j][i]) + cast(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(kv[j][i]) * exp_score; + sum_exp_value += exp_score; + } + + result[i] = cast(sum_product / sum_exp_value); + } + + // overlap the store with the next iteration's load + PDLTriggerSecondary(); + gmem_out.store(kv_out, result); +} + +template +SGL_DEVICE void c4_write_decode(InFloat* kv_buf, const InFloat* kv_src) { + using namespace device; + + using StorageIn = AlignedVector; + const auto gmem = tile::Memory::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 +C4_KERNEL void flash_c4_decode(const __grid_constant__ Compress4DecodeParams params) { + using namespace device; + using Trait = C4Trait; + + 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(params.kv_input) + split_offset; + const auto kv_output = static_cast(params.kv_output) + split_offset; + const auto kv_buffer = static_cast(params.kv_buffer) + split_offset; + const auto score_bias = static_cast(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(); + c4_write_decode(kv_dst, kv_src); + if (plan.seq_len % 4 == 0) { + const auto need_overlap = plan.seq_len > 4; + c4_forward(kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, 8); + } +} + +template +C4_KERNEL void flash_c4_prefill(const __grid_constant__ Compress4PrefillParams params) { + using namespace device; + using Trait = C4Trait; + + 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(params.kv_input) + split_offset; + const auto kv_output = static_cast(params.kv_output) + split_offset; + const auto kv_buffer = static_cast(params.kv_buffer) + split_offset; + const auto score_bias = static_cast(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(); + c4_forward(kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, plan.buffer_len); +} + +template +WRITE_KERNEL void write_c4_prefill(const __grid_constant__ Compress4PrefillParams params) { + using namespace device; + using Trait = C4Trait; + using StorageIn = AlignedVector; + + 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(params.kv_input) + split_offset; + const auto kv_buffer = static_cast(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::warp(); + + PDLWaitPrimary(); + StorageIn data[4]; +#pragma unroll + for (int32_t i = 0; i < 4; ++i) { + data[i] = gmem.load(kv_src, i); + } + PDLTriggerSecondary(); +#pragma unroll + for (int32_t i = 0; i < 4; ++i) { + gmem.store(kv_buf, data[i], i); + } +} + +template +struct FlashCompress4Kernel { + static constexpr auto decode_kernel = flash_c4_decode; + static constexpr auto prefill_c_kernel = flash_c4_prefill; + static constexpr auto prefill_w_kernel = write_c4_prefill; + 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; + + 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(); + + TensorMatcher({-1, 4, Trait::kElementSize}) // kv score + .with_dtype() + .with_device(device_) + .verify(kv_buffer); + TensorMatcher({N, Trait::kElementSize}) // kv score input + .with_dtype() + .with_device(device_) + .verify(kv_input); + TensorMatcher({N, kHeadDim}) // kv compressed output + .with_dtype() + .with_device(device_) + .verify(kv_output); + TensorMatcher({8, kHeadDim}) // ape + .with_dtype() + .with_device(device_) + .verify(ape); + + const auto plan_d = compress::verify_plan_d(plan_d_, N, device_); + const auto batch_size = static_cast(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(); + + TensorMatcher({-1, 4, Trait::kElementSize}) // kv score + .with_dtype() + .with_device(device_) + .verify(kv_buffer); + TensorMatcher({N, Trait::kElementSize}) // kv score input (ragged) + .with_dtype() + .with_device(device_) + .verify(kv_input); + TensorMatcher({C, kHeadDim}) // kv compressed output (compact) + .with_dtype() + .with_device(device_) + .verify(kv_output); + TensorMatcher({8, kHeadDim}) // ape + .with_dtype() + .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(N.unwrap()); + const auto num_c = static_cast(C.unwrap()); + const auto num_w = static_cast(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 diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh new file mode 100644 index 000000000..4518dcffa --- /dev/null +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh @@ -0,0 +1,827 @@ +#include +#include +#include + +#include +#include + +#include + +#include +#include + +#include +#include + +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(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(plan.ragged_id), static_cast(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::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(params.seq_lens_ptr[tx]); + extend_len = static_cast(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(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(j); + const uint32_t ragged_id = k; + + if ((position + 1) % cr == 0) { + const int32_t buffer_len = window_size - min(static_cast(j) + 1, window_size); + const uint32_t out_idx = atomicAdd(&counter_c, 1u); + params.plan_c[out_idx] = { + .seq_len = static_cast(position + 1), + .ragged_id = static_cast(ragged_id), + .buffer_len = static_cast(buffer_len), + .read_page_0 = -1, + .read_page_1 = static_cast(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(tx); j < el; j += static_cast(block_size)) { + const int32_t position = pl + j; + const uint32_t ragged_id = base_e + static_cast(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(position + 1), + .ragged_id = static_cast(ragged_id), + .buffer_len = static_cast(buffer_len), + .read_page_0 = -1, + .read_page_1 = static_cast(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(batch_id), position + 1); + } + } + base_e += static_cast(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(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(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(params.seq_ptr[idx]); + const auto position_1 = static_cast(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(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(params.rid_ptr[batch_id]); + // `seq_len` is ratio-aligned for compress events + const auto position_1 = static_cast(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(params.rid_ptr[batch_id]); + // `write_loc` carries (position + 1) at this stage; may not be ratio-aligned + const auto position = static_cast(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(params.rid_ptr[idx]); + const auto seq_len = static_cast(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(seq_len), + .write_loc = write_loc, + .read_page_0 = read_page_0, + .read_page_1 = read_page_1, + }; +} + +using PrefillPlan = tvm::ffi::Tuple; + +/** + * \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(); + device_.set_options(); + + TensorMatcher({B}) // + .with_dtype() + .with_device(device_) + .verify(req_pool_indices); + TensorMatcher({-1, -1}) // + .with_dtype() + .with_device(device_) + .verify(req_to_token); + TensorMatcher({-1}) // + .with_dtype() + .with_device(device_) + .verify(full_to_swa); + TensorMatcher({B}) // + .with_dtype() + .with_device(cpu_or_gpu) + .verify(seq_lens) + .verify(extend_lens); + TensorMatcher({-1}) // + .with_dtype() + .with_device() + .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(seq_lens.data_ptr()); + const auto ext_ptr = static_cast(extend_lens.data_ptr()); + const auto rid_ptr = static_cast(req_pool_indices.data_ptr()); + const auto r2t_ptr = static_cast(req_to_token.data_ptr()); + const auto f2s_ptr = static_cast(full_to_swa.data_ptr()); + + const auto batch_size = static_cast(B.unwrap()); + constexpr auto kMaxTokens = static_cast(std::numeric_limits::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(C.data_ptr()), + .plan_w = static_cast(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(C.data_ptr()), + .plan_w = static_cast(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(pin_buffer.numel()) * sizeof(uint8_t); + RuntimeCheck(pin_buffer_bytes >= num_q_tokens * (sizeof(PlanC) + sizeof(PlanW))); + const auto plan_c_ptr = reinterpret_cast(pin_buffer.data_ptr()); + const auto plan_w_ptr = reinterpret_cast(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(position + 1), + .ragged_id = static_cast(ragged_id), + .buffer_len = static_cast(buffer_len), + // to be filled by kernel + .read_page_0 = -1, + .read_page_1 = static_cast(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(C.data_ptr()), + .plan_w = static_cast(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(); + + TensorMatcher({B}) // + .with_dtype() + .with_device(device_) + .verify(req_pool_indices); + TensorMatcher({-1, -1}) // + .with_dtype() + .with_device(device_) + .verify(req_to_token); + TensorMatcher({-1}) // + .with_dtype() + .with_device(device_) + .verify(full_to_swa); + TensorMatcher({B}) // + .with_dtype() + .with_device(device_) + .verify(seq_lens); + + const auto batch_size = static_cast(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(D.data_ptr()), + .rid_ptr = static_cast(req_pool_indices.data_ptr()), + .r2t_ptr = static_cast(req_to_token.data_ptr()), + .f2s_ptr = static_cast(full_to_swa.data_ptr()), + .seq_ptr = static_cast(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(); + + TensorMatcher({B}) // + .with_dtype() + .with_device(device_) + .verify(req_pool_indices); + TensorMatcher({B}) // + .with_dtype() + .with_device() + .verify(seq_lens) + .verify(extend_lens); + TensorMatcher({-1}) // + .with_dtype() + .with_device() + .verify(pin_buffer); + + const auto pin_buffer_bytes = static_cast(pin_buffer.numel()) * sizeof(uint8_t); + RuntimeCheck(pin_buffer_bytes >= num_q_tokens * (sizeof(PlanC) + sizeof(PlanW))); + const auto plan_c_ptr = reinterpret_cast(pin_buffer.data_ptr()); + const auto plan_w_ptr = reinterpret_cast(plan_c_ptr + num_q_tokens); + + const bool is_overlap = (compress_ratio == 4); + const auto seq_ptr = static_cast(seq_lens.data_ptr()); + const auto ext_ptr = static_cast(extend_lens.data_ptr()); + const auto rid_ptr = static_cast(req_pool_indices.data_ptr()); + + const auto window_size = compress_ratio * (is_overlap ? 2 : 1); + const auto batch_size = static_cast(B.unwrap()); + constexpr auto kMaxTokens = static_cast(std::numeric_limits::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(position + 1), + .ragged_id = static_cast(ragged_id), + .buffer_len = static_cast(buffer_len), + // to be filled by kernel + .read_page_0 = -1, + .read_page_1 = static_cast(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(C.data_ptr()), + .plan_w = static_cast(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(); + + TensorMatcher({B}) // + .with_dtype() + .with_device(device_) + .verify(req_pool_indices); + TensorMatcher({B}) // + .with_dtype() + .with_device(device_) + .verify(seq_lens); + RuntimeCheck(compress_ratio == 4 || compress_ratio == 128); + + const auto batch_size = static_cast(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(D.data_ptr()), + .rid_ptr = static_cast(req_pool_indices.data_ptr()), + .seq_ptr = static_cast(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 diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh new file mode 100644 index 000000000..653ff8750 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh @@ -0,0 +1,419 @@ +#include +#include + +#include +#include +#include +#include +#include + +#include +#include + +#include + +#include + +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 +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; + using Float4 = AlignedVector; + + 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(params.input) + work_id * kHeadDim; + int32_t position; + int32_t out_loc; + if constexpr (kMode == CompressExtend) { + const auto plan = static_cast(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(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, "Unsupported Mode"); + } + const auto freqs_cis = params.freqs_cis + position * kRopeDim; + + PDLWaitPrimary(); + 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(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(input_vec[i]); + const auto fp32_weight = cast(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(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; + 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(); + 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(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 +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; + using Float2 = AlignedVector; + + 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(params.input) + work_id * kHeadDim; + int32_t position; + int32_t out_loc; + if constexpr (kMode == CompressExtend) { + const auto plan = static_cast(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(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, "Unsupported Mode"); + } + const auto freqs_cis = params.freqs_cis + position * kRopeDim; + + PDLWaitPrimary(); + 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(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(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(input_vec[i]); + const auto fp32_weight = cast(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(); + + // 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(fp32x2_t{data[0], data[1]}); + const auto rope_ptr = value_ptr + 448; + reinterpret_cast(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(value_ptr)[tx] = result; + // All lanes in this warp produce the same scale byte; let lane 0 publish. + if (lane_id == 0) static_cast(scale_ptr)[warp_id] = scale_ue8m0; + } +} + +template +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 + static constexpr auto select_kernel() { + if constexpr (kIsIndexer) { + return fused_norm_rope_indexer; + } else { + return fused_norm_rope_flashmla; + } + } + + 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(is_decode); + + auto N = SymbolicSize{"num_tokens"}; + auto device_ = SymbolicDevice{}; + device_.set_options(); + + TensorMatcher({N, kHeadDim}) // input + .with_dtype() + .with_device(device_) + .verify(input); + TensorMatcher({kHeadDim}) // weight + .with_dtype() + .with_device(device_) + .verify(weight); + TensorMatcher({-1, kRopeDim}) // freqs_cis + .with_dtype() + .with_device(device_) + .verify(freqs_cis); + TensorMatcher({-1}) // out_loc + .with_dtype() + .with_device(device_) + .verify(out_loc); + TensorMatcher({-1, -1}) // cache + .with_strides({kPageBytes, 1}) + .with_dtype() + .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(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(freqs_cis.data_ptr()), + .out_loc = static_cast(out_loc.data_ptr()), + .kvcache = static_cast(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() : select_kernel(); + LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params); + } +}; + +} // namespace diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh new file mode 100644 index 000000000..f8ee07ce5 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh @@ -0,0 +1,629 @@ +#include +#include + +#include +#include +#include +#include +#include +#include + +#include + +#include + +#include +#include + +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 +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; + + 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(params.q_input) + batch_id * params.q_input_stride_batch + head_id * kHeadDim; + const auto output_ptr = + static_cast(params.q_output) + batch_id * params.q_output_stride_batch + head_id * kHeadDim; + const auto position = static_cast(static_cast(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{lane_id, kWarpThreads}; + + PDLWaitPrimary(); + + // part 1: rmsnorm-self (no weight). + const auto gmem = tile::Memory{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(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(input_vec[i][j]); + input_vec[i][j] = cast(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(); + + // part 2: RoPE on all 32 lanes -- one (real, imag) bf16x2 pair per lane. + using DType2 = packed_t; + const auto mem_elem = tile::Memory{lane_id, kWarpThreads}; + const auto elem = mem_elem.load(s_rope[warp_id]); + const auto [x_real, x_imag] = cast(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(rotated)); +} + +template +struct FusedQNormRopeKernel { + template + static constexpr auto kernel = fused_q_norm_rope; + + 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(); + + TensorMatcher({B, H, kHeadDim}) // + .with_strides({-1, kHeadDim, 1}) + .with_dtype() + .with_device(device_) + .verify(q_input); + TensorMatcher({B, H, kHeadDim}) // + .with_strides({-1, kHeadDim, 1}) + .with_dtype() + .with_device(device_) + .verify(q_output); + TensorMatcher({-1, kRopeDim}) // + .with_dtype() + .with_device(device_) + .verify(freqs_cis); + auto pos_dtype = SymbolicDType{}; + TensorMatcher({B}) // + .with_dtype(pos_dtype) + .with_device(device_) + .verify(positions); + + const auto batch_size = static_cast(B.unwrap()); + const auto num_q_heads = static_cast(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(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; + const auto k_int64 = kernel; + const auto k = pos_dtype.is_type() ? 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 +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; + using Float2 = AlignedVector; + + 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(params.kv) + work_id * params.kv_stride_batch; + const auto position = static_cast(static_cast(params.positions)[work_id]); + const auto out_loc = params.out_loc[work_id]; + const auto freqs_cis = params.freqs_cis + position * kRopeDim; + + PDLWaitPrimary(); + 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(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(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(input_vec[i]); + const auto w = cast(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(); + + // 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(fp32x2_t{data[0], data[1]}); + const auto rope_ptr = value_ptr + 448; + reinterpret_cast(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(value_ptr)[tx] = result; + if (lane_id == 0) static_cast(scale_ptr)[warp_id] = scale_ue8m0; + } +} + +template +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 + static constexpr auto kernel = fused_k_norm_rope_flashmla; + + 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(); + + TensorMatcher({B, kHeadDim}) // + .with_strides({-1, 1}) + .with_dtype() + .with_device(device_) + .verify(kv); + TensorMatcher({kHeadDim}) // + .with_dtype() + .with_device(device_) + .verify(kv_weight); + TensorMatcher({-1, kRopeDim}) // + .with_dtype() + .with_device(device_) + .verify(freqs_cis); + auto pos_dtype = SymbolicDType{}; + TensorMatcher({B}) // + .with_dtype(pos_dtype) + .with_device(device_) + .verify(positions); + TensorMatcher({B}) // + .with_dtype() + .with_device(device_) + .verify(out_loc); + TensorMatcher({-1, -1}) // + .with_strides({kPageBytes, 1}) + .with_dtype() + .with_device(device_) + .verify(kvcache); + + const auto batch_size = static_cast(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(freqs_cis.data_ptr()), + .positions = positions.data_ptr(), + .out_loc = static_cast(out_loc.data_ptr()), + .kvcache = static_cast(kvcache.data_ptr()), + .kv_stride_batch = kv.stride(0), + .batch_size = batch_size, + .eps = eps, + }; + const auto k_int32 = kernel; + const auto k_int64 = kernel; + const auto k = pos_dtype.is_type() ? 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 +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; + using Float4 = AlignedVector; + using OutStorage = AlignedVector; // 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(params.q_input) + work_id * kHeadDim; + const auto position = static_cast(static_cast(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(); + Float4 data, freq; + const auto weight_val = cast(static_cast(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(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(); + + // 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(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(params.q_fp8) + work_id * kHeadDim; + result.store(out_row, lane_id); + params.weights_out[work_id] = weight_val * params.weight_scale * scale; + } +} + +template +struct FusedQIndexerRopeHadamardQuantKernel { + template + static constexpr auto kernel = fused_q_indexer_rope_hadamard_quant; + + 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(); + + // 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() + .with_device(device_) + .verify(q_input); + TensorMatcher({B, H, kHeadDim}) // + .with_strides({-1, kHeadDim, 1}) + .with_dtype() + .with_device(device_) + .verify(q_fp8); + TensorMatcher({B, H}) // + .with_dtype() + .with_device(device_) + .verify(weight); + TensorMatcher({B, H, 1}) // + .with_dtype() + .with_device(device_) + .verify(weights_out); + TensorMatcher({-1, kRopeDim}) // + .with_dtype() + .with_device(device_) + .verify(freqs_cis); + auto pos_dtype = SymbolicDType{}; + TensorMatcher({B}) // + .with_dtype(pos_dtype) + .with_device(device_) + .verify(positions); + + const auto batch_size = static_cast(B.unwrap()); + const auto num_heads = static_cast(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(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(weights_out.data_ptr()), + .weight_scale = static_cast(weight_scale), + .freqs_cis = static_cast(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; + const auto k_int64 = kernel; + const auto k = pos_dtype.is_type() ? k_int32 : k_int64; + LaunchKernel(num_blocks, kFusedQBlockSize, device_.unwrap()) // + .enable_pdl(kUsePDL)(k, params); + } +}; + +} // namespace diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/rmsnorm.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/rmsnorm.cuh index f9407ec84..a69642c5a 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/rmsnorm.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/rmsnorm.cuh @@ -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(); // 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(); TensorMatcher({N, H, D}) // input - .with_strides({Dh, Dn, 1}) + .with_strides({-1, -1, 1}) .with_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() .with_device(device) .verify(output); const auto batch_size = static_cast(N.unwrap()); const auto num_head = static_cast(H.unwrap()); - const auto stride_head_bytes = static_cast(Dn.unwrap() * sizeof(DType)); - const auto stride_batch_bytes = static_cast(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(input.stride(0) * sizeof(DType)), + .stride_head_bytes_0 = static_cast(input.stride(1) * sizeof(DType)), + .stride_batch_bytes_1 = static_cast(output.stride(0) * sizeof(DType)), + .stride_head_bytes_1 = static_cast(output.stride(1) * sizeof(DType)), .batch_size = batch_size, .num_head = num_head, .eps = eps, diff --git a/python/sglang/jit_kernel/deepseek_v4.py b/python/sglang/jit_kernel/deepseek_v4.py index 72192b533..6f9d772ad 100644 --- a/python/sglang/jit_kernel/deepseek_v4.py +++ b/python/sglang/jit_kernel/deepseek_v4.py @@ -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 diff --git a/python/sglang/jit_kernel/dsv4/__init__.py b/python/sglang/jit_kernel/dsv4/__init__.py new file mode 100644 index 000000000..0dcf60827 --- /dev/null +++ b/python/sglang/jit_kernel/dsv4/__init__.py @@ -0,0 +1,10 @@ +from .compress import * +from .utils import make_name + +__all__ = [ + "CompressorDecodePlan", + "CompressorPrefillPlan", + "compress_forward", + "compress_norm_rope_store", + "make_name", +] diff --git a/python/sglang/jit_kernel/dsv4/compress.py b/python/sglang/jit_kernel/dsv4/compress.py new file mode 100644 index 000000000..7885a82f6 --- /dev/null +++ b/python/sglang/jit_kernel/dsv4/compress.py @@ -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, + ) diff --git a/python/sglang/jit_kernel/dsv4/utils.py b/python/sglang/jit_kernel/dsv4/utils.py new file mode 100644 index 000000000..8085074f6 --- /dev/null +++ b/python/sglang/jit_kernel/dsv4/utils.py @@ -0,0 +1,2 @@ +def make_name(name: str) -> str: + return f"dpsk_v4_{name}" diff --git a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/compress_v2.cuh b/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/compress_v2.cuh new file mode 100644 index 000000000..3e87127c5 --- /dev/null +++ b/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/compress_v2.cuh @@ -0,0 +1,99 @@ +#pragma once + +#include +#include + +#include + +#include +#include + +#include + +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() + .with_device(device) + .verify(t); + return static_cast(t.data_ptr()); +} + +inline auto verify_plan_c(tvm::ffi::TensorView t, SymbolicSize& N, SymbolicDevice& device) -> const CompressPlan* { + TensorMatcher({N, sizeof(CompressPlan)}) // + .with_dtype() + .with_device(device) + .verify(t); + return static_cast(t.data_ptr()); +} + +inline auto verify_plan_w(tvm::ffi::TensorView t, SymbolicSize& N, SymbolicDevice& device) -> const WritePlan* { + TensorMatcher({N, sizeof(WritePlan)}) // + .with_dtype() + .with_device(device) + .verify(t); + return static_cast(t.data_ptr()); +} + +} // namespace host::compress diff --git a/python/sglang/jit_kernel/tests/deepseek_v4/__init__.py b/python/sglang/jit_kernel/tests/deepseek_v4/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/jit_kernel/tests/deepseek_v4/common.py b/python/sglang/jit_kernel/tests/deepseek_v4/common.py new file mode 100644 index 000000000..7288caa44 --- /dev/null +++ b/python/sglang/jit_kernel/tests/deepseek_v4/common.py @@ -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 diff --git a/python/sglang/jit_kernel/tests/deepseek_v4/test_c128_v2.py b/python/sglang/jit_kernel/tests/deepseek_v4/test_c128_v2.py new file mode 100644 index 000000000..e17de0855 --- /dev/null +++ b/python/sglang/jit_kernel/tests/deepseek_v4/test_c128_v2.py @@ -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"])) diff --git a/python/sglang/jit_kernel/tests/deepseek_v4/test_c4_v2.py b/python/sglang/jit_kernel/tests/deepseek_v4/test_c4_v2.py new file mode 100644 index 000000000..23a513da1 --- /dev/null +++ b/python/sglang/jit_kernel/tests/deepseek_v4/test_c4_v2.py @@ -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"])) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index cbd81b0d3..efd182b12 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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 diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 93e450765..f9f396428 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -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, diff --git a/python/sglang/srt/layers/attention/dsv4/compressor.py b/python/sglang/srt/layers/attention/dsv4/compressor.py index 0cc5bc312..332f52977 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor.py @@ -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, diff --git a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py new file mode 100644 index 000000000..ede264de4 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py @@ -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, + ) diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py index 3bc982446..5c538631e 100644 --- a/python/sglang/srt/layers/attention/dsv4/indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/indexer.py @@ -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) diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index e0e7dd56f..bc8aa8db1 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -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, diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 4461a0a27..3ad06d8c2 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -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,