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

Deepseek v4: support mixed dtype compression states (#27277)

Browse Source 
Co-authored-by: zhujunyu <zhujunyu.666@bytedance.com>
This commit is contained in:
Ryan Zzz
2026-06-17 00:56:41 -07:00
committed by GitHub
co-authored by zhujunyu
parent 7256ee9871
commit 8fd1694dd2
9 changed files with 1414 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh
+114 -50
View File
@@ -27,7 +27,9 @@
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/object.h>
#include <cfloat>
#include <cstdint>
#include <type_traits>
namespace {
@@ -89,12 +91,12 @@ struct C128Trait {
static_assert(kHeadDim % kTileDim == 0);
};
template <typename Trait, bool kUsePDL, typename InFloat, typename OutFloat>
template <typename Trait, bool kUsePDL, typename BufferFloat, typename InputFloat, typename OutFloat>
SGL_DEVICE void c128_forward(
const InFloat* kv_buf, // [128n, 128n + 127]
const InFloat* kv_src, // ragged pointer at position = 128n + 127
const BufferFloat* kv_buf, // [128n, 128n + 127]
const InputFloat* kv_src, // ragged pointer at position = 128n + 127
OutFloat* kv_out,
const InFloat* score_bias,
const InputFloat* score_bias,
const int32_t buffer_len) {
using namespace device;
@@ -102,7 +104,7 @@ SGL_DEVICE void c128_forward(
const auto lane_id = threadIdx.x % kWarpThreads;
/// NOTE: part 1: load kv + score
using StorageIn = AlignedVector<InFloat, kTileElements>;
using StorageIn = AlignedVector<InputFloat, kTileElements>;
const auto gmem_in = tile::Memory<StorageIn>{lane_id, kWarpThreads};
StorageIn kv[kElementsPerWarp];
StorageIn score[kElementsPerWarp];
@@ -117,13 +119,38 @@ SGL_DEVICE void c128_forward(
const auto kv_start = kv_src - 127 * Trait::kElementSize; // point to start
if constexpr (std::is_same_v<BufferFloat, InputFloat>) {
#pragma unroll
for (int32_t i = 0; i < kElementsPerWarp; ++i) {
const int32_t j = i + warp_offset;
__builtin_assume(j < 128);
const auto src = j < buffer_len ? kv_buf : kv_start;
kv[i] = gmem_in.load(src + j * Trait::kElementSize);
score[i] = gmem_in.load(src + j * Trait::kElementSize + Trait::kScoreOffset);
for (int32_t i = 0; i < kElementsPerWarp; ++i) {
const int32_t j = i + warp_offset;
__builtin_assume(j < 128);
const auto src = j < buffer_len ? kv_buf : kv_start;
kv[i] = gmem_in.load(src + j * Trait::kElementSize);
score[i] = gmem_in.load(src + j * Trait::kElementSize + Trait::kScoreOffset);
}
} else { // mixed dtype
using StorageBuffer = AlignedVector<BufferFloat, kTileElements>;
const auto gmem_buffer = tile::Memory<StorageBuffer>{lane_id, kWarpThreads};
#pragma unroll
for (int32_t i = 0; i < kElementsPerWarp; ++i) {
const int32_t j = i + warp_offset;
__builtin_assume(j < 128);
if (j < buffer_len) {
const auto src = kv_buf + j * Trait::kElementSize;
const auto kv_tmp = gmem_buffer.load(src);
const auto score_tmp = gmem_buffer.load(src + Trait::kScoreOffset);
#pragma unroll
for (int32_t k = 0; k < kTileElements; ++k) {
kv[i][k] = cast<InputFloat>(kv_tmp[k]);
score[i][k] = cast<InputFloat>(score_tmp[k]);
}
} else {
const auto src = kv_start + j * Trait::kElementSize;
kv[i] = gmem_in.load(src);
score[i] = gmem_in.load(src + Trait::kScoreOffset);
}
}
}
/// NOTE: part 2: safe online softmax + weighted sum
@@ -141,6 +168,7 @@ SGL_DEVICE void c128_forward(
// convert to fp32 and apply bias first
#pragma unroll
for (int32_t i = 0; i < kTileElements; ++i) {
#pragma unroll
for (int32_t j = 0; j < kElementsPerWarp; ++j) {
score_fp32[i][j] = cast<float>(score[j][i]) + cast<float>(bias[j][i]);
}
@@ -215,25 +243,41 @@ SGL_DEVICE void c128_forward(
}
}
template <typename Trait, typename InFloat>
SGL_DEVICE void c128_write_decode(InFloat* kv_buf, const InFloat* kv_src) {
template <typename Trait, typename BufferFloat, typename InputFloat>
SGL_DEVICE void c128_write_decode(BufferFloat* kv_buf, const InputFloat* kv_src) {
using namespace device;
using Storage = AlignedVector<InFloat, kTileElements>;
const auto gmem = tile::Memory<Storage>::warp();
using StorageInput = AlignedVector<InputFloat, kTileElements>;
const auto gmem_input = tile::Memory<StorageInput>::warp();
Storage data[2];
StorageInput data[2];
#pragma unroll
for (int32_t i = 0; i < 2; ++i) {
data[i] = gmem.load(kv_src + Trait::kHeadDim * i);
data[i] = gmem_input.load(kv_src + Trait::kHeadDim * i);
}
if constexpr (std::is_same_v<BufferFloat, InputFloat>) {
#pragma unroll
for (int32_t i = 0; i < 2; ++i) {
gmem.store(kv_buf + Trait::kHeadDim * i, data[i]);
for (int32_t i = 0; i < 2; ++i) {
gmem_input.store(kv_buf + Trait::kHeadDim * i, data[i]);
}
} else {
using StorageBuffer = AlignedVector<BufferFloat, kTileElements>;
const auto gmem_buffer = tile::Memory<StorageBuffer>::warp();
StorageBuffer data_cast[2];
#pragma unroll
for (int32_t i = 0; i < 2; ++i) {
#pragma unroll
for (int32_t j = 0; j < kTileElements; ++j) {
data_cast[i][j] = cast<BufferFloat>(data[i][j]);
}
gmem_buffer.store(kv_buf + Trait::kHeadDim * i, data_cast[i]);
}
}
}
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
C128_KERNEL void flash_c128_decode(const __grid_constant__ Compress128DecodeParams params) {
using namespace device;
using Trait = C128Trait<kHeadDim>;
@@ -245,10 +289,10 @@ C128_KERNEL void flash_c128_decode(const __grid_constant__ Compress128DecodePara
if (global_bid >= params.batch_size) return;
const auto plan = params.plan_d[global_bid];
const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
const auto kv_input = static_cast<const InputFloat*>(params.kv_input) + split_offset;
const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
const auto score_bias = static_cast<const InFloat*>(params.score_bias) + split_offset;
const auto kv_buffer = static_cast<BufferFloat*>(params.kv_buffer) + split_offset;
const auto score_bias = static_cast<const InputFloat*>(params.score_bias) + split_offset;
const auto kv_src = kv_input + global_bid * Trait::kElementSize;
const auto kv_out = kv_output + global_bid * Trait::kHeadDim;
@@ -258,15 +302,15 @@ C128_KERNEL void flash_c128_decode(const __grid_constant__ Compress128DecodePara
PDLWaitPrimary<kUsePDL>();
// the write warp must match the load warp in the following `c128_forward`
if (warp_id == kNumWarps - 1) {
c128_write_decode<Trait>(kv_dst, kv_src);
c128_write_decode<Trait, BufferFloat, InputFloat>(kv_dst, kv_src);
}
if (plan.write_loc % 128 == 127) {
c128_forward<Trait, kUsePDL>(kv_buf, kv_src, kv_out, score_bias, 128);
c128_forward<Trait, kUsePDL, BufferFloat, InputFloat, OutFloat>(kv_buf, kv_src, kv_out, score_bias, 128);
}
}
// compress kernel
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
C128_KERNEL void flash_c128_prefill(const __grid_constant__ Compress128PrefillParams params) {
using namespace device;
using Trait = C128Trait<kHeadDim>;
@@ -277,10 +321,10 @@ C128_KERNEL void flash_c128_prefill(const __grid_constant__ Compress128PrefillPa
if (global_pid >= params.num_compress) return;
const auto plan = params.plan_c[global_pid];
const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
const auto kv_input = static_cast<const InputFloat*>(params.kv_input) + split_offset;
const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
const auto score_bias = static_cast<const InFloat*>(params.score_bias) + split_offset;
const auto kv_buffer = static_cast<BufferFloat*>(params.kv_buffer) + split_offset;
const auto score_bias = static_cast<const InputFloat*>(params.score_bias) + split_offset;
if (plan.is_invalid()) return;
const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize;
@@ -288,14 +332,14 @@ C128_KERNEL void flash_c128_prefill(const __grid_constant__ Compress128PrefillPa
const auto kv_out = kv_output + global_pid * Trait::kHeadDim;
const auto kv_buf = kv_buffer + plan.read_page_1 * Trait::kPageElementSize;
PDLWaitPrimary<kUsePDL>();
c128_forward<Trait, kUsePDL>(kv_buf, kv_src, kv_out, score_bias, plan.buffer_len);
c128_forward<Trait, kUsePDL, BufferFloat, InputFloat, OutFloat>(kv_buf, kv_src, kv_out, score_bias, plan.buffer_len);
}
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
WRITE_KERNEL void write_c128_prefill(const __grid_constant__ Compress128PrefillParams params) {
using namespace device;
using Trait = C128Trait<kHeadDim>;
using StorageIn = AlignedVector<InFloat, kTileElements>;
using StorageInput = AlignedVector<InputFloat, kTileElements>;
const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
const uint32_t global_wid = global_tid / kWarpThreads; // warp id
@@ -307,33 +351,53 @@ WRITE_KERNEL void write_c128_prefill(const __grid_constant__ Compress128PrefillP
if (global_pid >= params.num_write) return;
const auto plan = params.plan_w[global_pid];
const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
const auto kv_input = static_cast<const InputFloat*>(params.kv_input) + split_offset;
const auto kv_buffer = static_cast<BufferFloat*>(params.kv_buffer) + split_offset;
if (plan.is_invalid()) return;
// each warp will handle a contiguous region
const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize;
const auto kv_buf = kv_buffer + plan.write_loc * Trait::kElementSize;
const auto gmem = tile::Memory<StorageIn>::warp();
const auto gmem_input = tile::Memory<StorageInput>::warp();
PDLWaitPrimary<kUsePDL>();
StorageIn data[2];
StorageInput data[2];
#pragma unroll
for (int32_t i = 0; i < 2; ++i) {
data[i] = gmem.load(kv_src, i);
data[i] = gmem_input.load(kv_src, i);
}
PDLTriggerSecondary<kUsePDL>();
if constexpr (std::is_same_v<BufferFloat, InputFloat>) {
PDLTriggerSecondary<kUsePDL>();
#pragma unroll
for (int32_t i = 0; i < 2; ++i) {
gmem.store(kv_buf, data[i], i);
for (int32_t i = 0; i < 2; ++i) {
gmem_input.store(kv_buf, data[i], i);
}
} else {
using StorageBuffer = AlignedVector<BufferFloat, kTileElements>;
const auto gmem_buffer = tile::Memory<StorageBuffer>::warp();
StorageBuffer data_cast[2];
#pragma unroll
for (int32_t i = 0; i < 2; ++i) {
#pragma unroll
for (int32_t j = 0; j < kTileElements; ++j) {
data_cast[i][j] = cast<BufferFloat>(data[i][j]);
}
}
PDLTriggerSecondary<kUsePDL>();
#pragma unroll
for (int32_t i = 0; i < 2; ++i) {
gmem_buffer.store(kv_buf, data_cast[i], i);
}
}
}
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
struct FlashCompress128Kernel {
static constexpr auto decode_kernel = flash_c128_decode<kHeadDim, InFloat, OutFloat, kUsePDL>;
static constexpr auto prefill_c_kernel = flash_c128_prefill<kHeadDim, InFloat, OutFloat, kUsePDL>;
static constexpr auto prefill_w_kernel = write_c128_prefill<kHeadDim, InFloat, OutFloat, kUsePDL>;
static constexpr auto decode_kernel = flash_c128_decode<kHeadDim, BufferFloat, InputFloat, OutFloat, kUsePDL>;
static constexpr auto prefill_c_kernel = flash_c128_prefill<kHeadDim, BufferFloat, InputFloat, OutFloat, kUsePDL>;
static constexpr auto prefill_w_kernel = write_c128_prefill<kHeadDim, BufferFloat, InputFloat, OutFloat, kUsePDL>;
static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads; // 64
static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
using Trait = C128Trait<kHeadDim>;
@@ -351,11 +415,11 @@ struct FlashCompress128Kernel {
device_.set_options<kDLGPU>();
TensorMatcher({-1, 128, Trait::kElementSize}) // kv score
.with_dtype<InFloat>()
.with_dtype<BufferFloat>()
.with_device(device_)
.verify(kv_buffer);
TensorMatcher({N, Trait::kElementSize}) // kv score input
.with_dtype<InFloat>()
.with_dtype<InputFloat>()
.with_device(device_)
.verify(kv_input);
TensorMatcher({N, kHeadDim}) // kv compressed output
@@ -363,7 +427,7 @@ struct FlashCompress128Kernel {
.with_device(device_)
.verify(kv_output);
TensorMatcher({128, kHeadDim}) // ape
.with_dtype<InFloat>()
.with_dtype<InputFloat>()
.with_device(device_)
.verify(ape);
@@ -398,11 +462,11 @@ struct FlashCompress128Kernel {
device_.set_options<kDLGPU>();
TensorMatcher({-1, 128, Trait::kElementSize}) // kv score
.with_dtype<InFloat>()
.with_dtype<BufferFloat>()
.with_device(device_)
.verify(kv_buffer);
TensorMatcher({N, Trait::kElementSize}) // kv score input (ragged)
.with_dtype<InFloat>()
.with_dtype<InputFloat>()
.with_device(device_)
.verify(kv_input);
TensorMatcher({C, kHeadDim}) // kv compressed output (compact)
@@ -410,7 +474,7 @@ struct FlashCompress128Kernel {
.with_device(device_)
.verify(kv_output);
TensorMatcher({128, kHeadDim}) // ape
.with_dtype<InFloat>()
.with_dtype<InputFloat>()
.with_device(device_)
.verify(ape);
python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh
+154 -68
View File
@@ -28,6 +28,7 @@
#include <cfloat>
#include <cstdint>
#include <type_traits>
namespace {
@@ -74,20 +75,18 @@ struct C4Trait {
static_assert(kHeadDim % kTileDim == 0);
};
template <typename Trait, bool kUsePDL, typename InFloat, typename OutFloat>
template <typename Trait, bool kUsePDL, typename BufferFloat, typename InputFloat, typename OutFloat>
SGL_DEVICE void c4_forward(
const InFloat* kv_buf_0, // overlap [4n - 4, 4n - 1]
const InFloat* kv_buf_1, // normal [4n + 0, 4n + 3]
const InFloat* kv_src, // ragged pointer at position = 4n + 3
const BufferFloat* kv_buf_0, // overlap [4n - 4, 4n - 1]
const BufferFloat* kv_buf_1, // normal [4n + 0, 4n + 3]
const InputFloat* kv_src, // ragged pointer at position = 4n + 3
OutFloat* kv_out,
const InFloat* score_bias,
const InputFloat* score_bias,
const bool should_overlap,
const int32_t buffer_len) {
using namespace device;
/// NOTE: part 1: load kv + score
using StorageIn = AlignedVector<InFloat, kTileElements>;
/// NOTE: load one tile_dim (< head_dim) at at time
using StorageIn = AlignedVector<InputFloat, kTileElements>;
const auto gmem_in = tile::Memory<StorageIn>::warp();
StorageIn kv[8];
StorageIn score[8];
@@ -98,32 +97,80 @@ SGL_DEVICE void c4_forward(
bias[i] = gmem_in.load(score_bias + i * Trait::kHeadDim);
}
if (should_overlap) {
const auto kv_start = kv_src - 7 * Trait::kElementSize; // point to start
if constexpr (std::is_same_v<BufferFloat, InputFloat>) {
if (should_overlap) {
const auto kv_start = kv_src - 7 * Trait::kElementSize; // point to start
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
const auto src = i < buffer_len ? kv_buf_0 : kv_start;
const auto base = src + i * Trait::kElementSize;
kv[i] = gmem_in.load(base);
score[i] = gmem_in.load(base + Trait::kScoreOffset);
}
} else {
[[unlikely]];
constexpr float kFloatNegInf = -FLT_MAX;
for (int32_t i = 0; i < 4; ++i) {
const auto src = i < buffer_len ? kv_buf_0 : kv_start;
const auto base = src + i * Trait::kElementSize;
kv[i] = gmem_in.load(base);
score[i] = gmem_in.load(base + Trait::kScoreOffset);
}
} else {
[[unlikely]];
constexpr float kFloatNegInf = -FLT_MAX;
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
kv[i].fill(cast<InFloat>(0.0f));
score[i].fill(cast<InFloat>(kFloatNegInf));
for (int32_t i = 0; i < 4; ++i) {
kv[i].fill(cast<InputFloat>(0.0f));
score[i].fill(cast<InputFloat>(kFloatNegInf));
}
}
}
const auto kv_start = kv_src - 3 * Trait::kElementSize; // point to start
const auto kv_start = kv_src - 3 * Trait::kElementSize; // point to start
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
const auto src = i + 4 < buffer_len ? kv_buf_1 : kv_start;
const auto base = src + i * Trait::kElementSize + Trait::kOverlapOffset;
kv[i + 4] = gmem_in.load(base);
score[i + 4] = gmem_in.load(base + Trait::kScoreOffset);
for (int32_t i = 0; i < 4; ++i) {
const auto src = i + 4 < buffer_len ? kv_buf_1 : kv_start;
const auto base = src + i * Trait::kElementSize + Trait::kOverlapOffset;
kv[i + 4] = gmem_in.load(base);
score[i + 4] = gmem_in.load(base + Trait::kScoreOffset);
}
} else { // mixed dtype
using StorageBuffer = AlignedVector<BufferFloat, kTileElements>;
const auto gmem_buffer = tile::Memory<StorageBuffer>::warp();
const auto kv_start_0 = kv_src - 7 * Trait::kElementSize; // point to start
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
if (should_overlap && i < buffer_len) {
const auto base = kv_buf_0 + i * Trait::kElementSize;
const auto kv_tmp = gmem_buffer.load(base);
const auto score_tmp = gmem_buffer.load(base + Trait::kScoreOffset);
#pragma unroll
for (int32_t j = 0; j < kTileElements; ++j) {
kv[i][j] = cast<InputFloat>(kv_tmp[j]);
score[i][j] = cast<InputFloat>(score_tmp[j]);
}
} else if (should_overlap) {
const auto base = kv_start_0 + i * Trait::kElementSize;
kv[i] = gmem_in.load(base);
score[i] = gmem_in.load(base + Trait::kScoreOffset);
} else {
[[unlikely]];
constexpr float kFloatNegInf = -FLT_MAX;
kv[i].fill(cast<InputFloat>(0.0f));
score[i].fill(cast<InputFloat>(kFloatNegInf));
}
}
const auto kv_start = kv_src - 3 * Trait::kElementSize; // point to start
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
if (i + 4 < buffer_len) {
const auto base = kv_buf_1 + i * Trait::kElementSize + Trait::kOverlapOffset;
const auto kv_tmp = gmem_buffer.load(base);
const auto score_tmp = gmem_buffer.load(base + Trait::kScoreOffset);
#pragma unroll
for (int32_t j = 0; j < kTileElements; ++j) {
kv[i + 4][j] = cast<InputFloat>(kv_tmp[j]);
score[i + 4][j] = cast<InputFloat>(score_tmp[j]);
}
} else {
const auto base = kv_start + i * Trait::kElementSize + Trait::kOverlapOffset;
kv[i + 4] = gmem_in.load(base);
score[i + 4] = gmem_in.load(base + Trait::kScoreOffset);
}
}
}
/// NOTE: part 2: safe online softmax + weighted sum
@@ -137,6 +184,7 @@ SGL_DEVICE void c4_forward(
// convert to fp32 and apply bias first
#pragma unroll
for (int32_t i = 0; i < kTileElements; ++i) {
#pragma unroll
for (int32_t j = 0; j < 8; ++j) {
score_fp32[i][j] = cast<float>(score[j][i]) + cast<float>(bias[j][i]);
}
@@ -171,25 +219,41 @@ SGL_DEVICE void c4_forward(
gmem_out.store(kv_out, result);
}
template <typename Trait, typename InFloat>
SGL_DEVICE void c4_write_decode(InFloat* kv_buf, const InFloat* kv_src) {
template <typename Trait, typename BufferFloat, typename InputFloat>
SGL_DEVICE void c4_write_decode(BufferFloat* kv_buf, const InputFloat* kv_src) {
using namespace device;
using StorageIn = AlignedVector<InFloat, kTileElements>;
const auto gmem = tile::Memory<StorageIn>::warp();
using StorageInput = AlignedVector<InputFloat, kTileElements>;
const auto gmem_input = tile::Memory<StorageInput>::warp();
StorageIn data[4];
StorageInput data[4];
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
data[i] = gmem.load(kv_src + Trait::kHeadDim * i);
data[i] = gmem_input.load(kv_src + Trait::kHeadDim * i);
}
if constexpr (std::is_same_v<BufferFloat, InputFloat>) {
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
gmem.store(kv_buf + Trait::kHeadDim * i, data[i]);
for (int32_t i = 0; i < 4; ++i) {
gmem_input.store(kv_buf + Trait::kHeadDim * i, data[i]);
}
} else {
using StorageBuffer = AlignedVector<BufferFloat, kTileElements>;
const auto gmem_buffer = tile::Memory<StorageBuffer>::warp();
StorageBuffer data_cast[4];
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
#pragma unroll
for (int32_t j = 0; j < kTileElements; ++j) {
data_cast[i][j] = cast<BufferFloat>(data[i][j]);
}
gmem_buffer.store(kv_buf + Trait::kHeadDim * i, data_cast[i]);
}
}
}
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
C4_KERNEL void flash_c4_decode(const __grid_constant__ Compress4DecodeParams params) {
using namespace device;
using Trait = C4Trait<kHeadDim>;
@@ -202,10 +266,10 @@ C4_KERNEL void flash_c4_decode(const __grid_constant__ Compress4DecodeParams par
if (global_bid >= params.batch_size) return;
const auto plan = params.plan_d[global_bid];
const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
const auto kv_input = static_cast<const InputFloat*>(params.kv_input) + split_offset;
const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
const auto score_bias = static_cast<const InFloat*>(params.score_bias) + split_offset;
const auto kv_buffer = static_cast<BufferFloat*>(params.kv_buffer) + split_offset;
const auto score_bias = static_cast<const InputFloat*>(params.score_bias) + split_offset;
const auto kv_src = kv_input + global_bid * Trait::kElementSize;
const auto kv_out = kv_output + global_bid * Trait::kHeadDim;
@@ -214,14 +278,15 @@ C4_KERNEL void flash_c4_decode(const __grid_constant__ Compress4DecodeParams par
const auto kv_dst = kv_buffer + plan.write_loc * Trait::kElementSize;
PDLWaitPrimary<kUsePDL>();
c4_write_decode<Trait>(kv_dst, kv_src);
c4_write_decode<Trait, BufferFloat, InputFloat>(kv_dst, kv_src);
if (plan.seq_len % 4 == 0) {
const auto need_overlap = plan.seq_len > 4;
c4_forward<Trait, kUsePDL>(kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, 8);
c4_forward<Trait, kUsePDL, BufferFloat, InputFloat, OutFloat>(
kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, 8);
}
}
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
C4_KERNEL void flash_c4_prefill(const __grid_constant__ Compress4PrefillParams params) {
using namespace device;
using Trait = C4Trait<kHeadDim>;
@@ -234,10 +299,10 @@ C4_KERNEL void flash_c4_prefill(const __grid_constant__ Compress4PrefillParams p
if (global_pid >= params.num_compress) return;
const auto plan = params.plan_c[global_pid];
const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
const auto kv_input = static_cast<const InputFloat*>(params.kv_input) + split_offset;
const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
const auto score_bias = static_cast<const InFloat*>(params.score_bias) + split_offset;
const auto kv_buffer = static_cast<BufferFloat*>(params.kv_buffer) + split_offset;
const auto score_bias = static_cast<const InputFloat*>(params.score_bias) + split_offset;
if (plan.is_invalid()) return;
const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize;
@@ -247,14 +312,15 @@ C4_KERNEL void flash_c4_prefill(const __grid_constant__ Compress4PrefillParams p
const auto kv_buf_1 = kv_buffer + plan.read_page_1 * Trait::kPageElementSize;
const bool need_overlap = plan.seq_len > 4;
PDLWaitPrimary<kUsePDL>();
c4_forward<Trait, kUsePDL>(kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, plan.buffer_len);
c4_forward<Trait, kUsePDL, BufferFloat, InputFloat, OutFloat>(
kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, plan.buffer_len);
}
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
WRITE_KERNEL void write_c4_prefill(const __grid_constant__ Compress4PrefillParams params) {
using namespace device;
using Trait = C4Trait<kHeadDim>;
using StorageIn = AlignedVector<InFloat, kTileElements>;
using StorageInput = AlignedVector<InputFloat, kTileElements>;
const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
const uint32_t global_wid = global_tid / kWarpThreads; // warp id
@@ -266,33 +332,53 @@ WRITE_KERNEL void write_c4_prefill(const __grid_constant__ Compress4PrefillParam
if (global_pid >= params.num_write) return;
const auto plan = params.plan_w[global_pid];
const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
const auto kv_input = static_cast<const InputFloat*>(params.kv_input) + split_offset;
const auto kv_buffer = static_cast<BufferFloat*>(params.kv_buffer) + split_offset;
if (plan.is_invalid()) return;
// each warp will handle a contiguous region
const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize;
const auto kv_buf = kv_buffer + plan.write_loc * Trait::kElementSize;
const auto gmem = tile::Memory<StorageIn>::warp();
const auto gmem_input = tile::Memory<StorageInput>::warp();
PDLWaitPrimary<kUsePDL>();
StorageIn data[4];
StorageInput data[4];
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
data[i] = gmem.load(kv_src, i);
data[i] = gmem_input.load(kv_src, i);
}
PDLTriggerSecondary<kUsePDL>();
if constexpr (std::is_same_v<BufferFloat, InputFloat>) {
PDLTriggerSecondary<kUsePDL>();
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
gmem.store(kv_buf, data[i], i);
for (int32_t i = 0; i < 4; ++i) {
gmem_input.store(kv_buf, data[i], i);
}
} else {
using StorageBuffer = AlignedVector<BufferFloat, kTileElements>;
const auto gmem_buffer = tile::Memory<StorageBuffer>::warp();
StorageBuffer data_cast[4];
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
#pragma unroll
for (int32_t j = 0; j < kTileElements; ++j) {
data_cast[i][j] = cast<BufferFloat>(data[i][j]);
}
}
PDLTriggerSecondary<kUsePDL>();
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
gmem_buffer.store(kv_buf, data_cast[i], i);
}
}
}
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
struct FlashCompress4Kernel {
static constexpr auto decode_kernel = flash_c4_decode<kHeadDim, InFloat, OutFloat, kUsePDL>;
static constexpr auto prefill_c_kernel = flash_c4_prefill<kHeadDim, InFloat, OutFloat, kUsePDL>;
static constexpr auto prefill_w_kernel = write_c4_prefill<kHeadDim, InFloat, OutFloat, kUsePDL>;
static constexpr auto decode_kernel = flash_c4_decode<kHeadDim, BufferFloat, InputFloat, OutFloat, kUsePDL>;
static constexpr auto prefill_c_kernel = flash_c4_prefill<kHeadDim, BufferFloat, InputFloat, OutFloat, kUsePDL>;
static constexpr auto prefill_w_kernel = write_c4_prefill<kHeadDim, BufferFloat, InputFloat, OutFloat, kUsePDL>;
static constexpr uint32_t kBlockSize = 128;
static constexpr uint32_t kTileDim = kTileElements * device::kWarpThreads;
static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
@@ -312,11 +398,11 @@ struct FlashCompress4Kernel {
device_.set_options<kDLGPU>();
TensorMatcher({-1, 4, Trait::kElementSize}) // kv score
.with_dtype<InFloat>()
.with_dtype<BufferFloat>()
.with_device(device_)
.verify(kv_buffer);
TensorMatcher({N, Trait::kElementSize}) // kv score input
.with_dtype<InFloat>()
.with_dtype<InputFloat>()
.with_device(device_)
.verify(kv_input);
TensorMatcher({N, kHeadDim}) // kv compressed output
@@ -324,7 +410,7 @@ struct FlashCompress4Kernel {
.with_device(device_)
.verify(kv_output);
TensorMatcher({8, kHeadDim}) // ape
.with_dtype<InFloat>()
.with_dtype<InputFloat>()
.with_device(device_)
.verify(ape);
@@ -359,11 +445,11 @@ struct FlashCompress4Kernel {
device_.set_options<kDLGPU>();
TensorMatcher({-1, 4, Trait::kElementSize}) // kv score
.with_dtype<InFloat>()
.with_dtype<BufferFloat>()
.with_device(device_)
.verify(kv_buffer);
TensorMatcher({N, Trait::kElementSize}) // kv score input (ragged)
.with_dtype<InFloat>()
.with_dtype<InputFloat>()
.with_device(device_)
.verify(kv_input);
TensorMatcher({C, kHeadDim}) // kv compressed output (compact)
@@ -371,7 +457,7 @@ struct FlashCompress4Kernel {
.with_device(device_)
.verify(kv_output);
TensorMatcher({8, kHeadDim}) // ape
.with_dtype<InFloat>()
.with_dtype<InputFloat>()
.with_device(device_)
.verify(ape);
const auto plan_c = compress::verify_plan_c(plan_c_, C, device_);
python/sglang/jit_kernel/dsv4/compress.py
+7 -2
View File
@@ -44,11 +44,14 @@ def _jit_compress_norm_rope_module(
@cache_once
def _jit_compress_module(
head_dim: int,
dtype_buffer: torch.dtype,
dtype_in: torch.dtype,
dtype_out: torch.dtype,
ratio: Literal[4, 128],
) -> Module:
args = make_cpp_args(head_dim, dtype_in, dtype_out, is_arch_support_pdl())
args = make_cpp_args(
head_dim, dtype_buffer, dtype_in, dtype_out, is_arch_support_pdl()
)
kernel_class = f"FlashCompress{ratio}Kernel<{args}>"
return load_jit(
make_name(f"compress_{ratio}_v2"),
@@ -336,7 +339,9 @@ def compress_forward(
module = _jit_compress_128_online_module(512)
else:
dtype_in, dtype_out = kv_score_input.dtype, out.dtype
module = _jit_compress_module(head_dim, dtype_in, dtype_out, compress_ratio)
module = _jit_compress_module(
head_dim, kv_score_buffer.dtype, dtype_in, dtype_out, compress_ratio
)
fn = module.decode if plan.is_decode else module.prefill
fn(kv_score_buffer, kv_score_input, out, ape, *plan[1:3])
return out
python/sglang/srt/environ.py
+1
View File
@@ -788,6 +788,7 @@ class Envs:
SGLANG_OPT_USE_JIT_INDEXER_METADATA = EnvBool(True)
SGLANG_OPT_USE_ONLINE_COMPRESS = EnvBool(False)
SGLANG_EXPERIMENTAL_ONLINE_C128_MTP = EnvBool(False)
SGLANG_DSV4_COMPRESS_STATE_DTYPE = EnvStr("float32")
SGLANG_OPT_USE_COMPRESSOR_V2 = EnvBool(True)
SGLANG_FP8_PAGED_MQA_LOGITS_TORCH = EnvBool(False)
SGLANG_TOPK_TRANSFORM_512_TORCH = EnvBool(False)
python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
+6 -4
View File
@@ -448,7 +448,8 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
page_size: int,
swa_page_size: int,
dtype: torch.dtype,
state_dtype: torch.dtype,
c4_state_dtype: torch.dtype,
c128_state_dtype: torch.dtype,
qk_nope_head_dim: int,
qk_rope_head_dim: int,
indexer_head_dim: int,
@@ -494,7 +495,8 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
self.c128_size = c128_size
self.c4_state_pool_size = c4_state_pool_size
self.c128_state_pool_size = c128_state_pool_size
self.state_dtype = state_dtype
self.c4_state_dtype = c4_state_dtype
self.c128_state_dtype = c128_state_dtype
self.compression_ratios = compression_ratios
self.online_mtp_max_draft_tokens = online_mtp_max_draft_tokens
self.online_c128_mtp_pending_seq_lens: Optional[torch.Tensor] = None
@@ -761,7 +763,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
ring_size=ring_size,
overlap=overlap,
head_dim=self.qk_nope_head_dim + self.qk_rope_head_dim,
dtype=self.state_dtype,
dtype=self.c4_state_dtype if ratio == 4 else self.c128_state_dtype,
device=self.device,
enable_memory_saver=enable_memory_saver,
ratio=ratio,
@@ -779,7 +781,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
overlap=overlap,
head_dim=self.indexer_head_dim,
device=self.device,
dtype=self.state_dtype,
dtype=self.c4_state_dtype,
enable_memory_saver=enable_memory_saver,
ratio=ratio,
swa_page_size=self.swa_page_size,
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
+25 -6
View File
@@ -58,6 +58,24 @@ MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_OVERLAP = 1
logger = logging.getLogger(__name__)
def _get_dsv4_compress_state_dtypes() -> tuple[torch.dtype, torch.dtype]:
dtype_name = envs.SGLANG_DSV4_COMPRESS_STATE_DTYPE.get().strip().lower()
if dtype_name in ("float32", "fp32"):
return torch.float32, torch.float32
if dtype_name in ("bfloat16", "bf16"):
if envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get():
raise ValueError(
"SGLANG_DSV4_COMPRESS_STATE_DTYPE=bf16 is not supported when "
"SGLANG_OPT_USE_ONLINE_COMPRESS=1; online c128 state must stay float32."
)
return torch.bfloat16, torch.bfloat16
raise ValueError(
"Unsupported SGLANG_DSV4_COMPRESS_STATE_DTYPE="
f"{dtype_name!r}. Expected one of: float32, fp32, bfloat16, bf16."
)
_is_npu = is_npu()
_is_hip = is_hip()
@@ -418,7 +436,8 @@ class ModelRunnerKVCacheMixin:
swa_page_size=swa_page_size,
sliding_window=self.model_config.window_size,
dtype=self.kv_cache_dtype,
state_dtype=self.state_dtype,
c4_state_dtype=self.c4_state_dtype,
c128_state_dtype=self.c128_state_dtype,
qk_nope_head_dim=self.model_config.qk_nope_head_dim,
qk_rope_head_dim=self.model_config.qk_rope_head_dim,
indexer_head_dim=self.model_config.index_head_dim,
@@ -947,12 +966,12 @@ class ModelRunnerKVCacheMixin:
self.c4_state_pool_size = config.c4_state_pool_size
self.c128_state_pool_size = config.c128_state_pool_size
# state_dtype is a DSV4 architectural constant (fp32 for c4/c128
# state buffers); set unconditionally so draft workers have it before
# _init_pools reads it (target path also overwrites this in the
# configurator's resolve() for parity, harmless here).
# Draft worker does not own the compression-state pools, but keep the
# dtype attributes initialized so _init_pools can share one code path.
if is_deepseek_v4(self.model_config.hf_config):
self.state_dtype = torch.float32
self.c4_state_dtype, self.c128_state_dtype = (
_get_dsv4_compress_state_dtypes()
)
self._init_pools()
python/sglang/srt/model_executor/pool_configurator.py
+23 -4
View File
@@ -62,6 +62,23 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
def _get_dsv4_compress_state_dtype_sizes() -> tuple[int, int]:
dtype_name = envs.SGLANG_DSV4_COMPRESS_STATE_DTYPE.get().strip().lower()
if dtype_name in ("float32", "fp32"):
return 4, 4
if dtype_name in ("bfloat16", "bf16"):
if envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get():
raise ValueError(
"SGLANG_DSV4_COMPRESS_STATE_DTYPE=bf16 is not supported when "
"SGLANG_OPT_USE_ONLINE_COMPRESS=1; online c128 state must stay float32."
)
return 2, 2
raise ValueError(
"Unsupported SGLANG_DSV4_COMPRESS_STATE_DTYPE="
f"{dtype_name!r}. Expected one of: float32, fp32, bfloat16, bf16."
)
class MemoryPoolConfigurator:
"""Base class for memory pool configurators.
@@ -419,16 +436,18 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
)
attn_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
state_dtype_size = 4
c4_state_bytes = 2 * 2 * attn_head_dim * state_dtype_size
c4_state_dtype_size, c128_state_dtype_size = (
_get_dsv4_compress_state_dtype_sizes()
)
c4_state_bytes = 2 * 2 * attn_head_dim * c4_state_dtype_size
# Online c128 stores (max, sum, kv) per slot (3*head_dim) instead of
# raw (kv, score) (2*head_dim). Combined with ring_size=1 this still
# nets a large reduction (~3/256x) but the per-slot bytes go up.
c128_online = envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get()
c128_state_bytes = (
(3 if c128_online else 2 * 1) * attn_head_dim * state_dtype_size
(3 if c128_online else 2 * 1) * attn_head_dim * c128_state_dtype_size
)
c4_indexer_state_bytes = 2 * 2 * self.indexer_head_dim * state_dtype_size
c4_indexer_state_bytes = 2 * 2 * self.indexer_head_dim * c4_state_dtype_size
c4_state_ratio = self.c4_ring_size / self.swa_page_size
c128_state_ratio = self.c128_ring_size / self.swa_page_size
python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py
+2 -1
View File
@@ -390,7 +390,8 @@ class MockDSV4ModelRunner:
page_size=case.page_size,
swa_page_size=DSV4_SWA_WINDOW,
dtype=torch.float8_e4m3fn,
state_dtype=dtype,
c4_state_dtype=dtype,
c128_state_dtype=dtype,
qk_nope_head_dim=DSV4_QK_NOPE_HEAD_DIM,
qk_rope_head_dim=DSV4_QK_ROPE_HEAD_DIM,
indexer_head_dim=128,
test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py
+1082
View File
File diff suppressed because it is too large Load Diff