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

dsv4.1: compression, KV I/O, and metadata kernels (#39652)

Browse Source 
Co-authored-by: BBuf <1182563586@qq.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: DarkSharpness <76582120+DarkSharpness@users.noreply.github.com>
This commit is contained in:
Liangsheng Yin
2026-09-16 13:54:07 -07:00
committed by GitHub
co-authored by BBuf Claude Opus 5 DarkSharpness
parent 869674b3a7
commit 13d593b6cf
19 changed files with 2102 additions and 279 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/deepseek_v4/c1.cuh
+312
View File
@@ -0,0 +1,312 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/math.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/fp4_utils.cuh>
#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
#include <sgl_kernel/deepseek_v4/kv_layout.cuh>
#include <tvm/ffi/container/tensor.h>
#include <bit>
#include <cstdint>
namespace sglang {
/// \brief Ratio-1 decode compressor: RMSNorm and the whole main-KV write.
///
/// At ratio 1 the latent stands for the token itself, so the kernel's input is the
/// `wkv` GEMM output and its RoPE position is `positions`, not `positions - 1`.
/// `kv_output` is the pre-RoPE latent, for the index-K branch's `wk` projection.
struct Compress1DecodeParams {
const bf16_t* __restrict__ kv_input; // [num_tokens, kHeadDim] bf16
bf16_t* __restrict__ kv_output; // [num_tokens, kHeadDim] bf16, pre-RoPE
const bf16_t* __restrict__ norm_weight; // [kHeadDim] bf16
const float* __restrict__ freqs_cis; // [max_pos, kRopeDim] fp32, real/imag interleaved
const void* __restrict__ positions; // [num_tokens] PosT
const void* __restrict__ out_loc; // [num_tokens] LocT compressed slot; 0 marks a padded row
uint8_t* __restrict__ kvcache; // [npages, kPageBytes] uint8
float eps;
};
/// Elements per thread; 256 threads per token measured fastest on B200 decode batches.
/// At (512, 64) it also keeps the nope/rope split warp-aligned, as the fp8 amax reduction requires.
constexpr uint32_t kC1VecSize = 2;
/// \brief RMSNorm + RoPE tail + fp4 fake-quant + the FlashMLA store.
///
/// One CTA per token, `kHeadDim / kC1VecSize` threads over the row.
///
/// The three reductions have different widths and are not interchangeable: the RMSNorm
/// statistic spans the row, an fp8 store scale 64 elements, an fp4 block 16.
///
/// kLayout is the cache's page format: V4 and V41 store the fake-quantized value; V41_FP4
/// stores the e2m1 codes and their e4m3 scales directly, so the fp4 rounding happens once.
template <
int64_t kHeadDim,
int64_t kRopeDim,
int32_t kPageBits,
typename PosT,
typename LocT,
deepseek_v4::KVLayout kLayout,
bool kUsePDL>
__global__ __launch_bounds__(kHeadDim / kC1VecSize) void flash_c1_decode_kernel(
const __grid_constant__ Compress1DecodeParams params) {
using namespace device;
using deepseek_v4::KVLayout;
using deepseek_v4::fp8::cast_to_ue8m0;
using deepseek_v4::fp8::inv_scale_ue8m0;
using deepseek_v4::fp8::pack_fp8;
/// Threads over one token; the leading kNopeLanes carry the fp8 nope part, the rest the bf16 RoPE tail.
constexpr uint32_t kVecSize = kC1VecSize;
constexpr uint32_t kRowLanes = kHeadDim / kVecSize;
constexpr uint32_t kNopeLanes = (kHeadDim - kRopeDim) / kVecSize;
constexpr uint32_t kRowWarps = kRowLanes / kWarpThreads;
constexpr uint32_t kFp8Lanes = 64 / kVecSize;
constexpr uint32_t kFp4Lanes = deepseek_v4::fp4::kCompressedKVBlockSize / kVecSize;
using Paged = deepseek_v4::PagedKV<kLayout, kPageBits>;
static_assert(kHeadDim == 512 && kRopeDim == 64, "the FlashMLA layouts require (512, 64)");
static_assert(kHeadDim % kVecSize == 0 && kVecSize % 2 == 0);
static_assert(kRowLanes % kWarpThreads == 0, "a token owns a whole number of warps");
static_assert(kNopeLanes % kFp8Lanes == 0, "the nope part must end on an fp8 scale block");
static_assert(
(kHeadDim - kRopeDim) % deepseek_v4::fp4::kCompressedKVBlockSize == 0,
"no fp4 block may straddle the nope/rope seam");
static_assert(kFp8Lanes <= kWarpThreads && kFp4Lanes <= kWarpThreads);
using bf16_vec_t = AlignedVector<bf16x2_t, kVecSize / 2>;
using fp8_vec_t = AlignedVector<fp8x2_e4m3_t, kVecSize / 2>;
using freq_vec_t = AlignedVector<float, kVecSize>;
const uint32_t tx = threadIdx.x;
const uint32_t row = blockIdx.x;
// `out_loc` and `positions` are step metadata, independent of the PDL producer.
// Slots fit in int32; padded rows are suppressed at the cache store.
const auto out_loc = static_cast<int32_t>(static_cast<const LocT*>(params.out_loc)[row]);
const auto position = static_cast<int64_t>(static_cast<const PosT*>(params.positions)[row]);
PDLWaitPrimary<kUsePDL>();
float data[kVecSize];
bf16_vec_t latent;
{
bf16_vec_t input, weight;
input.load(params.kv_input + row * kHeadDim, tx);
weight.load(params.norm_weight, tx);
// `project` already returns bf16 at ratio 1, so `finish`'s `.to(bfloat16)`
// is a no-op and the statistic is taken over the loaded values as they are.
float local_sqrsum = 0.0f;
#pragma unroll
for (uint32_t j = 0; j < kVecSize / 2; ++j) {
const auto [x, y] = cast<fp32x2_t>(input[j]);
local_sqrsum += x * x;
local_sqrsum += y * y;
data[j * 2 + 0] = x;
data[j * 2 + 1] = y;
}
__shared__ float s_warp_sum[kRowWarps];
s_warp_sum[tx / kWarpThreads] = warp::reduce_sum(local_sqrsum);
__syncthreads();
float sqrsum = 0.0f;
#pragma unroll
for (uint32_t i = 0; i < kRowWarps; ++i) {
sqrsum += s_warp_sum[i];
}
constexpr float kInvHeadDim = 1.0f / static_cast<float>(kHeadDim);
const auto norm_factor = math::rsqrt(sqrsum * kInvHeadDim + params.eps);
#pragma unroll
for (uint32_t j = 0; j < kVecSize / 2; ++j) {
const auto [wx, wy] = cast<fp32x2_t>(weight[j]);
const auto x = data[j * 2 + 0] * norm_factor * wx;
const auto y = data[j * 2 + 1] * norm_factor * wy;
latent[j] = cast<bf16x2_t>(fp32x2_t{x, y});
}
}
latent.store(params.kv_output + row * kHeadDim, tx);
PDLTriggerSecondary<kUsePDL>();
// Match finish()'s bf16 rounding before the main-KV RoPE.
#pragma unroll
for (uint32_t j = 0; j < kVecSize / 2; ++j) {
const auto [x, y] = cast<fp32x2_t>(latent[j]);
data[j * 2 + 0] = x;
data[j * 2 + 1] = y;
}
if (tx >= kNopeLanes) {
// Match rope_tail()'s bf16 rounding: it ends in `.to(x.dtype)` before the fake-quant.
freq_vec_t freq;
freq.load(params.freqs_cis + position * kRopeDim, tx - kNopeLanes);
#pragma unroll
for (uint32_t j = 0; j < kVecSize / 2; ++j) {
const auto k = j * 2;
const auto x_real = data[k + 0];
const auto x_imag = data[k + 1];
const auto f_real = freq[k + 0];
const auto f_imag = freq[k + 1];
const auto rotated =
cast<bf16x2_t>(fp32x2_t{x_real * f_real - x_imag * f_imag, x_real * f_imag + x_imag * f_real});
const auto [r0, r1] = cast<fp32x2_t>(rotated);
data[k + 0] = r0;
data[k + 1] = r1;
}
}
if constexpr (kLayout == KVLayout::V41_FP4) {
// The fp4 cache takes the rotated bf16 value as is: its row quantizer is the fake quant, minus the dequant.
if (out_loc <= 0) return;
const auto kv_row = Paged::row(params.kvcache, out_loc);
return deepseek_v4::v41::store_row<kLayout>(kv_row.data, kv_row.scale, tx, data);
}
// FP4/E4M3 fake-quant over 16 elements, i.e. kFp4Lanes threads.
{
float amax = fabsf(data[0]);
#pragma unroll
for (uint32_t i = 1; i < kVecSize; ++i) {
amax = fmaxf(amax, fabsf(data[i]));
}
amax = warp::reduce_max<kFp4Lanes>(amax);
const auto scale = deepseek_v4::fp4::compressed_kv_scale(amax);
#pragma unroll
for (uint32_t i = 0; i < kVecSize / 2; ++i) {
const auto [x, y] = deepseek_v4::fp4::fake_quant_compressed_kv_x2({data[i * 2 + 0], data[i * 2 + 1]}, scale);
data[i * 2 + 0] = x;
data[i * 2 + 1] = y;
}
}
// A padded CUDA-graph row carries `out_loc == 0`, the reserved dummy slot, and must publish
// nothing: at ratio 1 the compressed slot is the FULL slot, so there is no other marker to read.
if (out_loc <= 0) return;
const auto kv_row = Paged::row(params.kvcache, out_loc);
if constexpr (kLayout == KVLayout::V41) {
// fp8 with one ue8m0 scale per 32 elements over the whole row, RoPE included.
return deepseek_v4::v41::store_row<kLayout>(kv_row.data, kv_row.scale, tx, data);
}
const auto value_ptr = kv_row.data;
if (tx >= kNopeLanes) {
bf16_vec_t rope_out;
#pragma unroll
for (uint32_t j = 0; j < kVecSize / 2; ++j) {
rope_out[j] = cast<bf16x2_t>(fp32x2_t{data[j * 2 + 0], data[j * 2 + 1]});
}
rope_out.store(value_ptr + (kHeadDim - kRopeDim), tx - kNopeLanes);
} else {
// fp8 e4m3 with one ue8m0 scale per 64 elements.
float abs_max = fabsf(data[0]);
#pragma unroll
for (uint32_t i = 1; i < kVecSize; ++i) {
abs_max = fmaxf(abs_max, fabsf(data[i]));
}
abs_max = warp::reduce_max<kFp8Lanes>(abs_max);
const auto scale_ue8m0 = cast_to_ue8m0(fmaxf(1e-4f, abs_max) / math::FP8_E4M3_MAX);
const auto inv_scale = inv_scale_ue8m0(scale_ue8m0);
fp8_vec_t nope_out;
#pragma unroll
for (uint32_t i = 0; i < kVecSize / 2; ++i) {
nope_out[i] = pack_fp8(data[i * 2 + 0] * inv_scale, data[i * 2 + 1] * inv_scale);
}
nope_out.store(value_ptr, tx);
kv_row.scale[tx / kFp8Lanes] = scale_ue8m0;
}
}
/// \brief Host side of `flash_c1_decode_kernel`.
template <int64_t kHeadDim, int64_t kRopeDim, uint32_t kPageSize, deepseek_v4::KVLayout kLayout, bool kUsePDL>
struct FlashCompress1Kernel {
static constexpr int32_t kPageBits = std::bit_width(kPageSize) - 1;
static constexpr int64_t kPageBytes = deepseek_v4::kv_page_bytes<kLayout>(kPageSize);
static constexpr uint32_t kBlockSize = kHeadDim / kC1VecSize;
static_assert(std::has_single_bit(kPageSize), "the page/slot split needs a power-of-two page");
static_assert(kBlockSize % device::kWarpThreads == 0 && kBlockSize <= 1024);
static_assert(kLayout != deepseek_v4::KVLayout::V4 || kPageBytes == host::div_ceil(584ll * kPageSize, 576) * 576);
template <typename PosT, typename LocT>
static constexpr auto kernel = flash_c1_decode_kernel<kHeadDim, kRopeDim, kPageBits, PosT, LocT, kLayout, kUsePDL>;
/// \brief The (`positions`, `out_loc`) dtype pair, resolved at run time.
static auto select(const bool pos_i32, const bool loc_i32) {
if (pos_i32) return loc_i32 ? kernel<int32_t, int32_t> : kernel<int32_t, int64_t>;
return loc_i32 ? kernel<int64_t, int32_t> : kernel<int64_t, int64_t>;
}
/// \brief RMSNorm + RoPE + fp4 fake-quant + the FlashMLA store, one launch.
///
/// \param kv_input `[num_tokens, kHeadDim]` bf16, the `wkv` projection.
/// \param kv_output `[num_tokens, kHeadDim]` bf16, the pre-RoPE latent.
/// \param norm_weight `[kHeadDim]` bf16.
/// \param freqs_cis `[max_pos, kRopeDim]` fp32, real/imag interleaved.
/// \param positions `[num_tokens]` int32 or int64, indexed as-is.
/// \param out_loc `[num_tokens]` int32 or int64, the compressed slot; `0` is a padded row.
/// \param kvcache `[npages, kPageBytes]` uint8, or the pool's fp8 view of it.
static void run_decode_fusion(
const tvm::ffi::TensorView kv_input,
const tvm::ffi::TensorView kv_output,
const tvm::ffi::TensorView norm_weight,
const tvm::ffi::TensorView freqs_cis,
const tvm::ffi::TensorView positions,
const tvm::ffi::TensorView out_loc,
const tvm::ffi::TensorView kvcache,
const float eps) {
using namespace host;
auto N = SymbolicSize{"num_tokens"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({N, kHeadDim}) //
.with_dtype<bf16_t>()
.with_device(device_)
.verify(kv_input)
.verify(kv_output);
TensorMatcher({kHeadDim}).with_dtype<bf16_t>().with_device(device_).verify(norm_weight);
// Real/imag interleaved, so the trailing dim is kRopeDim, not kRopeDim / 2.
TensorMatcher({-1, kRopeDim}).with_dtype<fp32_t>().with_device(device_).verify(freqs_cis);
// The scheduler's `out_cache_loc` (which `c1_out_loc` aliases at ratio 1)
// is int64; the unit tests hand int32. Both are indexed as-is.
auto pos_dtype = SymbolicDType{};
auto loc_dtype = SymbolicDType{};
TensorMatcher({N}).with_dtype<int32_t, int64_t>(pos_dtype).with_device(device_).verify(positions);
TensorMatcher({N}).with_dtype<int32_t, int64_t>(loc_dtype).with_device(device_).verify(out_loc);
// The pool allocates the buffer as uint8 and hands it out viewed as its fp8
// dtype (`get_extra_key_buffer`); both are one byte per element.
TensorMatcher({-1, kPageBytes}).with_dtype<uint8_t, fp8_e4m3_t>().with_device(device_).verify(kvcache);
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
if (num_tokens == 0) return;
const auto params = Compress1DecodeParams{
.kv_input = static_cast<const bf16_t*>(kv_input.data_ptr()),
.kv_output = static_cast<bf16_t*>(kv_output.data_ptr()),
.norm_weight = static_cast<const bf16_t*>(norm_weight.data_ptr()),
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
.positions = positions.data_ptr(),
.out_loc = out_loc.data_ptr(),
.kvcache = static_cast<uint8_t*>(kvcache.data_ptr()),
.eps = eps,
};
const auto k = select(pos_dtype.is_type<int32_t>(), loc_dtype.is_type<int32_t>());
LaunchKernel(num_tokens, kBlockSize, device_.unwrap()) //
.enable_pdl(kUsePDL)(k, params);
}
};
// The JIT module names and wrappers spell the layouts as bare enumerators.
using enum deepseek_v4::KVLayout;
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/c2.cuh
+398
View File
@@ -0,0 +1,398 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/math.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/fp4_utils.cuh>
#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
#include <sgl_kernel/deepseek_v4/kv_layout.cuh>
#include <tvm/ffi/container/tensor.h>
#include <bit>
#include <cstdint>
namespace sglang {
/// \brief Ratio-2 decode compressor: pair-pool, RMSNorm and the main-KV write.
///
/// `kv_input` and `kv_state` rows are `2 * kHeadDim` floats, kv then score.
/// `kv_output` is the pre-RoPE latent, for the index-K branch's `wk`.
struct Compress2DecodeParams {
const float* __restrict__ kv_input; // [num_tokens, 2 * kHeadDim] fp32
/// `CompressStatePool`'s flat `KVAndScore` buffer, `[size, 2 * kHeadDim]` fp32, kv in the low
/// half and score in the high half; a request's pending pair lives at `req * ring_size + pos % ring_size`.
float* __restrict__ kv_state;
bf16_t* __restrict__ kv_output; // [num_tokens, kHeadDim] bf16, pre-RoPE
const bf16_t* __restrict__ norm_weight; // [kHeadDim] bf16
const float* __restrict__ freqs_cis; // [max_pos, kRopeDim] fp32, real/imag interleaved
const void* __restrict__ positions; // [num_tokens] PosT
const int64_t* __restrict__ req; // [num_tokens], req_pool_idx per token
const void* __restrict__ raw_out_loc; // [num_tokens] LocT, the FULL slot; 0 marks a padded row
uint8_t* __restrict__ kvcache; // [npages, kPageBytes] uint8
/// Positions per request slot in the pair-state ring.
uint32_t ring_size;
float eps;
};
/// Elements per thread; 256 threads per token measured fastest on B200 decode batches.
/// At (512, 64) it also keeps the nope/rope split warp-aligned, as the fp8 amax reduction requires.
constexpr uint32_t kC2VecSize = 2;
/// \brief grid = num_tokens, block = kHeadDim / kC2VecSize.
///
/// An odd position completes a group with its even predecessor; an even one parks in the state.
///
/// Under target-verify, `draft_len` consecutive positions per request: `blockIdx.x` is the
/// position inside the block, `blockIdx.y` the request, and every row but the first takes its
/// partner from the previous `kv_input` row instead of the ring.
///
/// The three reductions have different widths and are not interchangeable: the RMSNorm
/// statistic spans the row, an fp8 store scale 64 elements, an fp4 block 16.
///
/// kLayout is the cache's page format: V4 and V41 store the fake-quantized value; V41_FP4
/// stores the e2m1 codes and their e4m3 scales directly, so the fp4 rounding happens once.
template <
bool kVerify,
int64_t kHeadDim,
int64_t kRopeDim,
int32_t kPageBits,
typename PosT,
typename LocT,
deepseek_v4::KVLayout kLayout,
bool kUsePDL>
__global__ __launch_bounds__(kHeadDim / kC2VecSize) void flash_c2_decode_kernel(const Compress2DecodeParams params) {
using namespace device;
using deepseek_v4::KVLayout;
using deepseek_v4::fp8::cast_to_ue8m0;
using deepseek_v4::fp8::inv_scale_ue8m0;
using deepseek_v4::fp8::pack_fp8;
constexpr uint32_t kVecSize = kC2VecSize;
constexpr uint32_t kCTASize = kHeadDim / kVecSize;
constexpr int64_t kStride = kHeadDim * 2;
/// Threads covering the fp8 nope part; the rest carry the bf16 RoPE tail.
constexpr uint32_t kNopeThreads = (kHeadDim - kRopeDim) / kVecSize;
constexpr uint32_t kFp8Lanes = 64 / kVecSize;
constexpr uint32_t kFp4Lanes = deepseek_v4::fp4::kCompressedKVBlockSize / kVecSize;
using Paged = deepseek_v4::PagedKV<kLayout, kPageBits>;
static_assert(kHeadDim == (kVecSize * kCTASize));
static_assert(kCTASize % kWarpThreads == 0);
static_assert(kNopeThreads % kFp8Lanes == 0, "the nope part must end on an fp8 scale block");
static_assert(kWarpThreads % kFp8Lanes == 0 && kWarpThreads % kFp4Lanes == 0);
static_assert(kHeadDim == 512 && kRopeDim == 64, "the FlashMLA layouts require (512, 64)");
using fp32_vec_t = AlignedVector<float, kVecSize>;
using bf16_vec_t = AlignedVector<bf16x2_t, kVecSize / 2>;
const auto tx = threadIdx.x;
// Verify gives each request a CTA column; decode a flat grid of one row each.
const auto row = kVerify ? blockIdx.y * gridDim.x + blockIdx.x : blockIdx.x;
PDLWaitPrimary<kUsePDL>();
// Slots fit in int32 whatever width the scheduler hands them in.
const auto raw_out_loc = static_cast<int32_t>(static_cast<const LocT*>(params.raw_out_loc)[row]);
// CUDA graph padding is a completely inert row: do not read schedule, input,
// state, or RoPE data, and do not publish output or update the cache.
if (raw_out_loc == 0) return PDLTriggerSecondary<kUsePDL>();
const auto pos = static_cast<const PosT*>(params.positions)[row];
// A completing row reads the slot left by `pos - 1`;
// a pending row writes its own slot, so reads and writes stay disjoint.
const auto rid = params.req[row];
fp32_vec_t kv_new, score_new;
kv_new.load(params.kv_input + row * kStride, tx);
score_new.load(params.kv_input + row * kStride, tx + kCTASize);
const auto ring = static_cast<int64_t>(rid) * params.ring_size;
const auto read_row = ring + (pos - 1 + params.ring_size) % params.ring_size;
const auto write_row = ring + pos % params.ring_size;
fp32_vec_t kv_old, score_old;
const float* partner = params.kv_state + read_row * kStride;
if constexpr (kVerify) {
if (blockIdx.x != 0) partner = params.kv_input + static_cast<int64_t>(row - 1) * kStride;
}
kv_old.load(partner, tx);
score_old.load(partner, tx + kCTASize);
if ((pos & 1) == 0) {
kv_new.store(params.kv_state + write_row * kStride, tx);
score_new.store(params.kv_state + write_row * kStride, tx + kCTASize);
return PDLTriggerSecondary<kUsePDL>();
}
constexpr uint32_t kNumWarps = kCTASize / kWarpThreads;
__shared__ float s_warp_sum[kNumWarps];
fp32_vec_t staged, freq;
bf16_vec_t weight, out;
weight.load(params.norm_weight, tx);
if (tx >= kNopeThreads) freq.load(params.freqs_cis + (pos - 1) * kRopeDim, tx - kNopeThreads);
// With two scores `exp(-|s0 - s1|)` is the whole softmax: one exp, argument
// always <= 0, so no max-subtraction pass and no overflow.
#pragma unroll
for (uint32_t i = 0; i < kVecSize; ++i) {
const auto delta = score_old[i] - score_new[i];
const auto scale = expf(-fabsf(delta));
const auto scale_0 = delta > 0 ? 1.0f : scale;
const auto scale_1 = delta > 0 ? scale : 1.0f;
staged[i] = (kv_old[i] * scale_0 + kv_new[i] * scale_1) / (1.0f + scale);
}
// `finish` casts to bf16 before the norm, so the sum of squares must see the rounded values.
float local_sqrsum = 0.0f;
#pragma unroll
for (uint32_t i = 0; i < kVecSize / 2; ++i) {
const auto packed = fp32x2_t{staged[i * 2 + 0], staged[i * 2 + 1]};
const auto [x, y] = cast<fp32x2_t>(cast<bf16x2_t>(packed));
local_sqrsum += x * x;
local_sqrsum += y * y;
staged[i * 2 + 0] = x;
staged[i * 2 + 1] = y;
}
const auto warp_sum = warp::reduce_sum(local_sqrsum);
s_warp_sum[tx / kWarpThreads] = warp_sum;
__syncthreads();
float sqrsum = 0.0f;
#pragma unroll
for (uint32_t i = 0; i < kNumWarps; ++i) {
sqrsum += s_warp_sum[i];
}
constexpr float kInvScale = 1.0f / static_cast<float>(kHeadDim);
const auto norm_factor = math::rsqrt(sqrsum * kInvScale + params.eps);
#pragma unroll
for (uint32_t i = 0; i < kVecSize / 2; ++i) {
const auto [wx, wy] = cast<fp32x2_t>(weight[i]);
const auto x = staged[i * 2 + 0] * norm_factor * wx;
const auto y = staged[i * 2 + 1] * norm_factor * wy;
out[i] = cast<bf16x2_t>(fp32x2_t{x, y});
}
// The pre-RoPE latent, for the index-K branch's `wk` projection. Published
// before the trigger because that GEMM is the successor that reads it.
out.store(params.kv_output, static_cast<int64_t>(row) * kCTASize + tx);
PDLTriggerSecondary<kUsePDL>();
// ---- main-KV branch: RoPE tail, fp4 fake-quant, cache store ----
// Match finish()'s bf16 rounding before RoPE.
#pragma unroll
for (uint32_t i = 0; i < kVecSize / 2; ++i) {
const auto [x, y] = cast<fp32x2_t>(out[i]);
staged[i * 2 + 0] = x;
staged[i * 2 + 1] = y;
}
if (tx >= kNopeThreads) {
// Match rope_tail()'s bf16 rounding before fake quantization.
// Only odd positions reach here; the latent represents `pos - 1`.
#pragma unroll
for (uint32_t i = 0; i < kVecSize / 2; ++i) {
const auto x_real = staged[i * 2 + 0];
const auto x_imag = staged[i * 2 + 1];
const auto f_real = x_real * freq[i * 2 + 0] - x_imag * freq[i * 2 + 1];
const auto f_imag = x_real * freq[i * 2 + 1] + x_imag * freq[i * 2 + 0];
const auto rotated = cast<bf16x2_t>(fp32x2_t{f_real, f_imag});
const auto [r0, r1] = cast<fp32x2_t>(rotated);
staged[i * 2 + 0] = r0;
staged[i * 2 + 1] = r1;
}
}
if constexpr (kLayout == KVLayout::V41_FP4) {
// The fp4 cache takes the rotated bf16 value as is: its row quantizer is the fake quant, minus the dequant.
const int32_t out_loc = raw_out_loc >> 1;
const auto kv_row = Paged::row(params.kvcache, out_loc);
return deepseek_v4::v41::store_row<kLayout>(kv_row.data, kv_row.scale, tx, staged);
}
// FP4/E4M3 fake-quant over 16 elements, i.e. kFp4Lanes threads.
{
float amax = fabsf(staged[0]);
#pragma unroll
for (uint32_t i = 1; i < kVecSize; ++i) {
amax = fmaxf(amax, fabsf(staged[i]));
}
amax = warp::reduce_max<kFp4Lanes>(amax);
const auto scale = deepseek_v4::fp4::compressed_kv_scale(amax);
#pragma unroll
for (uint32_t i = 0; i < kVecSize / 2; ++i) {
const auto [x, y] = deepseek_v4::fp4::fake_quant_compressed_kv_x2({staged[i * 2 + 0], staged[i * 2 + 1]}, scale);
staged[i * 2 + 0] = x;
staged[i * 2 + 1] = y;
}
}
// `raw_out_loc / ratio`; ratio 2 makes it a shift.
const int32_t out_loc = raw_out_loc >> 1;
const auto kv_row = Paged::row(params.kvcache, out_loc);
if constexpr (kLayout == KVLayout::V41) {
// fp8 with one ue8m0 scale per 32 elements over the whole row, RoPE included.
return deepseek_v4::v41::store_row<kLayout>(kv_row.data, kv_row.scale, tx, staged);
}
const auto value_ptr = kv_row.data;
if (tx >= kNopeThreads) {
bf16_vec_t rope_out;
#pragma unroll
for (uint32_t i = 0; i < kVecSize / 2; ++i) {
rope_out[i] = cast<bf16x2_t>(fp32x2_t{staged[i * 2 + 0], staged[i * 2 + 1]});
}
rope_out.store(value_ptr + (kHeadDim - kRopeDim), tx - kNopeThreads);
} else {
// fp8 e4m3 with one ue8m0 scale per 64 elements.
auto abs_max = fabsf(staged[0]);
#pragma unroll
for (uint32_t i = 1; i < kVecSize; ++i) {
abs_max = fmaxf(abs_max, fabsf(staged[i]));
}
abs_max = warp::reduce_max<kFp8Lanes>(abs_max);
const auto scale_ue8m0 = cast_to_ue8m0(fmaxf(1e-4f, abs_max) / math::FP8_E4M3_MAX);
const auto inv_scale = inv_scale_ue8m0(scale_ue8m0);
#pragma unroll
for (uint32_t i = 0; i < kVecSize / 2; ++i) {
reinterpret_cast<fp8x2_e4m3_t*>(value_ptr)[tx * (kVecSize / 2) + i] =
pack_fp8(staged[i * 2 + 0] * inv_scale, staged[i * 2 + 1] * inv_scale);
}
kv_row.scale[tx / kFp8Lanes] = scale_ue8m0;
}
}
template <int64_t kHeadDim, int64_t kRopeDim, uint32_t kPageSize, deepseek_v4::KVLayout kLayout, bool kUsePDL>
struct FlashCompress2Kernel {
static constexpr uint32_t kBlockSize = kHeadDim / kC2VecSize;
static constexpr int32_t kPageBits = std::bit_width(kPageSize) - 1;
static constexpr int64_t kPageBytes = deepseek_v4::kv_page_bytes<kLayout>(kPageSize);
static_assert(kLayout != deepseek_v4::KVLayout::V4 || kPageBytes == host::div_ceil(584ll * kPageSize, 576) * 576);
template <bool kVerify, typename PosT, typename LocT>
static constexpr auto kernel =
flash_c2_decode_kernel<kVerify, kHeadDim, kRopeDim, kPageBits, PosT, LocT, kLayout, kUsePDL>;
/// \brief The (`positions`, `raw_out_loc`) dtype pair, resolved at run time.
template <bool kVerify>
static auto select(const bool pos_i32, const bool loc_i32) {
if (pos_i32) return loc_i32 ? kernel<kVerify, int32_t, int32_t> : kernel<kVerify, int32_t, int64_t>;
return loc_i32 ? kernel<kVerify, int64_t, int32_t> : kernel<kVerify, int64_t, int64_t>;
}
// The sum of squares reduces through a fixed-size shared array, so the CTA must be whole warps.
static_assert(kHeadDim % (4 * device::kWarpThreads) == 0, "head_dim must be a multiple of 128");
static_assert(std::has_single_bit(kPageSize), "the page/slot split needs a power-of-two page");
/// \brief `run_decode_fusion` for a target-verify block.
///
/// `draft_len` consecutive positions per request, request-major, which the
/// grid reproduces as `draft_len x batch`.
static void run_decode_fusion(
const tvm::ffi::TensorView kv_input,
const tvm::ffi::TensorView kv_state,
const tvm::ffi::TensorView kv_output,
const tvm::ffi::TensorView norm_weight,
const tvm::ffi::TensorView positions,
const tvm::ffi::TensorView req,
const tvm::ffi::TensorView raw_out_loc,
const float eps,
const tvm::ffi::TensorView freqs_cis,
const tvm::ffi::TensorView kvcache,
const int64_t ring_size,
const int64_t draft_len) {
launch(
kv_input,
kv_state,
kv_output,
norm_weight,
positions,
req,
raw_out_loc,
eps,
ring_size,
freqs_cis,
kvcache,
draft_len);
}
private:
static void launch(
const tvm::ffi::TensorView kv_input,
const tvm::ffi::TensorView kv_state,
const tvm::ffi::TensorView kv_output,
const tvm::ffi::TensorView norm_weight,
const tvm::ffi::TensorView positions,
const tvm::ffi::TensorView req,
const tvm::ffi::TensorView raw_out_loc,
const float eps,
const int64_t ring_size,
const tvm::ffi::TensorView freqs_cis,
const tvm::ffi::TensorView kvcache,
const int64_t draft_len) {
using namespace host;
auto N = SymbolicSize{"num_tokens"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLGPU>();
TensorMatcher({N, kHeadDim * 2}).with_dtype<fp32_t>().with_device(device_).verify(kv_input);
TensorMatcher({-1, kHeadDim * 2}).with_dtype<fp32_t>().with_device(device_).verify(kv_state);
// Only rows that complete a group are written.
TensorMatcher({N, kHeadDim}).with_dtype<bf16_t>().with_device(device_).verify(kv_output);
TensorMatcher({kHeadDim}).with_dtype<bf16_t>().with_device(device_).verify(norm_weight);
// Metadata retains its original dtypes: the scheduler uses int64 locations,
// while callers may also supply int32 locations and positions.
auto pos_dtype = SymbolicDType{};
auto loc_dtype = SymbolicDType{};
TensorMatcher({N}).with_dtype<int32_t, int64_t>(pos_dtype).with_device(device_).verify(positions);
TensorMatcher({N}).with_dtype<int64_t>().with_device(device_).verify(req);
TensorMatcher({N}).with_dtype<int32_t, int64_t>(loc_dtype).with_device(device_).verify(raw_out_loc);
// Real/imag interleaved, so the trailing dim is kRopeDim, not kRopeDim / 2.
TensorMatcher({-1, kRopeDim}).with_dtype<fp32_t>().with_device(device_).verify(freqs_cis);
// The pool allocates the buffer as uint8 and hands it out viewed as its
// fp8 dtype (`get_extra_key_buffer`); both are one byte per element.
TensorMatcher({-1, kPageBytes}).with_dtype<uint8_t, fp8_e4m3_t>().with_device(device_).verify(kvcache);
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
if (num_tokens == 0) return;
const auto is_verify = draft_len > 1;
CHECK_HOST(ring_size > 0 && draft_len >= 1);
CHECK_HOST(!is_verify || num_tokens % draft_len == 0);
CHECK_HOST(!is_verify || ring_size > draft_len)
<< "the pair-state ring (" << ring_size << ") must be wider than the draft length (" << draft_len << ")";
const auto params = Compress2DecodeParams{
.kv_input = static_cast<const float*>(kv_input.data_ptr()),
.kv_state = static_cast<float*>(kv_state.data_ptr()),
.kv_output = static_cast<bf16_t*>(kv_output.data_ptr()),
.norm_weight = static_cast<const bf16_t*>(norm_weight.data_ptr()),
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
.positions = positions.data_ptr(),
.req = static_cast<const int64_t*>(req.data_ptr()),
.raw_out_loc = raw_out_loc.data_ptr(),
.kvcache = static_cast<uint8_t*>(kvcache.data_ptr()),
.ring_size = static_cast<uint32_t>(ring_size),
.eps = eps,
};
// `LaunchKernel` is move-only, so each arm builds its own.
const auto pos_i32 = pos_dtype.is_type<int32_t>();
const auto loc_i32 = loc_dtype.is_type<int32_t>();
if (is_verify) {
const auto block = static_cast<uint32_t>(draft_len);
const auto k = select<true>(pos_i32, loc_i32);
LaunchKernel(dim3{block, num_tokens / block}, kBlockSize, device_.unwrap()) //
.enable_pdl(kUsePDL)(k, params);
} else {
const auto k = select<false>(pos_i32, loc_i32);
LaunchKernel(num_tokens, kBlockSize, device_.unwrap()) //
.enable_pdl(kUsePDL)(k, params);
}
}
};
// The JIT module names and wrappers spell the layouts as bare enumerators.
using enum deepseek_v4::KVLayout;
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh
+6 -4
View File
@@ -505,10 +505,12 @@ inline PrefillPlan plan_compress_prefill(
const auto f2s_ptr = static_cast<const F2S_T*>(full_to_state.data_ptr());
const auto batch_size = static_cast<uint32_t>(B.unwrap());
constexpr auto kMaxTokens = static_cast<uint32_t>(std::numeric_limits<uint16_t>::max());
// ragged_id is a zero-based uint16 index, so a 64K-token batch is valid.
constexpr auto kMaxTokens = static_cast<uint32_t>(std::numeric_limits<uint16_t>::max()) + 1;
RuntimeCheck(compress_ratio == 4 || compress_ratio == 128);
RuntimeCheck(!use_req_ring || compress_ratio == 4);
RuntimeCheck(batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens);
// Keep batch_id below 65535: pack_w(65535, 65535, ...) is the invalid sentinel.
RuntimeCheck(batch_size < kMaxTokens && 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);
// Write pad: trailing tokens kept resident so a verify batch's committed tail survives
@@ -750,9 +752,9 @@ inline PrefillPlan plan_compress_prefill_legacy(
const auto window_size = compress_ratio * (is_overlap ? 2 : 1);
const auto batch_size = static_cast<uint32_t>(B.unwrap());
constexpr auto kMaxTokens = static_cast<uint32_t>(std::numeric_limits<uint16_t>::max());
constexpr auto kMaxTokens = static_cast<uint32_t>(std::numeric_limits<uint16_t>::max()) + 1;
RuntimeCheck(compress_ratio == 4 || compress_ratio == 128);
RuntimeCheck(batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens);
RuntimeCheck(batch_size < kMaxTokens && batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens);
uint32_t counter = 0;
uint32_t counter_c = 0;
python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope_v2.cuh
+50 -11
View File
@@ -9,6 +9,7 @@
#include <sgl_kernel/deepseek_v4/compress_v2.cuh>
#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
#include <sgl_kernel/deepseek_v4/kv_layout.cuh>
#include <tvm/ffi/container/tensor.h>
@@ -386,14 +387,16 @@ constexpr int64_t kFp8TwoPoolRowBytes = 512;
// ----------------------------------------------------------------------------
// 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.
// Cache layout (kLayout): V4 = 584 bytes/token = 448 fp8 nope + 64 (=32 bf16x2) rope + 8 scale;
// V41 / V41_FP4 = the fully quantized fp8 (528 B) / fp4 (288 B) rows, one scale per 32 / 16 values.
// ----------------------------------------------------------------------------
template <
typename DType,
ForwardMode kMode,
int32_t kPageBits,
bool kBf16Store,
deepseek_v4::KVLayout kLayout,
bool kUsePDL,
bool kBf16Store = false,
bool kFp8TwoPool = false>
FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormRopeStoreParams params) {
using namespace device;
@@ -407,12 +410,15 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR
constexpr uint32_t kRopeWarp = kNumWarps - 1;
// kBf16Store: write the whole head_dim as plain BF16 (no fp8 / no scale) into a
// [num_slots, head_dim] bf16 cache (page_size==1) at row out_loc
static_assert(!(kBf16Store && kLayout != deepseek_v4::KVLayout::V4), "the bf16 store is not a paged layout");
using Paged = deepseek_v4::PagedKV<kLayout, kPageBits>;
// kFp8TwoPool: 512 B row holding the 448 fp8 nope + its UE8M0 scales, with rope
// split off into a second [num_slots, kRopeDim] bf16 pool at the same row
static_assert(!(kBf16Store && kFp8TwoPool));
constexpr int64_t kRowBytes = kBf16Store ? (kHeadDim * 2ll) : (kFp8TwoPool ? kFp8TwoPoolRowBytes : 576ll);
constexpr int64_t kPageBytes =
(kBf16Store || kFp8TwoPool) ? (kRowBytes << kPageBits) : host::div_ceil(584ll << kPageBits, 576) * 576;
static_assert(!(kFp8TwoPool && kLayout != deepseek_v4::KVLayout::V4), "the fp8 two-pool store is a V4 cache");
static_assert(kHeadDim == kBlockSize * kVecSize);
static_assert(kRopeDim == kWarpThreads * kVecSize);
static_assert(kHeadDim - kRopeDim == kRopeWarp * kWarpThreads * kVecSize);
@@ -480,10 +486,36 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR
}
}
const auto row = Paged::row(params.kvcache, out_loc);
if constexpr (kLayout != deepseek_v4::KVLayout::V4) {
// V4.1 layouts: the whole row is quantized. Match the unfused path, which quantizes the bf16
// the norm produces: round the normed values, rotate in bf16, round again, quantize the row.
using Packed = packed_t<DType>;
PDLTriggerSecondary<kUsePDL>();
auto rounded = cast<fp32x2_t>(cast<Packed>(fp32x2_t{data[0], data[1]}));
if (warp_id == kRopeWarp) {
const auto x_real = rounded.x;
const auto x_imag = rounded.y;
const auto freq_real = freq[0];
const auto freq_imag = freq[1];
rounded = cast<fp32x2_t>(
cast<Packed>(fp32x2_t{x_real * freq_real - x_imag * freq_imag, x_real * freq_imag + x_imag * freq_real}));
}
const float v[2] = {rounded.x, rounded.y};
deepseek_v4::v41::store_row<kLayout>(row.data, row.scale, tx, v);
return;
}
// The bf16 cache is dense [num_slots, head_dim] rows and the fp8 two-pool cache kRowBytes rows,
// both addressed by out_loc directly rather than through the paged helper.
const int64_t page = out_loc >> kPageBits;
const int64_t offset = out_loc & ((1 << kPageBits) - 1);
const auto page_ptr = params.kvcache + page * kPageBytes;
const auto value_ptr = page_ptr + offset * kRowBytes;
const auto value_ptr = kBf16Store ? params.kvcache + static_cast<int64_t>(out_loc) * (kHeadDim * 2)
: kFp8TwoPool ? page_ptr + offset * kRowBytes
: row.data;
PDLTriggerSecondary<kUsePDL>();
@@ -534,8 +566,7 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR
scale_ptr[0] = scale_ue8m0;
scale_ptr[1] = scale_ue8m0;
} else {
const auto scale_ptr = page_ptr + (576 << kPageBits) + offset * 8;
static_cast<uint8_t*>(scale_ptr)[warp_id] = scale_ue8m0;
static_cast<uint8_t*>(row.scale)[warp_id] = scale_ue8m0;
}
}
}
@@ -546,15 +577,19 @@ template <
int64_t kHeadDim,
int64_t kRopeDim,
uint32_t kPageSize,
bool kUsePDL,
int32_t kPreshuffleSize = 0,
bool kBf16Store = false>
int32_t kPreshuffleSize,
bool kBf16Store,
deepseek_v4::KVLayout kLayout,
bool kUsePDL>
struct FusedNormRopeKernel {
static constexpr int32_t kLogPageSize = std::countr_zero(kPageSize);
static constexpr bool kIsIndexer = (kHeadDim == 128);
static_assert(!(kIsIndexer && kBf16Store), "bf16 store only for flashmla head_dim=512");
static_assert(
!(kIsIndexer && kLayout != deepseek_v4::KVLayout::V4), "the V4.1 layouts are FlashMLA (head_dim=512) caches");
static constexpr int64_t kIndexerBytes = 132 * kPageSize;
static constexpr int64_t kFlashMLABytes = host::div_ceil(584 * kPageSize, 576) * 576;
static constexpr int64_t kFlashMLABytes = deepseek_v4::kv_page_bytes<kLayout>(kPageSize);
static_assert(kLayout != deepseek_v4::KVLayout::V4 || kFlashMLABytes == host::div_ceil(584 * kPageSize, 576) * 576);
static constexpr int64_t kBf16Bytes = kHeadDim * 2 * kPageSize; // plain bf16 cache
static constexpr int64_t kPageBytes = kBf16Store ? kBf16Bytes : (kIsIndexer ? kIndexerBytes : kFlashMLABytes);
@@ -567,7 +602,7 @@ struct FusedNormRopeKernel {
if constexpr (kIsIndexer) {
return fused_norm_rope_indexer<DType, kMode, kLogPageSize, kUsePDL, kPreshuffleSize>;
} else {
return fused_norm_rope_flashmla<DType, kMode, kLogPageSize, kUsePDL, kBf16Store>;
return fused_norm_rope_flashmla<DType, kMode, kLogPageSize, kBf16Store, kLayout, kUsePDL>;
}
}
@@ -575,7 +610,8 @@ struct FusedNormRopeKernel {
static constexpr auto select_fp8_2buff_kernel() {
static_assert(!kIsIndexer, "fp8 two-pool store is only defined for the flashmla latent");
static_assert(!kBf16Store, "fp8 two-pool store and bf16 store are separate layouts");
return fused_norm_rope_flashmla<DType, kMode, kLogPageSize, kUsePDL, false, true>;
static_assert(kLayout == deepseek_v4::KVLayout::V4, "the fp8 two-pool store is a V4 cache");
return fused_norm_rope_flashmla<DType, kMode, kLogPageSize, false, kLayout, kUsePDL, true>;
}
template <ForwardMode kMode>
@@ -791,4 +827,7 @@ struct FusedNormRopeKernel {
}
};
// The JIT module names and wrappers spell the layouts as bare enumerators.
using enum deepseek_v4::KVLayout;
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/main_norm_rope.cuh
+47 -10
View File
@@ -9,6 +9,7 @@
#include <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
#include <sgl_kernel/deepseek_v4/kv_layout.cuh>
#include <tvm/ffi/container/tensor.h>
@@ -259,13 +260,20 @@ struct FusedKNormRopeFlashMLAParams {
float eps;
};
template <typename DType, int64_t kHeadDim, int64_t kRopeDim, typename PosT, int32_t kPageBits, bool kUsePDL>
template <
typename DType,
int64_t kHeadDim,
int64_t kRopeDim,
typename PosT,
int32_t kPageBits,
deepseek_v4::KVLayout kLayout,
bool kUsePDL>
K_KERNEL void fused_k_norm_rope_flashmla(const __grid_constant__ FusedKNormRopeFlashMLAParams params) {
using namespace device;
constexpr int64_t kVecSize = 2;
constexpr uint32_t kRopeWarp = kFusedKNumWarps - 1;
constexpr int64_t kPageBytes = host::div_ceil(584ll << kPageBits, 576) * 576;
using Paged = deepseek_v4::PagedKV<kLayout, kPageBits>;
static_assert(kHeadDim == kFusedKBlockSize * kVecSize);
static_assert(kRopeDim == kWarpThreads * kVecSize);
static_assert(kHeadDim - kRopeDim == kRopeWarp * kWarpThreads * kVecSize);
@@ -325,10 +333,28 @@ K_KERNEL void fused_k_norm_rope_flashmla(const __grid_constant__ FusedKNormRopeF
// here, not at the load, so the out_loc prefetch overlaps the norm above.
if (out_loc < 0) return;
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;
const auto row = Paged::row(params.kvcache, out_loc);
if constexpr (kLayout != deepseek_v4::KVLayout::V4) {
// V4.1 layouts: every dim is quantized, one scale per 32 (fp8) or 16 (fp4) values. The
// reference rotates in bf16, so round the normed values, rotate, round again, then quantize.
using Packed = packed_t<DType>;
PDLTriggerSecondary<kUsePDL>();
auto rounded = cast<fp32x2_t>(cast<Packed>(fp32x2_t{data[0], data[1]}));
if (warp_id == kRopeWarp) {
const auto x_real = rounded.x;
const auto x_imag = rounded.y;
const auto freq_real = freq[0];
const auto freq_imag = freq[1];
rounded = cast<fp32x2_t>(
cast<Packed>(fp32x2_t{x_real * freq_real - x_imag * freq_imag, x_real * freq_imag + x_imag * freq_real}));
}
const float v[2] = {rounded.x, rounded.y};
return deepseek_v4::v41::store_row<kLayout>(row.data, row.scale, tx, v);
}
const auto value_ptr = row.data;
PDLTriggerSecondary<kUsePDL>();
@@ -351,22 +377,30 @@ K_KERNEL void fused_k_norm_rope_flashmla(const __grid_constant__ FusedKNormRopeF
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;
const auto scale_ptr = row.scale;
reinterpret_cast<fp8x2_e4m3_t*>(value_ptr)[tx] = result;
if (lane_id == 0) static_cast<uint8_t*>(scale_ptr)[warp_id] = scale_ue8m0;
}
}
template <typename DType, int64_t kHeadDim, int64_t kRopeDim, uint32_t kPageSize, bool kUsePDL>
template <
typename DType,
int64_t kHeadDim,
int64_t kRopeDim,
uint32_t kPageSize,
deepseek_v4::KVLayout kLayout,
bool kUsePDL>
struct FusedKNormRopeFlashMLAKernel {
static constexpr int32_t kLogPageSize = std::countr_zero(kPageSize);
static constexpr int64_t kPageBytes = host::div_ceil(584 * kPageSize, 576) * 576;
static constexpr int64_t kPageBytes = deepseek_v4::kv_page_bytes<kLayout>(kPageSize);
static_assert(kLayout != deepseek_v4::KVLayout::V4 || kPageBytes == host::div_ceil(584 * kPageSize, 576) * 576);
static_assert(std::has_single_bit(kPageSize), "kPageSize must be a power of 2");
static_assert(1 << kLogPageSize == kPageSize);
static_assert(kHeadDim == 512 && kRopeDim == 64, "FlashMLA layout requires (512, 64)");
template <typename PosT>
static constexpr auto kernel = fused_k_norm_rope_flashmla<DType, kHeadDim, kRopeDim, PosT, kLogPageSize, kUsePDL>;
static constexpr auto kernel =
fused_k_norm_rope_flashmla<DType, kHeadDim, kRopeDim, PosT, kLogPageSize, kLayout, kUsePDL>;
static void forward(
const tvm::ffi::TensorView kv,
@@ -881,4 +915,7 @@ struct FusedQIndexerRopeHadamardFp4QuantKernel {
}
};
// The JIT module names and wrappers spell the layouts as bare enumerators.
using enum deepseek_v4::KVLayout;
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/store.cuh
+137 -24
View File
@@ -8,6 +8,7 @@
#include <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
#include <sgl_kernel/deepseek_v4/kv_layout.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
@@ -15,6 +16,7 @@
#include <bit>
#include <cstdint>
#include <cuda_fp8.h>
#include <optional>
namespace sglang {
@@ -29,12 +31,21 @@ struct FusedStoreCacheParam {
uint32_t num_tokens;
};
/// Parameters of the V4.1 (fp8 / fp4) FlashMLA store; `freqs_cis` is the per-token
/// (real, imag) pairs of the 64 RoPE dims, or nullptr when the input is already rotated.
struct FusedStoreCacheV41Param {
const void* __restrict__ input;
void* __restrict__ cache;
const void* __restrict__ indices;
const float* __restrict__ freqs_cis;
uint32_t num_tokens;
};
template <typename Float, typename IndicesT, uint32_t kPageBits, bool kUsePDL>
__global__ void fused_store_flashmla_cache(const __grid_constant__ FusedStoreCacheParam param) {
using namespace device;
/// NOTE: 584 = 576 + 8
constexpr int64_t kPageBytes = host::div_ceil(584 << kPageBits, 576) * 576;
using Paged = deepseek_v4::PagedKV<deepseek_v4::KVLayout::V4, kPageBits>;
// each warp handles 64 elements, 8 warps, each block handles 1 row
const auto& [input, cache, indices, num_tokens] = param;
@@ -56,25 +67,80 @@ __global__ void fused_store_flashmla_cache(const __grid_constant__ FusedStoreCac
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 int32_t page = index >> kPageBits;
const int32_t offset = index & ((1 << kPageBits) - 1);
const auto page_ptr = pointer::offset(cache, page * kPageBytes);
const auto value_ptr = pointer::offset(page_ptr, offset * 576);
const auto scale_ptr = pointer::offset(page_ptr, 576 << kPageBits, offset * 8);
static_cast<fp8x2_e4m3_t*>(value_ptr)[tid] = result;
static_cast<uint8_t*>(scale_ptr)[wid] = scale_ue8m0;
const auto row = Paged::row(static_cast<uint8_t*>(cache), index);
reinterpret_cast<fp8x2_e4m3_t*>(row.data)[tid] = result;
row.scale[wid] = scale_ue8m0;
} else {
const auto result = cast<bf16x2_t>(elems);
const int32_t page = index >> kPageBits;
const int32_t offset = index & ((1 << kPageBits) - 1);
const auto page_ptr = pointer::offset(cache, page * kPageBytes);
const auto value_ptr = pointer::offset(page_ptr, offset * 576, 448);
static_cast<bf16x2_t*>(value_ptr)[tid - 7 * 32] = result;
const auto row = Paged::row(static_cast<uint8_t*>(cache), index);
reinterpret_cast<bf16x2_t*>(row.data + 448)[tid - 7 * 32] = result;
}
PDLTriggerSecondary<kUsePDL>();
}
/// Elements per thread of the V4.1 store (`512 / vec` threads per token): 4 for fp8 rows, 2 for
/// fp4, whose per-element IEEE divisions spread better over more threads (measured on B200, bs 1..512).
constexpr uint32_t v41_store_vec_size(deepseek_v4::KVLayout layout) {
return layout == deepseek_v4::KVLayout::V41 ? 4 : 2;
}
template <
typename Float,
typename IndicesT,
uint32_t kPageBits,
deepseek_v4::KVLayout kLayout,
bool kRope,
bool kUsePDL>
__global__ void fused_store_flashmla_cache_v41(const __grid_constant__ FusedStoreCacheV41Param param) {
using namespace device;
using Paged = deepseek_v4::PagedKV<kLayout, kPageBits>;
static_assert(kLayout != deepseek_v4::KVLayout::V4, "the V4 layout has its own kernel above");
constexpr uint32_t kVecSize = v41_store_vec_size(kLayout);
constexpr uint32_t kNopeLanes = (512 - 64) / kVecSize; // threads from here on hold the RoPE tail
using Packed = packed_t<Float>;
using Vec = AlignedVector<Packed, kVecSize / 2>;
const auto& [input, cache, indices, freqs_cis, num_tokens] = param;
const uint32_t bid = blockIdx.x;
const uint32_t tid = threadIdx.x;
PDLWaitPrimary<kUsePDL>();
const auto index = static_cast<const IndicesT*>(indices)[bid];
Vec elems;
elems.load(static_cast<const Float*>(input) + bid * 512, tid);
float v[kVecSize];
#pragma unroll
for (uint32_t i = 0; i < kVecSize / 2; ++i) {
const auto [x, y] = cast<fp32x2_t>(elems[i]);
v[2 * i] = x;
v[2 * i + 1] = y;
}
if constexpr (kRope) {
if (tid >= kNopeLanes) {
// (real, imag) pairs of the tail, rotated and rounded back to the input dtype as `rope_tail` does.
AlignedVector<float, kVecSize> freq;
freq.load(freqs_cis + bid * 64, tid - kNopeLanes);
#pragma unroll
for (uint32_t i = 0; i < kVecSize / 2; ++i) {
const auto x = v[2 * i];
const auto y = v[2 * i + 1];
const auto rotated = cast<fp32x2_t>(
cast<Packed>(fp32x2_t{x * freq[2 * i] - y * freq[2 * i + 1], x * freq[2 * i + 1] + y * freq[2 * i]}));
v[2 * i] = rotated.x;
v[2 * i + 1] = rotated.y;
}
}
}
const auto row = Paged::row(static_cast<uint8_t*>(cache), index);
deepseek_v4::v41::store_row<kLayout>(row.data, row.scale, tid, v);
PDLTriggerSecondary<kUsePDL>();
}
template <typename Float, typename IndicesT, uint32_t kPageBits, bool kUsePDL>
__global__ void fused_store_indexer_cache(const __grid_constant__ FusedStoreCacheParam param) {
using namespace device;
@@ -120,16 +186,41 @@ __global__ void fused_store_indexer_cache(const __grid_constant__ FusedStoreCach
PDLTriggerSecondary<kUsePDL>();
}
template <typename Float, typename IndicesT, uint32_t kPageSize, bool kUsePDL>
template <typename Float, typename IndicesT, uint32_t kPageSize, deepseek_v4::KVLayout kLayout, bool kUsePDL>
struct FusedStoreCacheFlashMLAKernel {
static constexpr int32_t kLogSize = std::countr_zero(kPageSize);
static constexpr int64_t kPageBytes = host::div_ceil(584 * kPageSize, 576) * 576;
static constexpr auto kernel = fused_store_flashmla_cache<Float, IndicesT, kLogSize, kUsePDL>;
static constexpr bool kIsV4 = kLayout == deepseek_v4::KVLayout::V4;
static constexpr int64_t kPageBytes = deepseek_v4::kv_page_bytes<kLayout>(kPageSize);
static_assert(!kIsV4 || 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 << kLogSize == kPageSize);
template <bool kRope>
static constexpr auto v41_kernel = fused_store_flashmla_cache_v41<Float, IndicesT, kLogSize, kLayout, kRope, kUsePDL>;
/// Store rows that are already normed and rotated.
static void run(tvm::ffi::TensorView input, tvm::ffi::TensorView cache, tvm::ffi::TensorView indices) {
launch(input, cache, indices, std::nullopt);
}
/// V4.1 layouts only: rotate the RoPE tail in-kernel with the per-token `freqs_cis`
/// (`[num_tokens, 64]` fp32, real / imag interleaved) before quantizing.
static void run_rope(
tvm::ffi::TensorView input,
tvm::ffi::TensorView cache,
tvm::ffi::TensorView indices,
tvm::ffi::TensorView freqs_cis) {
static_assert(!kIsV4, "the V4 layout keeps its RoPE dims in bf16 and has no in-kernel RoPE");
launch(input, cache, indices, freqs_cis);
}
private:
static void launch(
tvm::ffi::TensorView input,
tvm::ffi::TensorView cache,
tvm::ffi::TensorView indices,
std::optional<tvm::ffi::TensorView> freqs_cis) {
using namespace host;
auto N = SymbolicSize{"num_tokens"};
@@ -148,16 +239,35 @@ struct FusedStoreCacheFlashMLAKernel {
.with_dtype<IndicesT>()
.with_device(device_)
.verify(indices);
if (freqs_cis.has_value()) {
// Real / imag interleaved, so the trailing dim is 64, not 32.
TensorMatcher({N, 64}).with_dtype<float>().with_device(device_).verify(*freqs_cis);
}
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
const auto params = FusedStoreCacheParam{
.input = input.data_ptr(),
.cache = cache.data_ptr(),
.indices = indices.data_ptr(),
.num_tokens = num_tokens,
};
if (num_tokens == 0) return;
const auto kBlockSize = 256;
const auto num_blocks = num_tokens;
LaunchKernel(num_blocks, kBlockSize, device_.unwrap()).enable_pdl(kUsePDL)(kernel, params);
if constexpr (kIsV4) {
RuntimeCheck(!freqs_cis.has_value(), "the V4 layout has no in-kernel RoPE");
const auto params = FusedStoreCacheParam{
.input = input.data_ptr(),
.cache = cache.data_ptr(),
.indices = indices.data_ptr(),
.num_tokens = num_tokens,
};
constexpr auto kernel = fused_store_flashmla_cache<Float, IndicesT, kLogSize, kUsePDL>;
LaunchKernel(num_blocks, kBlockSize, device_.unwrap()).enable_pdl(kUsePDL)(kernel, params);
} else {
const auto params = FusedStoreCacheV41Param{
.input = input.data_ptr(),
.cache = cache.data_ptr(),
.indices = indices.data_ptr(),
.freqs_cis = freqs_cis.has_value() ? static_cast<const float*>(freqs_cis->data_ptr()) : nullptr,
.num_tokens = num_tokens,
};
const auto kernel = freqs_cis.has_value() ? v41_kernel<true> : v41_kernel<false>;
LaunchKernel(num_blocks, 512 / v41_store_vec_size(kLayout), device_.unwrap()).enable_pdl(kUsePDL)(kernel, params);
}
}
};
@@ -202,4 +312,7 @@ struct FusedStoreCacheIndexerKernel {
}
};
// The JIT module names and wrappers spell the layouts as bare enumerators.
using enum deepseek_v4::KVLayout;
} // namespace sglang
python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/fp4_utils.cuh
+60
View File
@@ -0,0 +1,60 @@
#pragma once
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
#ifndef USE_ROCM
#include <cuda_fp4.h>
#endif
// FP4 (e2m1) helpers: per-32 UE8M0 for the indexer, per-16 E4M3 for compressed KV.
namespace sglang {
namespace deepseek_v4::fp4 {
/// Largest finite e2m1 value.
constexpr float kMax = 6.0f;
/// `6 * 2^-126`, the amax floor `torch_quant.fake_quant_fp4` clamps to.
constexpr float kAmaxFloor = 6.0f * 1.1754943508222875e-38f;
/// Elements sharing one ue8m0 scale.
constexpr uint32_t kBlockSize = 32;
/// Compressed-KV elements sharing one E4M3 scale.
constexpr uint32_t kCompressedKVBlockSize = 16;
/// \brief Round amax / 6 to a positive finite E4M3 scale, ties to even.
SGL_DEVICE float compressed_kv_scale(float amax) {
const auto raw = fminf(fmaxf(amax * (1.0f / kMax), 0x1p-9f), 448.0f);
return static_cast<float>(__nv_fp8_e4m3(raw));
}
/// \brief Quantize compressed KV with its E4M3 scale and return dequantized values.
SGL_DEVICE fp32x2_t fake_quant_compressed_kv_x2(fp32x2_t x, float scale) {
const fp32x2_t scaled{__fdiv_rn(x.x, scale) + 0.0f, __fdiv_rn(x.y, scale) + 0.0f};
const auto code = __nv_cvt_float2_to_fp4x2(scaled, __NV_E2M1, cudaRoundNearest);
const auto grid = device::cast<fp32x2_t>(fp16x2_t{__nv_cvt_fp4x2_to_halfraw2(code, __NV_E2M1)});
return {grid.x * scale, grid.y * scale};
}
/// \brief Per-block ue8m0 scale and its reciprocal, from the block's absmax.
SGL_DEVICE fp32x2_t block_scale(float amax) {
const auto exponent = fp8::cast_to_ue8m0(fmaxf(amax, kAmaxFloor) * (1.0f / kMax));
return {__uint_as_float(static_cast<uint32_t>(exponent) << 23), fp8::inv_scale_ue8m0(exponent)};
}
/// \brief Round a pair onto the e2m1 grid and back, through `scale`.
///
/// Every e2m1 value is exact in fp16, so the roundtrip is lossless. Adding `0.0f` during
/// scaling clears negative zero to match `torch.sign(0) == 0` in `torch_quant.round_fp4`.
SGL_DEVICE fp32x2_t fake_quant_x2(fp32x2_t x, float scale, float inv_scale) {
const fp32x2_t scaled{__fmaf_rn(x.x, inv_scale, 0.0f), __fmaf_rn(x.y, inv_scale, 0.0f)};
const auto code = __nv_cvt_float2_to_fp4x2(scaled, __NV_E2M1, cudaRoundNearest);
const auto grid = device::cast<fp32x2_t>(fp16x2_t{__nv_cvt_fp4x2_to_halfraw2(code, __NV_E2M1)});
return {grid.x * scale, grid.y * scale};
}
} // namespace deepseek_v4::fp4
} // namespace sglang
python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/kv_layout.cuh
+212
View File
@@ -0,0 +1,212 @@
#pragma once
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
#include <cstdint>
#ifndef USE_ROCM
#include <cuda_fp4.h>
#include <cuda_fp8.h>
#endif
// Paged fp8 / fp4 KV cache layouts read by the d_qk = 512 sparse MLA decode
// kernels. A page block is `page_size` data rows followed by `page_size` scale
// rows, so the scale region starts at byte `page_size * kDataBytes`:
//
// V4 584 B/token: 448 e4m3 + 64 bf16 (RoPE) data, 7 ue8m0 scales + 1 pad,
// one scale per 64 e4m3 values.
// V41 528 B/token: 512 e4m3 data (the RoPE dims are quantized too),
// 16 ue8m0 scales, one per 32 values.
// V41_FP4 288 B/token: 512 e2m1 data packed two per byte (even index in the
// low nibble), 32 e4m3 scales, one per 16 values.
//
// The reader requires the rows of a page to be contiguous and the page stride
// to be a multiple of kPageAlign (its TMA row stride), which is what
// kv_page_bytes pads to. The pure-torch reference of the V41 (fp8) quantizer is
// `sglang.kernels.ops.attention.dsv4.torch_quant`.
namespace sglang {
namespace deepseek_v4 {
enum class KVLayout : int32_t { V4 = 0, V41 = 1, V41_FP4 = 2 };
template <KVLayout kLayout>
struct KVLayoutTraits;
template <>
struct KVLayoutTraits<KVLayout::V4> {
static constexpr int64_t kDataBytes = 576;
static constexpr int64_t kScaleBytes = 8;
static constexpr int64_t kTileSize = 64;
static constexpr int64_t kPageAlign = 576;
static constexpr int64_t kBytesPerToken = kDataBytes + kScaleBytes;
};
template <>
struct KVLayoutTraits<KVLayout::V41> {
static constexpr int64_t kDataBytes = 512;
static constexpr int64_t kScaleBytes = 16;
static constexpr int64_t kTileSize = 32;
static constexpr int64_t kPageAlign = 512;
static constexpr int64_t kBytesPerToken = kDataBytes + kScaleBytes;
};
template <>
struct KVLayoutTraits<KVLayout::V41_FP4> {
static constexpr int64_t kDataBytes = 256;
static constexpr int64_t kScaleBytes = 32;
static constexpr int64_t kTileSize = 16;
static constexpr int64_t kPageAlign = 256;
static constexpr int64_t kBytesPerToken = kDataBytes + kScaleBytes;
};
/// Bytes of one page block: `page_size` tokens, padded up to the reader's row stride.
template <KVLayout kLayout>
constexpr int64_t kv_page_bytes(int64_t page_size) {
using Traits = KVLayoutTraits<kLayout>;
return (page_size * Traits::kBytesPerToken + Traits::kPageAlign - 1) / Traits::kPageAlign * Traits::kPageAlign;
}
/// Addressing of a paged cache: `1 << kPageBits` data rows then as many scale rows, padded to kPageAlign.
template <KVLayout kLayout, uint32_t kPageBits>
struct PagedKV {
using Traits = KVLayoutTraits<kLayout>;
static constexpr int64_t kPageSize = int64_t{1} << kPageBits;
static constexpr int64_t kPageBytes = kv_page_bytes<kLayout>(kPageSize);
/// Byte offset of the scale rows inside a page.
static constexpr int64_t kScaleBase = Traits::kDataBytes << kPageBits;
template <typename LocT>
static constexpr LocT page_of(LocT loc) {
return loc >> kPageBits;
}
template <typename LocT>
static constexpr LocT slot_of(LocT loc) {
return loc & (static_cast<LocT>(kPageSize) - 1);
}
/// Byte offsets of token `loc`'s data row and scale row from the cache base.
template <typename LocT>
static constexpr int64_t data_offset(LocT loc) {
return page_of(loc) * kPageBytes + slot_of(loc) * Traits::kDataBytes;
}
template <typename LocT>
static constexpr int64_t scale_offset(LocT loc) {
return page_of(loc) * kPageBytes + kScaleBase + slot_of(loc) * Traits::kScaleBytes;
}
struct Row {
uint8_t* data;
uint8_t* scale;
};
/// The data row and scale row of token `loc`.
template <typename LocT>
SGL_DEVICE static Row row(uint8_t* cache, LocT loc) {
uint8_t* page = cache + page_of(loc) * kPageBytes;
return {page + slot_of(loc) * Traits::kDataBytes, page + kScaleBase + slot_of(loc) * Traits::kScaleBytes};
}
};
#ifndef USE_ROCM
namespace v41 {
/// The row helpers below quantize one 512-wide token spread over `512 / kVecSize` threads,
/// thread `tx` holding elements `[kVecSize * tx, kVecSize * (tx + 1))` as fp32. They are
/// warp-collective: every thread of the token must call them together. NaN / inf are not handled.
/// Per-thread |max| over the vector.
template <uint32_t kVecSize>
SGL_DEVICE float vec_amax(const float (&v)[kVecSize]) {
float amax = fabsf(v[0]);
#pragma unroll
for (uint32_t i = 1; i < kVecSize; ++i) {
amax = fmaxf(amax, fabsf(v[i]));
}
return amax;
}
/// V4.1 fp8 row: one ue8m0 scale per 32-element tile. `data_row` is the token's 512 B,
/// `scale_row` its 16 scale bytes. Scale `2^ceil(log2(max(amax / 448, 1e-4)))` stored as the
/// ue8m0 byte, payload `e4m3(x / scale)` rounded to nearest even (|x / scale| <= 448, so it
/// never saturates).
template <uint32_t kVecSize>
SGL_DEVICE void store_row_fp8(uint8_t* data_row, uint8_t* scale_row, uint32_t tx, const float (&v)[kVecSize]) {
using namespace device;
constexpr uint32_t kTileLanes = KVLayoutTraits<KVLayout::V41>::kTileSize / kVecSize;
static_assert(kVecSize % 2 == 0 && kTileLanes >= 1 && (kTileLanes & (kTileLanes - 1)) == 0);
const float amax = warp::reduce_max<kTileLanes>(vec_amax(v));
// ceil(log2(max(amax / 448, 1e-4))) exactly, from the bits of amax: 448 = 1.75 * 2^8, so the
// quotient's exponent is amax's minus 8, plus one when the mantissa exceeds 1.75, floored at 2^-13.
const uint32_t bits = __float_as_uint(amax);
const int32_t exponent =
max(static_cast<int32_t>(bits >> 23) - 8 + static_cast<int32_t>((bits & 0x7FFFFFu) > 0x600000u), 114);
// The scale is a power of two, so the multiply by its reciprocal is the exact quotient.
const float inv_scale = fp8::inv_scale_ue8m0(exponent);
AlignedVector<fp8x2_e4m3_t, kVecSize / 2> out;
#pragma unroll
for (uint32_t i = 0; i < kVecSize / 2; ++i) {
// `cvt.rn.satfinite.e4m3x2` directly: |x / scale| <= 448 needs no clamp.
out[i] = fp8x2_e4m3_t{fp32x2_t{v[2 * i] * inv_scale, v[2 * i + 1] * inv_scale}};
}
out.store(data_row, tx);
// Every lane of the tile holds the exponent; they all store the same byte.
scale_row[tx / kTileLanes] = static_cast<uint8_t>(exponent);
}
/// V4.1 fp4 row: one e4m3 scale per 16-element tile. `data_row` is the token's 256 B,
/// `scale_row` its 32 scale bytes. Scale `e4m3(clamp(amax / 6, 2^-9, 448))` rounded to
/// nearest even; codes `cvt.rn.satfinite.e2m1x2` of `x / scale` (ties to even, saturating
/// at 6, the sign kept for a value that rounds to zero), the even element in the low nibble.
template <uint32_t kVecSize>
SGL_DEVICE void store_row_fp4(uint8_t* data_row, uint8_t* scale_row, uint32_t tx, const float (&v)[kVecSize]) {
using namespace device;
constexpr uint32_t kTileLanes = KVLayoutTraits<KVLayout::V41_FP4>::kTileSize / kVecSize;
static_assert(kVecSize % 2 == 0 && kTileLanes >= 1 && (kTileLanes & (kTileLanes - 1)) == 0);
const float amax = warp::reduce_max<kTileLanes>(vec_amax(v));
const __nv_fp8_e4m3 scale_e4m3{fminf(fmaxf(__fdiv_rn(amax, 6.0f), 0x1p-9f), 448.0f)};
const float scale = static_cast<float>(scale_e4m3);
AlignedVector<uint8_t, kVecSize / 2> out;
#pragma unroll
for (uint32_t i = 0; i < kVecSize / 2; ++i) {
// IEEE division by the rounded scale, as the reference does; a reciprocal multiply could cross an e2m1 tie.
out[i] = static_cast<uint8_t>(__nv_cvt_float2_to_fp4x2(
fp32x2_t{__fdiv_rn(v[2 * i], scale), __fdiv_rn(v[2 * i + 1], scale)}, __NV_E2M1, cudaRoundNearest));
}
out.store(data_row, tx);
// Every lane of the tile holds the scale; they all store the same byte.
scale_row[tx / kTileLanes] = scale_e4m3.__x;
}
/// Dispatch on the layout for a 512-wide row held as `kVecSize` consecutive fp32 per thread.
/// V4 has no row helper here: its writers keep their nope / RoPE split code.
template <KVLayout kLayout, uint32_t kVecSize>
SGL_DEVICE void store_row(uint8_t* data_row, uint8_t* scale_row, uint32_t tx, const float (&v)[kVecSize]) {
static_assert(kLayout != KVLayout::V4, "V4 rows are written by the caller");
if constexpr (kLayout == KVLayout::V41) {
store_row_fp8(data_row, scale_row, tx, v);
} else {
store_row_fp4(data_row, scale_row, tx, v);
}
}
/// Same, for a row kept in an `AlignedVector<float, kVecSize>`.
template <KVLayout kLayout, std::size_t kVecSize>
SGL_DEVICE void
store_row(uint8_t* data_row, uint8_t* scale_row, uint32_t tx, const device::AlignedVector<float, kVecSize>& v) {
store_row<kLayout>(data_row, scale_row, tx, *reinterpret_cast<const float (*)[kVecSize]>(v.data()));
}
} // namespace v41
#endif // USE_ROCM
} // namespace deepseek_v4
} // namespace sglang
python/sglang/kernels/ops/attention/dsv4/attn.py
+43 -6
View File
@@ -1,4 +1,4 @@
from typing import Literal, Tuple
from typing import Literal, Optional, Tuple, Union
import torch
import triton
@@ -12,6 +12,7 @@ from sglang.kernels.jit.utils import (
make_cpp_args,
)
from .kv_layout import KVLayout
from .utils import make_name
@@ -30,20 +31,32 @@ def _jit_fused_store_module(
input_dtype: torch.dtype,
index_dtype: torch.dtype,
page_size: int,
layout: KVLayout,
):
args = make_cpp_args(input_dtype, index_dtype, page_size, is_arch_support_pdl())
cname = "FlashMLA" if name == "flashmla" else "Indexer"
if name == "flashmla":
args = make_cpp_args(
input_dtype, index_dtype, page_size, layout.cpp_name, is_arch_support_pdl()
)
# The V4 layout keeps its RoPE dims in bf16 and has no in-kernel RoPE.
cname = "FlashMLA"
wrappers = ["run"] if layout is KVLayout.V4 else ["run", "run_rope"]
else:
assert layout is KVLayout.V4, "only the FlashMLA cache has V4.1 layouts"
args = make_cpp_args(input_dtype, index_dtype, page_size, is_arch_support_pdl())
cname, wrappers = "Indexer", ["run"]
kernel_class = f"FusedStoreCache{cname}Kernel<{args}>"
return load_jit(
make_name("store_" + name),
*args,
cuda_files=["deepseek_v4/store.cuh"],
cuda_wrappers=[("run", f"{kernel_class}::run")],
cuda_wrappers=[(w, f"{kernel_class}::{w}") for w in wrappers],
)
def get_paged_mqa_logits_metadata(seq_lens: torch.Tensor, page_size: int, num_sm: int):
assert page_size == 64
# The schedule only depends on the sequence lengths (256-token splits), not
# on the page size.
assert page_size in (64, 128), page_size
seq_lens = seq_lens.view(-1).to(torch.int32)
bs = int(seq_lens.shape[0])
metadata = seq_lens.new_empty(num_sm + 1, 2)
@@ -67,8 +80,25 @@ def fused_store_cache(
*,
page_size: int,
type: Literal["flashmla", "indexer"],
layout: Union[KVLayout, str] = KVLayout.V4,
freqs_cis: Optional[torch.Tensor] = None,
) -> None:
"""Quantize ``input`` ``[num_tokens, 512]`` (bf16, normed and rotated) into the
paged cache at ``indices``.
:param layout: the cache's :class:`KVLayout`. ``V4`` is the 584-byte layout
(fp8 nope, bf16 rope); ``V41`` (528 B) and ``V41_FP4`` (288 B) are the
V4.1 formats, fp8 with per-32 ue8m0 scales and e2m1 with per-16 e4m3
scales over all 512 dims.
:param freqs_cis: V4.1 layouts only. ``[num_tokens, 32]`` complex or
``[num_tokens, 64]`` fp32 (real / imag interleaved); rotates the 64-dim
RoPE tail in-kernel, so ``input`` must then be the un-rotated latent.
"""
layout = KVLayout.parse(layout)
if is_hip_runtime():
assert layout is KVLayout.V4 and freqs_cis is None, (
"the V4.1 KV layouts are CUDA (sm100) only"
)
from sglang.kernels.ops.kvcache.triton_store_cache import (
triton_fused_store_cache,
)
@@ -80,8 +110,15 @@ def fused_store_cache(
input_dtype=input.dtype,
index_dtype=indices.dtype,
page_size=page_size,
layout=layout,
)
module.run(input, cache, indices)
if freqs_cis is None:
module.run(input, cache, indices)
else:
assert layout is not KVLayout.V4, "the V4 layout has no in-kernel RoPE"
if freqs_cis.is_complex():
freqs_cis = torch.view_as_real(freqs_cis).flatten(-2)
module.run_rope(input, cache, indices, freqs_cis.contiguous())
@triton.jit
python/sglang/kernels/ops/attention/dsv4/compress.py
+14 -2
View File
@@ -16,6 +16,7 @@ from sglang.srt.layers.attention.dsa.utils import (
)
from sglang.srt.utils import is_hip, is_xpu
from .kv_layout import KVLayout
from .utils import make_name
_is_xpu = is_xpu()
@@ -48,7 +49,8 @@ def _jit_compress_norm_rope_module(
head_dim: int,
rope_dim: int,
page_size: int,
bf16_store: bool = False,
bf16_store: bool,
layout: KVLayout,
fp8_2buff: bool = False,
) -> Module:
args = make_cpp_args(
@@ -56,9 +58,10 @@ def _jit_compress_norm_rope_module(
head_dim,
rope_dim,
page_size,
is_arch_support_pdl(),
INDEXER_K_CACHE_PRESHUFFLE_TILE if aiter_can_use_preshuffle_paged_mqa() else 0,
bf16_store,
layout.cpp_name,
is_arch_support_pdl(),
)
cuda_wrappers = [("forward", f"FusedNormRopeKernel<{args}>::forward")]
if head_dim == 128:
@@ -455,9 +458,16 @@ def compress_norm_rope_store(
kvcache_scale: Optional[torch.Tensor] = None,
rope_cache: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
fp4_k_write_metadata=None,
layout: Union[KVLayout, str] = KVLayout.V4,
fp8_2buff: bool = False,
kvcache_rope: Optional[torch.Tensor] = None,
) -> None:
layout = KVLayout.parse(layout)
if layout is not KVLayout.V4:
assert kv.shape[-1] == 512 and not use_fp4 and not bf16_store, (
"the V4.1 layouts are paged FlashMLA main-KV caches"
)
assert not is_hip() and not _is_xpu, "the V4.1 KV layouts are CUDA (sm100) only"
if use_fp4:
assert kv.shape[-1] == 128
if is_hip() and use_fp4:
@@ -482,6 +492,7 @@ def compress_norm_rope_store(
if fp8_2buff:
assert not (use_fp4 or bf16_store), "fp8 two-pool store is its own layout"
assert layout is KVLayout.V4, "fp8 two-pool store is a V4 (584 B page) cache"
assert kv.shape[-1] != 128, "fp8 two-pool store is the latent, not the indexer"
assert kvcache_rope is not None, "fp8 two-pool store needs the rope pool"
assert not _is_xpu, "fp8 two-pool store is only wired for the CUDA/HIP kernel"
@@ -507,6 +518,7 @@ def compress_norm_rope_store(
freq_cis.shape[-1],
page_size,
bf16_store,
layout,
fp8_2buff,
)
if use_fp4:
python/sglang/kernels/ops/attention/dsv4/dequant_k_cache.py
+170 -1
View File
@@ -1,9 +1,11 @@
from typing import Optional
from typing import Optional, Union
import torch
import triton
import triton.language as tl
from sglang.kernels.jit.utils import get_jit_cuda_arch, is_hip_runtime
from sglang.kernels.ops.attention.dsv4.kv_layout import KVLayout
from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz
fp8_dtype = torch.float8_e4m3fnuz if is_fp8_fnuz() else torch.float8_e4m3fn
@@ -26,6 +28,7 @@ def dequantize_k_cache_paged(
page_table_1_flattened: torch.Tensor,
page_size: int,
out: Optional[torch.Tensor] = None,
layout: Union[KVLayout, str] = KVLayout.V4,
) -> torch.Tensor:
"""Dequantize the DeepSeek v4 paged KV cache for a list of token IDs.
@@ -36,10 +39,16 @@ def dequantize_k_cache_paged(
out: optional (num_tokens, 1, DIM_NOPE + DIM_ROPE) bf16 destination.
May be a slice of a larger workspace; the kernel uses out.stride(0)
so contiguous-along-dim-0 slices work.
layout: the cache's :class:`KVLayout`.
Returns:
(num_tokens, 1, DIM_NOPE + DIM_ROPE) bfloat16.
"""
layout = KVLayout.parse(layout)
if layout is not KVLayout.V4:
return dequantize_k_cache_paged_v41(
quant_k_cache, page_table_1_flattened, page_size, out=out, layout=layout
)
assert quant_k_cache.is_contiguous()
assert page_table_1_flattened.dtype in (torch.int32, torch.int64)
@@ -85,6 +94,69 @@ def dequantize_k_cache_paged(
return out
def dequantize_k_cache_paged_v41(
quant_k_cache: torch.Tensor,
page_table_1_flattened: torch.Tensor,
page_size: int,
out: Optional[torch.Tensor] = None,
layout: KVLayout = KVLayout.V41,
) -> torch.Tensor:
"""Dequantize a V4.1 paged cache (fp8 ``V41`` or fp4 ``V41_FP4``) for a list
of token IDs into ``(num_tokens, 1, 512)`` bf16.
Bit-exact with the pure-torch dequantizer of these formats.
"""
layout = KVLayout.parse(layout)
assert layout in (KVLayout.V41, KVLayout.V41_FP4), layout
if is_hip_runtime() or get_jit_cuda_arch().major < 10:
raise RuntimeError(
"DeepSeek V4.1 KV cache dequantization requires CUDA SM100 or newer"
)
assert quant_k_cache.is_contiguous()
assert page_table_1_flattened.dtype in (torch.int32, torch.int64)
quant_k_cache_u8 = quant_k_cache.view(torch.uint8)
num_tokens = page_table_1_flattened.shape[0]
bytes_per_page = quant_k_cache_u8.shape[-1]
assert bytes_per_page >= page_size * layout.bytes_per_token, (
f"{bytes_per_page=} cannot hold {page_size} tokens of {layout}"
)
buf_fp8 = quant_k_cache_u8.view(fp8_dtype).reshape(-1)
buf_uint8 = quant_k_cache_u8.reshape(-1)
if out is None:
out = torch.empty(
(num_tokens, 1, DIM_NOPE + DIM_ROPE),
dtype=torch.bfloat16,
device=quant_k_cache.device,
)
else:
assert out.shape == (num_tokens, 1, DIM_NOPE + DIM_ROPE)
assert out.dtype == torch.bfloat16
if num_tokens == 0:
return out
kernel = (
_dequantize_k_cache_paged_v41_fp8_kernel
if layout is KVLayout.V41
else _dequantize_k_cache_paged_v41_fp4_kernel
)
kernel[(num_tokens,)](
out,
buf_fp8,
buf_uint8,
page_table_1_flattened,
out.stride(0),
BYTES_PER_PAGE=bytes_per_page,
PAGE_SIZE=page_size,
DATA_BYTES=layout.data_bytes,
SCALE_BYTES=layout.scale_bytes,
TILE_SIZE=layout.tile_size,
S_OFFSET_BYTES=layout.scale_offset(page_size),
)
return out
def gather_dequant_requant_fp8_paged(
quant_k_cache: torch.Tensor,
page_table_1_flattened: torch.Tensor,
@@ -270,6 +342,103 @@ def _dequantize_k_cache_paged_kernel(
tl.store(output_ptr + out_row_base + DIM_NOPE + rope_offs, rope_data)
@triton.jit
def _ue8m0_to_fp32(scale_u8):
# Follows torch.float8_e8m0fnu: byte 0 is the denormal 2 ** -127, byte 255
# is NaN, otherwise 2 ** (byte - 127).
normal = (scale_u8.to(tl.int32) << 23).to(tl.float32, bitcast=True)
denormal = tl.full(scale_u8.shape, 0x00400000, tl.int32).to(
tl.float32, bitcast=True
)
scale = tl.where(scale_u8 == 0, denormal, normal)
return tl.where(scale_u8 == 255, float("nan"), scale)
@triton.jit
def _e2m1_code_to_fp32(code):
# The 4-bit e2m1 code: bit 3 is the sign, bits 0-2 index into
# [0, 0.5, 1, 1.5, 2, 3, 4, 6].
m = code & 7
e = m >> 1
f = (m & 1).to(tl.float32)
mag = tl.where(e == 0, 0.5 * f, tl.exp2((e - 1).to(tl.float32)) * (1.0 + 0.5 * f))
# Set the sign bit directly: a negated zero must stay -0.0 (code 0x8).
sign = (code & 8).to(tl.int32) << 28
return (mag.to(tl.int32, bitcast=True) | sign).to(tl.float32, bitcast=True)
@triton.jit
def _dequantize_k_cache_paged_v41_fp8_kernel(
output_ptr,
buf_fp8_ptr,
buf_uint8_ptr,
page_table_ptr,
output_stride_0,
BYTES_PER_PAGE: tl.constexpr,
PAGE_SIZE: tl.constexpr,
DATA_BYTES: tl.constexpr,
SCALE_BYTES: tl.constexpr,
TILE_SIZE: tl.constexpr,
S_OFFSET_BYTES: tl.constexpr,
):
# V41: 512 e4m3 values per token, then 16 ue8m0 scales (one per 32 values).
tl.static_assert(DATA_BYTES == 512 and SCALE_BYTES == 16 and TILE_SIZE == 32)
token_id = tl.program_id(0).to(tl.int64)
loc = tl.load(page_table_ptr + token_id).to(tl.int64)
page_idx = loc // PAGE_SIZE
in_page = loc % PAGE_SIZE
page_byte_base = page_idx * BYTES_PER_PAGE
token_data_base = page_byte_base + in_page * DATA_BYTES
token_scale_base = page_byte_base + S_OFFSET_BYTES + in_page * SCALE_BYTES
offs = tl.arange(0, DATA_BYTES)
vals = tl.load(buf_fp8_ptr + token_data_base + offs).to(tl.float32)
scale_u8 = tl.load(buf_uint8_ptr + token_scale_base + offs // TILE_SIZE)
out = vals * _ue8m0_to_fp32(scale_u8)
tl.store(
output_ptr + token_id * output_stride_0 + offs,
out.to(output_ptr.dtype.element_ty),
)
@triton.jit
def _dequantize_k_cache_paged_v41_fp4_kernel(
output_ptr,
buf_fp8_ptr,
buf_uint8_ptr,
page_table_ptr,
output_stride_0,
BYTES_PER_PAGE: tl.constexpr,
PAGE_SIZE: tl.constexpr,
DATA_BYTES: tl.constexpr,
SCALE_BYTES: tl.constexpr,
TILE_SIZE: tl.constexpr,
S_OFFSET_BYTES: tl.constexpr,
):
# V41_FP4: 512 e2m1 codes packed two per byte (even index in the low nibble),
# then 32 e4m3 scales (one per 16 values).
tl.static_assert(DATA_BYTES == 256 and SCALE_BYTES == 32 and TILE_SIZE == 16)
token_id = tl.program_id(0).to(tl.int64)
loc = tl.load(page_table_ptr + token_id).to(tl.int64)
page_idx = loc // PAGE_SIZE
in_page = loc % PAGE_SIZE
page_byte_base = page_idx * BYTES_PER_PAGE
token_data_base = page_byte_base + in_page * DATA_BYTES
token_scale_base = page_byte_base + S_OFFSET_BYTES + in_page * SCALE_BYTES
boffs = tl.arange(0, DATA_BYTES)
packed = tl.load(buf_uint8_ptr + token_data_base + boffs)
# Byte j holds elements 2j (low nibble) and 2j + 1, both in tile (2j) // 16.
scale = tl.load(buf_fp8_ptr + token_scale_base + (2 * boffs) // TILE_SIZE).to(
tl.float32
)
lo = _e2m1_code_to_fp32(packed & 0xF) * scale
hi = _e2m1_code_to_fp32(packed >> 4) * scale
out_base = output_ptr + token_id * output_stride_0
tl.store(out_base + 2 * boffs, lo.to(output_ptr.dtype.element_ty))
tl.store(out_base + 2 * boffs + 1, hi.to(output_ptr.dtype.element_ty))
@triton.jit
def _gather_dequant_requant_fp8_paged_kernel(
output_ptr,
python/sglang/kernels/ops/attention/dsv4/elementwise.py
+12 -3
View File
@@ -1,4 +1,4 @@
from typing import Optional, Tuple
from typing import Optional, Tuple, Union
import torch
@@ -10,6 +10,7 @@ from sglang.kernels.jit.utils import (
)
from sglang.srt.utils import is_hip, is_xpu
from .kv_layout import KVLayout
from .utils import make_name
_is_hip = is_hip()
@@ -55,9 +56,12 @@ def _jit_main_k_norm_rope_flashmla_module(
head_dim: int,
rope_dim: int,
page_size: int,
layout: KVLayout,
):
"""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())
args = make_cpp_args(
dtype, head_dim, rope_dim, page_size, layout.cpp_name, is_arch_support_pdl()
)
return load_jit(
make_name("main_k_norm_rope_flashmla"),
*args,
@@ -273,16 +277,21 @@ def fused_k_norm_rope_flashmla(
out_loc: torch.Tensor,
kvcache: torch.Tensor,
page_size: int,
layout: Union[KVLayout, str] = KVLayout.V4,
) -> None:
"""RMSNorm + RoPE ``kv`` and write it into the ``layout`` paged FlashMLA
cache at ``out_loc``."""
layout = KVLayout.parse(layout)
freqs_real = torch.view_as_real(freqs_cis).flatten(-2)
head_dim = kv.shape[-1]
rope_dim = freqs_real.shape[-1]
if _is_xpu:
assert layout is KVLayout.V4, "the V4.1 KV layouts are CUDA (sm100) only"
fused_k_norm_rope_flashmla_xpu(
kv, kv_weight, freqs_real, positions, out_loc, kvcache, eps, page_size
)
else:
module = _jit_main_k_norm_rope_flashmla_module(
kv.dtype, head_dim, rope_dim, page_size
kv.dtype, head_dim, rope_dim, page_size, layout
)
module.forward(kv, kv_weight, freqs_real, positions, out_loc, kvcache, eps)
python/sglang/kernels/ops/attention/dsv4/kv_layout.py
+74
View File
@@ -0,0 +1,74 @@
"""Paged fp8 / fp4 KV cache layouts of the DeepSeek-V4 family sparse MLA decode kernels.
A page block stores ``page_size`` data rows followed by ``page_size`` scale rows.
The reader selects the format from the bytes per token (the last dim of the
``(num_pages, page_size, 1, bytes_per_token)`` view) and requires the page
stride to be a multiple of its TMA row stride, which :meth:`KVLayout.page_bytes`
pads to. Mirrors ``sgl_kernel/deepseek_v4/kv_layout.cuh``.
"""
from __future__ import annotations
import enum
from typing import Union
class KVLayout(str, enum.Enum):
# 448 fp8 nope + 64 bf16 rope, 7 ue8m0 scales (+1 pad) per 64 values.
V4 = "v4"
# 512 fp8 (rope quantized too), 16 ue8m0 scales per 32 values.
V41 = "v41"
# 512 e2m1 packed two per byte (even index low nibble), 32 e4m3 scales per 16 values.
V41_FP4 = "v41_fp4"
@property
def data_bytes(self) -> int:
return {KVLayout.V4: 576, KVLayout.V41: 512, KVLayout.V41_FP4: 256}[self]
@property
def scale_bytes(self) -> int:
return {KVLayout.V4: 8, KVLayout.V41: 16, KVLayout.V41_FP4: 32}[self]
@property
def tile_size(self) -> int:
"""Values sharing one scale."""
return {KVLayout.V4: 64, KVLayout.V41: 32, KVLayout.V41_FP4: 16}[self]
@property
def bytes_per_token(self) -> int:
return self.data_bytes + self.scale_bytes
@property
def page_align(self) -> int:
"""Unit the page stride is padded to: the reader's TMA row stride."""
return {KVLayout.V4: 576, KVLayout.V41: 512, KVLayout.V41_FP4: 256}[self]
@property
def is_fp4(self) -> bool:
return self is KVLayout.V41_FP4
def page_bytes(self, page_size: int) -> int:
raw = page_size * self.bytes_per_token
return -(-raw // self.page_align) * self.page_align
def scale_offset(self, page_size: int) -> int:
"""Byte offset of the scale rows inside a page."""
return page_size * self.data_bytes
@property
def cpp_name(self) -> str:
"""The C++ enumerator, for JIT template arguments. Bare, because it is
also part of the JIT module name; the headers `using enum` it in."""
return self.name
@classmethod
def parse(cls, value: Union[str, KVLayout]) -> KVLayout:
if isinstance(value, KVLayout):
return value
return cls(str(value).lower())
def is_valid_kv_layout_pair(kv: KVLayout, extra_kv: KVLayout) -> bool:
"""The (main, extra) cache pairs the decode kernel accepts: identical layouts,
or the fp4 extra cache next to a V4.1 fp8 main cache."""
return extra_kv is kv or (kv is KVLayout.V41 and extra_kv is KVLayout.V41_FP4)
python/sglang/kernels/ops/attention/dsv4/low_ratio_compress.py
+183
View File
@@ -0,0 +1,183 @@
"""Fused ratio-1 and ratio-2 decode compressors: RMSNorm, RoPE and the FlashMLA
cache write in one launch.
Ratio 1 takes the bf16 ``wkv`` projection as is: RoPE uses the token's own
position and the compressed slot equals the FULL slot. Ratio 2 pair-pools the
token against the pending partner in the state ring first. ``out_loc == 0``
marks a padded graph row on both paths, and both return the pre-RoPE latent for
the index-key projection.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Optional, Union
import torch
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
from .kv_layout import KVLayout
from .utils import make_name
if TYPE_CHECKING:
from tvm_ffi.module import Module
@cache_once
def _jit_c1_module(head_dim: int, rope_dim: int, page_size: int, layout: KVLayout):
args = make_cpp_args(
head_dim,
rope_dim,
page_size,
layout.cpp_name,
is_arch_support_pdl(),
)
return load_jit(
make_name("c1_decode"),
*args,
cuda_files=["deepseek_v4/c1.cuh"],
cuda_wrappers=[
("decode_fusion", f"FlashCompress1Kernel<{args}>::run_decode_fusion"),
],
)
def c1_decode_norm_rope_store(
kv_input: torch.Tensor,
norm_weight: torch.Tensor,
positions: torch.Tensor,
out_loc: torch.Tensor,
eps: float,
freqs_cis: torch.Tensor,
k_cache: torch.Tensor,
*,
page_size: int,
layout: Union[KVLayout, str] = KVLayout.V4,
out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""RMSNorm ``kv_input`` and write the main KV slot, in a single launch.
:param kv_input: ``[num_tokens, head_dim]`` bf16, the raw ``wkv`` projection
output.
:param norm_weight: ``[head_dim]`` bf16, promoted to fp32 for the multiply.
:param positions: ``[num_tokens]`` int32 or int64, indexed into ``freqs_cis``
as-is: at ratio 1 the latent stands for the token itself,
so there is no ``- 1``.
:param out_loc: ``[num_tokens]`` int32 or int64 ``c1_out_loc``, which at ratio 1
equals ``raw_out_loc`` (the scheduler's int64 ``out_cache_loc``).
``0`` marks a padded graph row: its latent is still computed
and published, but nothing is written to the cache.
:param eps: RMSNorm epsilon.
:param freqs_cis: ``[max_pos, rope_dim]`` fp32, real/imag interleaved --
``torch.view_as_real(freqs).flatten(-2)``.
:param k_cache: the compressed KV pool buffer for this layer.
:param page_size: slots per page of that pool (``page_size // ratio``, i.e.
the FULL page size at ratio 1).
:param layout: the pool's :class:`KVLayout`. The fp8 layouts (``V4``,
``V41``) store the fp4 fake-quantized value; ``V41_FP4``
stores the e2m1 codes themselves, rounding once.
:param out: ``[num_tokens, head_dim]`` bf16 destination for the pre-RoPE
latent. Pass a persistent buffer under CUDA graphs.
:return: ``out``, the pre-RoPE post-norm latent.
"""
num_tokens, head_dim = kv_input.shape
if out is None:
out = kv_input.new_empty((num_tokens, head_dim))
layout = KVLayout.parse(layout)
module = _jit_c1_module(head_dim, freqs_cis.shape[-1], page_size, layout)
module.decode_fusion(
kv_input,
out,
norm_weight,
freqs_cis,
positions,
out_loc,
k_cache,
float(eps),
)
return out
@cache_once
def _jit_c2_module(
head_dim: int,
rope_dim: int,
page_size: int,
layout: KVLayout,
) -> Module:
args = make_cpp_args(
head_dim,
rope_dim,
page_size,
layout.cpp_name,
is_arch_support_pdl(),
)
return load_jit(
make_name("c2_decode"),
*args,
cuda_files=["deepseek_v4/c2.cuh"],
cuda_wrappers=[
("decode_fusion", f"FlashCompress2Kernel<{args}>::run_decode_fusion"),
],
)
def c2_decode_norm_rope_store(
kv_input: torch.Tensor,
kv_state: torch.Tensor,
norm_weight: torch.Tensor,
positions: torch.Tensor,
req: torch.Tensor,
raw_out_loc: torch.Tensor,
eps: float,
freqs_cis: torch.Tensor,
k_cache: torch.Tensor,
*,
page_size: int,
ring_size: int,
draft_len: int = 1,
layout: Union[KVLayout, str] = KVLayout.V4,
out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Pair-pool ``kv_input`` against ``kv_state``, RMSNorm, and write the main KV slot.
The cache store uses ``raw_out_loc // 2`` as its slot.
:param freqs_cis: ``[max_pos, rope_dim]`` fp32, real/imag interleaved --
``torch.view_as_real(freqs).flatten(-2)``. Indexed
in-kernel at ``positions - 1``, the position the latent
stands for.
:param k_cache: the compressed KV pool buffer for this layer.
:param page_size: slots per page of that pool (``page_size // ratio``).
:param layout: the pool's :class:`KVLayout`. The fp8 layouts (``V4``,
``V41``) store the fp4 fake-quantized value; ``V41_FP4``
stores the e2m1 codes themselves, rounding once.
"""
num_tokens, fused_dim = kv_input.shape
head_dim = fused_dim // 2
if out is None:
out = kv_input.new_empty((num_tokens, head_dim), dtype=torch.bfloat16)
layout = KVLayout.parse(layout)
module = _jit_c2_module(head_dim, freqs_cis.shape[-1], page_size, layout)
module.decode_fusion(
kv_input,
kv_state,
out,
norm_weight,
positions,
req,
raw_out_loc,
eps,
freqs_cis,
k_cache,
ring_size,
draft_len,
)
return out
python/sglang/kernels/ops/attention/dsv4/metadata_kernel.py
+126 -1
View File
@@ -1,10 +1,70 @@
from typing import Optional, Tuple
from typing import NamedTuple, Optional, Tuple
import torch
import triton
import triton.language as tl
@triton.jit
def _fill_all_compressed_indices_kernel(
page_table,
seq_lens,
page_indices,
raw_indices,
PAGE_STRIDE: tl.constexpr,
TOPK: tl.constexpr,
RATIO: tl.constexpr,
PAGE_SIZE: tl.constexpr,
BLOCK: tl.constexpr,
):
row = tl.program_id(0)
positions = tl.arange(0, BLOCK)
length = tl.load(seq_lens + row)
valid = (positions < length) & (positions < TOPK)
slots_per_page = PAGE_SIZE // RATIO
pages = tl.load(
page_table + row * PAGE_STRIDE + positions // slots_per_page,
mask=valid,
other=0,
)
slots = pages * slots_per_page + positions % slots_per_page
tl.store(
page_indices + row * TOPK + positions,
tl.where(valid, slots, -1),
positions < TOPK,
)
if raw_indices is not None:
tl.store(
raw_indices + row * TOPK + positions,
tl.where(valid, positions, -1),
positions < TOPK,
)
def fill_all_compressed_indices(
page_table: torch.Tensor,
compressed_seq_lens: torch.Tensor,
page_indices: torch.Tensor,
*,
compress_ratio: int,
page_size: int,
raw_indices: Optional[torch.Tensor] = None,
) -> None:
"""Fill all reachable slots; the caller guarantees compressed length <= top-k."""
topk = page_indices.shape[1]
_fill_all_compressed_indices_kernel[(compressed_seq_lens.numel(),)](
page_table,
compressed_seq_lens,
page_indices,
raw_indices,
page_table.stride(0),
topk,
compress_ratio,
page_size,
triton.next_power_of_2(topk),
)
@triton.jit(do_not_specialize=["bs", "num_write_tokens", "c128_cur_max_seq_len"])
def _init_compressed_attn_metadata_kernel(
seq_lens_ptr,
@@ -215,3 +275,68 @@ def init_compression_metadata(
page_size,
compute_page_indices,
)
@triton.jit
def _low_ratio_metadata(
LENS,
LOC,
OUT1,
LEN1,
SPARSE1,
PAGE1,
OUT2,
LEN2,
SPARSE2,
PAGE2,
TOPK: tl.constexpr,
PADDED: tl.constexpr,
BLOCK: tl.constexpr,
):
row = tl.program_id(0)
length = tl.load(LENS + row).to(tl.int32)
loc = tl.load(LOC + row).to(tl.int64)
len1, len2 = tl.maximum(length, 1), tl.maximum(length >> 1, 1)
tl.store(OUT1 + row, loc)
tl.store(OUT2 + row, tl.where((length & 1) == 0, loc >> 1, -1))
tl.store(LEN1 + row, len1)
tl.store(LEN2 + row, len2)
tl.store(SPARSE1 + row, tl.minimum(len1, TOPK))
tl.store(SPARSE2 + row, tl.minimum(len2, TOPK))
cols = tl.arange(0, BLOCK)
tl.store(PAGE1 + row * PADDED + cols, -1, cols < PADDED)
tl.store(PAGE2 + row * PADDED + cols, -1, cols < PADDED)
class LowRatioMetadata(NamedTuple):
"""Per-request slots and lengths of the ratio-1 and ratio-2 compressed caches."""
c1_out_loc: torch.Tensor
c1_seq_lens: torch.Tensor
c1_sparse_lens: torch.Tensor
c1_page_indices: torch.Tensor
c2_out_loc: torch.Tensor
c2_seq_lens: torch.Tensor
c2_sparse_lens: torch.Tensor
c2_page_indices: torch.Tensor
def build_low_ratio_metadata(seq_lens, out_loc, topk) -> LowRatioMetadata:
assert seq_lens.numel() == out_loc.numel()
rows = seq_lens.numel()
kw = dict(device=seq_lens.device, dtype=torch.int32)
padded = triton.cdiv(topk, 64) * 64
outputs = []
for _ in range(2):
outputs.extend(
[
torch.empty(rows, device=out_loc.device, dtype=torch.int64),
torch.empty(rows, **kw),
torch.empty(rows, **kw),
torch.empty((rows, padded), **kw),
]
)
_low_ratio_metadata[(rows,)](
seq_lens, out_loc, *outputs, topk, padded, triton.next_power_of_2(padded)
)
return LowRatioMetadata(*outputs)
python/sglang/kernels/ops/attention/dsv4/torch_quant.py
+118
View File
@@ -0,0 +1,118 @@
"""Pure-torch FP4 fake quantization for DeepSeek-V4.1.
Indexer values use per-32 UE8M0 scales; compressed KV uses per-16 E4M3 scales.
Both paths round to the E2M1 grid with ties to even.
"""
from typing import Optional
import torch
FP8_MAX = 448.0
FP4_MAX = 6.0
FP8_BLOCK_SIZE = 32
FP4_BLOCK_SIZE = 32
FP4_AMAX_FLOOR = 6 * 2.0**-126
def ceil_pow2(x: torch.Tensor) -> torch.Tensor:
"""2 ** ceil(log2(x)) for positive fp32 x, computed on the IEEE bits so the
result is exact at powers of two."""
bits = x.contiguous().view(torch.int32)
exponent = ((bits >> 23) & 0xFF) - 127
has_mantissa = (bits & 0x7FFFFF) != 0
exponent = exponent + has_mantissa.to(torch.int32)
return ((exponent + 127) << 23).view(torch.float32)
def block_scale(x: torch.Tensor, block_size: int, fmax: float, amax_floor: float):
"""Per-block ue8m0 scale, as fp32 powers of two, shape [..., N // block_size]."""
amax = x.float().unflatten(-1, (-1, block_size)).abs().amax(dim=-1)
amax = amax.clamp_min(amax_floor)
# The kernel multiplies by the fp32 reciprocal rather than dividing. A Python
# scalar keeps this free of host tensors, so it can run under CUDA graph capture.
return ceil_pow2(amax * (1.0 / fmax))
def round_fp4(x: torch.Tensor) -> torch.Tensor:
"""Round fp32 values in [-6, 6] onto the e2m1 grid with round-to-nearest-even."""
magnitude = x.abs()
step = torch.where(magnitude < 2.0, 0.5, torch.where(magnitude < 4.0, 1.0, 2.0))
return torch.round(magnitude / step) * step * torch.sign(x)
def fake_quant_fp4(x: torch.Tensor, block_size: int = FP4_BLOCK_SIZE) -> torch.Tensor:
"""Quantize to fp4 (per-block ue8m0 scale) and back, in x's dtype."""
scale = block_scale(x, block_size, FP4_MAX, FP4_AMAX_FLOOR)
scaled = x.float().unflatten(-1, (-1, block_size)) / scale.unsqueeze(-1)
deq = round_fp4(scaled.clamp(-FP4_MAX, FP4_MAX)) * scale.unsqueeze(-1)
return deq.flatten(-2).to(x.dtype)
def fake_quant_compressed_kv(x: torch.Tensor) -> torch.Tensor:
"""FP4 round-trip with one E4M3FN scale per 16 compressed-KV elements.
Round amax / 6 to E4M3 with ties to even, clamping the scale to its
positive finite range [2**-9, 448]. Zero blocks remain zero.
"""
blocks = x.float().unflatten(-1, (-1, 16))
amax = blocks.abs().amax(dim=-1, keepdim=True)
scale = (amax * (1.0 / FP4_MAX)).clamp(min=2**-9, max=FP8_MAX)
scale = scale.to(torch.float8_e4m3fn).float()
scaled = (blocks / scale).clamp(-FP4_MAX, FP4_MAX)
deq = round_fp4(scaled) * scale
return deq.flatten(-2).to(x.dtype)
# ---------------------------------------------------------------------------
# Pure-torch reference of the paged V4.1 fp8 KV cache format read by the sparse
# decode kernel (528 B/token, "V41").
# ---------------------------------------------------------------------------
def ceil_pow2_scale(x: torch.Tensor) -> torch.Tensor:
"""``2 ** ceil(log2(max(x, 1e-4)))`` as fp32, computed on the IEEE bits so
that it is exact at (and just above) powers of two."""
x = x.float()
scale = ceil_pow2(torch.clamp_min(x, 1e-4))
# ceil_pow2 works on the bits of a finite value; a NaN or inf amax passes
# through (both become the ue8m0 NaN byte, but only the NaN one turns the
# whole tile's payload into NaN).
return torch.where(torch.isfinite(x), scale, x)
def quantize_k_cache_v41(
k: torch.Tensor, page_bytes: Optional[int] = None
) -> torch.Tensor:
"""``k`` ``[num_pages, page_size, 512]`` -> uint8 ``[num_pages, page_bytes]``
pages of the V41 layout: 512 e4m3 per token, then 16 ue8m0 scales per token
(one per 32 values), ``scale = 2 ** ceil(log2(max(amax / 448, 1e-4)))``."""
num_pages, page_size, d = k.shape
assert d == 512
x = k.float().view(num_pages, page_size, 16, 32)
scale = ceil_pow2_scale(x.abs().amax(dim=-1) / 448.0)
data = (x / scale.unsqueeze(-1)).to(torch.float8_e4m3fn).view(torch.uint8)
scale_u8 = scale.to(torch.float8_e8m0fnu).view(torch.uint8)
raw = page_size * 528
if page_bytes is None:
page_bytes = -(-raw // 512) * 512
assert page_bytes >= raw
out = torch.zeros((num_pages, page_bytes), dtype=torch.uint8, device=k.device)
out[:, : page_size * 512] = data.reshape(num_pages, page_size * 512)
out[:, page_size * 512 : raw] = scale_u8.reshape(num_pages, page_size * 16)
return out
def dequantize_k_cache_v41(pages: torch.Tensor, page_size: int) -> torch.Tensor:
"""Inverse of :func:`quantize_k_cache_v41`: ``[num_pages, page_size, 512]`` bf16."""
num_pages = pages.shape[0]
pages = pages.view(torch.uint8)
data = pages[:, : page_size * 512].reshape(num_pages, page_size, 512)
scale = pages[:, page_size * 512 : page_size * 528].reshape(
num_pages, page_size, 16
)
values = data.view(torch.float8_e4m3fn).to(torch.bfloat16)
scale_bf16 = scale.view(torch.float8_e8m0fnu).to(torch.bfloat16)
return (values.view(num_pages, page_size, 16, 32) * scale_bf16.unsqueeze(-1)).view(
num_pages, page_size, 512
)
python/sglang/kernels/ops/attention/dsv4_attn_metadata_kernels.py
+68
View File
@@ -525,3 +525,71 @@ def build_causal_swa_page_indices_triton(
BLOCK_K=BLOCK_K,
)
return out
@triton.jit
def _small_page_table(
REQ_TO_TOKEN,
REQS,
LENS,
OUT_LENS,
POS,
PAGES,
SWA,
STRIDE: tl.constexpr,
NUM_PAGES: tl.constexpr,
PAGE_SIZE: tl.constexpr,
WINDOW: tl.constexpr,
BLOCK: tl.constexpr,
):
row, tile = tl.program_id(0), tl.program_id(1)
if tile == 0:
length = tl.load(LENS + row).to(tl.int32)
tl.store(OUT_LENS + row, length)
tl.store(POS + row, length - 1)
tl.store(SWA + row, tl.minimum(length, WINDOW))
req = tl.load(REQS + row).to(tl.int64)
p = tile * BLOCK + tl.arange(0, BLOCK)
slot = tl.load(
REQ_TO_TOKEN + req * STRIDE + p.to(tl.int64) * PAGE_SIZE,
mask=p < NUM_PAGES,
other=0,
).to(tl.int32)
tl.store(PAGES + row * NUM_PAGES + p, slot // PAGE_SIZE, mask=p < NUM_PAGES)
def build_page_table_positions_small(
*,
req_to_token,
req_pool_indices_repeated,
seq_lens_casual,
max_seq_len,
page_size,
swa_window,
):
assert page_size > 0 and page_size & (page_size - 1) == 0
rows = seq_lens_casual.numel()
pages = triton.cdiv(max_seq_len, page_size)
kw = dict(device=seq_lens_casual.device, dtype=torch.int32)
lengths, positions, swa = [torch.empty(rows, **kw) for _ in range(3)]
table = torch.empty((rows, pages), **kw)
_small_page_table[(rows, triton.cdiv(pages, 256))](
req_to_token,
req_pool_indices_repeated,
seq_lens_casual,
lengths,
positions,
table,
swa,
req_to_token.stride(0),
pages,
page_size,
swa_window,
256,
)
return PageTablePositionsResult(
seq_lens_casual=lengths,
positions_casual=positions,
page_table=table,
swa_topk_lengths=swa,
)