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

Revert "[AMD] Fix DeepSeek V4 Pro c128 state tensor dtype mismatch error and c4_sparse_raw_indices attribute error in cuda graph phase" (#27919)

Browse Source
This commit is contained in:
Wang, FangYuan
2026-06-11 14:25:32 -07:00
committed by GitHub
parent 12d3f02be8
commit 6e885c844f
4 changed files with 88 additions and 177 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh
+39 -65
View File
@@ -89,10 +89,10 @@ struct C128Trait {
static_assert(kHeadDim % kTileDim == 0);
};
template <typename Trait, bool kUsePDL, typename BufFloat, typename InFloat, typename OutFloat>
template <typename Trait, bool kUsePDL, typename InFloat, typename OutFloat>
SGL_DEVICE void c128_forward(
const BufFloat* kv_buf, // [128n, 128n + 127]
const InFloat* kv_src, // ragged pointer at position = 128n + 127
const InFloat* kv_buf, // [128n, 128n + 127]
const InFloat* kv_src, // ragged pointer at position = 128n + 127
OutFloat* kv_out,
const InFloat* score_bias,
const int32_t buffer_len) {
@@ -101,15 +101,11 @@ SGL_DEVICE void c128_forward(
const auto warp_id = threadIdx.x / kWarpThreads;
const auto lane_id = threadIdx.x % kWarpThreads;
/// NOTE: part 1: load kv + score. kv_score_buffer (fp32, runtime state pool)
/// keeps its own BufFloat dtype; input/ape share InFloat (ape is cast to bf16
/// at load). Every value is converted to fp32 right after load.
using StorageBuf = AlignedVector<BufFloat, kTileElements>;
/// NOTE: part 1: load kv + score
using StorageIn = AlignedVector<InFloat, kTileElements>;
const auto gmem_buf = tile::Memory<StorageBuf>{lane_id, kWarpThreads};
const auto gmem_in = tile::Memory<StorageIn>{lane_id, kWarpThreads};
float kv[kElementsPerWarp][kTileElements];
float score[kElementsPerWarp][kTileElements];
StorageIn kv[kElementsPerWarp];
StorageIn score[kElementsPerWarp];
StorageIn bias[kElementsPerWarp];
const int32_t warp_offset = warp_id * kElementsPerWarp;
@@ -125,23 +121,9 @@ SGL_DEVICE void c128_forward(
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 k = gmem_buf.load(kv_buf + j * Trait::kElementSize);
const auto s = gmem_buf.load(kv_buf + j * Trait::kElementSize + Trait::kScoreOffset);
#pragma unroll
for (int32_t t = 0; t < kTileElements; ++t) {
kv[i][t] = cast<float>(k[t]);
score[i][t] = cast<float>(s[t]);
}
} else {
const auto k = gmem_in.load(kv_start + j * Trait::kElementSize);
const auto s = gmem_in.load(kv_start + j * Trait::kElementSize + Trait::kScoreOffset);
#pragma unroll
for (int32_t t = 0; t < kTileElements; ++t) {
kv[i][t] = cast<float>(k[t]);
score[i][t] = cast<float>(s[t]);
}
}
const auto src = j < buffer_len ? kv_buf : kv_start;
kv[i] = gmem_in.load(src + j * Trait::kElementSize);
score[i] = gmem_in.load(src + j * Trait::kElementSize + Trait::kScoreOffset);
}
/// NOTE: part 2: safe online softmax + weighted sum
@@ -156,11 +138,11 @@ SGL_DEVICE void c128_forward(
float score_fp32[kTileElements][kElementsPerWarp];
// kv/score already fp32 (converted at load); just add the bias
// convert to fp32 and apply bias first
#pragma unroll
for (int32_t i = 0; i < kTileElements; ++i) {
for (int32_t j = 0; j < kElementsPerWarp; ++j) {
score_fp32[i][j] = score[j][i] + cast<float>(bias[j][i]);
score_fp32[i][j] = cast<float>(score[j][i]) + cast<float>(bias[j][i]);
}
}
@@ -181,7 +163,7 @@ SGL_DEVICE void c128_forward(
for (int32_t j = 0; j < 8; ++j) {
const auto fp32_score = score[j];
const auto exp_score = expf(fp32_score - max_value);
sum_product += kv[j][i] * exp_score;
sum_product += cast<float>(kv[j][i]) * exp_score;
sum_exp_value += exp_score;
}
@@ -233,27 +215,25 @@ SGL_DEVICE void c128_forward(
}
}
template <typename Trait, typename BufFloat, typename InFloat>
SGL_DEVICE void c128_write_decode(BufFloat* kv_buf, const InFloat* kv_src) {
template <typename Trait, typename InFloat>
SGL_DEVICE void c128_write_decode(InFloat* kv_buf, const InFloat* kv_src) {
using namespace device;
using StorageIn = AlignedVector<InFloat, kTileElements>;
using StorageBuf = AlignedVector<BufFloat, kTileElements>;
const auto gmem_in = tile::Memory<StorageIn>::warp();
const auto gmem_buf = tile::Memory<StorageBuf>::warp();
using Storage = AlignedVector<InFloat, kTileElements>;
const auto gmem = tile::Memory<Storage>::warp();
Storage data[2];
#pragma unroll
for (int32_t i = 0; i < 2; ++i) {
const auto d = gmem_in.load(kv_src + Trait::kHeadDim * i);
StorageBuf o;
data[i] = gmem.load(kv_src + Trait::kHeadDim * i);
}
#pragma unroll
for (int32_t t = 0; t < kTileElements; ++t)
o[t] = cast<BufFloat>(d[t]);
gmem_buf.store(kv_buf + Trait::kHeadDim * i, o);
for (int32_t i = 0; i < 2; ++i) {
gmem.store(kv_buf + Trait::kHeadDim * i, data[i]);
}
}
template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
C128_KERNEL void flash_c128_decode(const __grid_constant__ Compress128DecodeParams params) {
using namespace device;
using Trait = C128Trait<kHeadDim>;
@@ -267,7 +247,7 @@ C128_KERNEL void flash_c128_decode(const __grid_constant__ Compress128DecodePara
const auto plan = params.plan_d[global_bid];
const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
const auto kv_buffer = static_cast<BufFloat*>(params.kv_buffer) + split_offset;
const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
const auto score_bias = static_cast<const InFloat*>(params.score_bias) + split_offset;
const auto kv_src = kv_input + global_bid * Trait::kElementSize;
@@ -278,15 +258,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, BufFloat, InFloat>(kv_dst, kv_src);
c128_write_decode<Trait>(kv_dst, kv_src);
}
if (plan.write_loc % 128 == 127) {
c128_forward<Trait, kUsePDL, BufFloat, InFloat, OutFloat>(kv_buf, kv_src, kv_out, score_bias, 128);
c128_forward<Trait, kUsePDL>(kv_buf, kv_src, kv_out, score_bias, 128);
}
}
// compress kernel
template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
C128_KERNEL void flash_c128_prefill(const __grid_constant__ Compress128PrefillParams params) {
using namespace device;
using Trait = C128Trait<kHeadDim>;
@@ -299,7 +279,7 @@ C128_KERNEL void flash_c128_prefill(const __grid_constant__ Compress128PrefillPa
const auto plan = params.plan_c[global_pid];
const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
const auto kv_buffer = static_cast<BufFloat*>(params.kv_buffer) + split_offset;
const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
const auto score_bias = static_cast<const InFloat*>(params.score_bias) + split_offset;
if (plan.is_invalid()) return;
@@ -308,15 +288,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, BufFloat, InFloat, OutFloat>(kv_buf, kv_src, kv_out, score_bias, plan.buffer_len);
c128_forward<Trait, kUsePDL>(kv_buf, kv_src, kv_out, score_bias, plan.buffer_len);
}
template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
WRITE_KERNEL void write_c128_prefill(const __grid_constant__ Compress128PrefillParams params) {
using namespace device;
using Trait = C128Trait<kHeadDim>;
using StorageIn = AlignedVector<InFloat, kTileElements>;
using StorageBuf = AlignedVector<BufFloat, kTileElements>;
const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
const uint32_t global_wid = global_tid / kWarpThreads; // warp id
@@ -329,37 +308,32 @@ WRITE_KERNEL void write_c128_prefill(const __grid_constant__ Compress128PrefillP
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<BufFloat*>(params.kv_buffer) + split_offset;
const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
if (plan.is_invalid()) return;
// each warp will handle a contiguous region
const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize;
const auto kv_buf = kv_buffer + plan.write_loc * Trait::kElementSize;
const auto gmem_in = tile::Memory<StorageIn>::warp();
const auto gmem_buf = tile::Memory<StorageBuf>::warp();
const auto gmem = tile::Memory<StorageIn>::warp();
PDLWaitPrimary<kUsePDL>();
StorageIn data[2];
#pragma unroll
for (int32_t i = 0; i < 2; ++i) {
data[i] = gmem_in.load(kv_src, i);
data[i] = gmem.load(kv_src, i);
}
PDLTriggerSecondary<kUsePDL>();
#pragma unroll
for (int32_t i = 0; i < 2; ++i) {
StorageBuf o;
#pragma unroll
for (int32_t t = 0; t < kTileElements; ++t)
o[t] = cast<BufFloat>(data[i][t]);
gmem_buf.store(kv_buf, o, i);
gmem.store(kv_buf, data[i], i);
}
}
template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
struct FlashCompress128Kernel {
static constexpr auto decode_kernel = flash_c128_decode<kHeadDim, BufFloat, InFloat, OutFloat, kUsePDL>;
static constexpr auto prefill_c_kernel = flash_c128_prefill<kHeadDim, BufFloat, InFloat, OutFloat, kUsePDL>;
static constexpr auto prefill_w_kernel = write_c128_prefill<kHeadDim, BufFloat, InFloat, OutFloat, kUsePDL>;
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 int64_t kTileDim = kTileElements * device::kWarpThreads; // 64
static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
using Trait = C128Trait<kHeadDim>;
@@ -377,7 +351,7 @@ struct FlashCompress128Kernel {
device_.set_options<kDLGPU>();
TensorMatcher({-1, 128, Trait::kElementSize}) // kv score
.with_dtype<BufFloat>()
.with_dtype<InFloat>()
.with_device(device_)
.verify(kv_buffer);
TensorMatcher({N, Trait::kElementSize}) // kv score input
@@ -424,7 +398,7 @@ struct FlashCompress128Kernel {
device_.set_options<kDLGPU>();
TensorMatcher({-1, 128, Trait::kElementSize}) // kv score
.with_dtype<BufFloat>()
.with_dtype<InFloat>()
.with_device(device_)
.verify(kv_buffer);
TensorMatcher({N, Trait::kElementSize}) // kv score input (ragged)
python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh
+46 -93
View File
@@ -74,27 +74,23 @@ struct C4Trait {
static_assert(kHeadDim % kTileDim == 0);
};
template <typename Trait, bool kUsePDL, typename BufFloat, typename InFloat, typename OutFloat>
template <typename Trait, bool kUsePDL, typename InFloat, typename OutFloat>
SGL_DEVICE void c4_forward(
const BufFloat* kv_buf_0, // overlap [4n - 4, 4n - 1]
const BufFloat* kv_buf_1, // normal [4n + 0, 4n + 3]
const InFloat* kv_src, // ragged pointer at position = 4n + 3
const InFloat* kv_buf_0, // overlap [4n - 4, 4n - 1]
const InFloat* kv_buf_1, // normal [4n + 0, 4n + 3]
const InFloat* kv_src, // ragged pointer at position = 4n + 3
OutFloat* kv_out,
const InFloat* score_bias,
const bool should_overlap,
const int32_t buffer_len) {
using namespace device;
/// NOTE: part 1: load kv + score. kv_score_buffer (fp32, runtime state pool)
/// keeps its own BufFloat dtype; input/ape share InFloat (ape is cast to bf16
/// at load). Values are converted to fp32 at load.
using StorageBuf = AlignedVector<BufFloat, kTileElements>;
/// NOTE: part 1: load kv + score
using StorageIn = AlignedVector<InFloat, kTileElements>;
/// NOTE: load one tile_dim (< head_dim) at at time
const auto gmem_buf = tile::Memory<StorageBuf>::warp();
const auto gmem_in = tile::Memory<StorageIn>::warp();
float kv[8][kTileElements];
float score[8][kTileElements];
StorageIn kv[8];
StorageIn score[8];
StorageIn bias[8];
#pragma unroll
@@ -106,61 +102,28 @@ SGL_DEVICE void c4_forward(
const auto kv_start = kv_src - 7 * Trait::kElementSize; // point to start
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
if (i < buffer_len) {
const auto base = kv_buf_0 + i * Trait::kElementSize;
const auto k = gmem_buf.load(base);
const auto s = gmem_buf.load(base + Trait::kScoreOffset);
#pragma unroll
for (int32_t t = 0; t < kTileElements; ++t) {
kv[i][t] = cast<float>(k[t]);
score[i][t] = cast<float>(s[t]);
}
} else {
const auto base = kv_start + i * Trait::kElementSize;
const auto k = gmem_in.load(base);
const auto s = gmem_in.load(base + Trait::kScoreOffset);
#pragma unroll
for (int32_t t = 0; t < kTileElements; ++t) {
kv[i][t] = cast<float>(k[t]);
score[i][t] = cast<float>(s[t]);
}
}
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) {
#pragma unroll
for (int32_t t = 0; t < kTileElements; ++t) {
kv[i][t] = 0.0f;
score[i][t] = kFloatNegInf;
}
kv[i].fill(cast<InFloat>(0.0f));
score[i].fill(cast<InFloat>(kFloatNegInf));
}
}
const auto kv_start = kv_src - 3 * Trait::kElementSize; // point to start
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
if (i + 4 < buffer_len) {
const auto base = kv_buf_1 + i * Trait::kElementSize + Trait::kOverlapOffset;
const auto k = gmem_buf.load(base);
const auto s = gmem_buf.load(base + Trait::kScoreOffset);
#pragma unroll
for (int32_t t = 0; t < kTileElements; ++t) {
kv[i + 4][t] = cast<float>(k[t]);
score[i + 4][t] = cast<float>(s[t]);
}
} else {
const auto base = kv_start + i * Trait::kElementSize + Trait::kOverlapOffset;
const auto k = gmem_in.load(base);
const auto s = gmem_in.load(base + Trait::kScoreOffset);
#pragma unroll
for (int32_t t = 0; t < kTileElements; ++t) {
kv[i + 4][t] = cast<float>(k[t]);
score[i + 4][t] = cast<float>(s[t]);
}
}
const auto src = i + 4 < buffer_len ? kv_buf_1 : kv_start;
const auto base = src + i * Trait::kElementSize + Trait::kOverlapOffset;
kv[i + 4] = gmem_in.load(base);
score[i + 4] = gmem_in.load(base + Trait::kScoreOffset);
}
/// NOTE: part 2: safe online softmax + weighted sum
@@ -171,11 +134,11 @@ SGL_DEVICE void c4_forward(
// consume 32 fp registers
float score_fp32[kTileElements][8];
// kv/score already fp32 (converted at load); just add the bias
// convert to fp32 and apply bias first
#pragma unroll
for (int32_t i = 0; i < kTileElements; ++i) {
for (int32_t j = 0; j < 8; ++j) {
score_fp32[i][j] = score[j][i] + cast<float>(bias[j][i]);
score_fp32[i][j] = cast<float>(score[j][i]) + cast<float>(bias[j][i]);
}
}
@@ -196,7 +159,7 @@ SGL_DEVICE void c4_forward(
for (int32_t j = 0; j < 8; ++j) {
const auto fp32_score = score[j];
const auto exp_score = expf(fp32_score - max_value);
sum_product += kv[j][i] * exp_score;
sum_product += cast<float>(kv[j][i]) * exp_score;
sum_exp_value += exp_score;
}
@@ -208,27 +171,25 @@ SGL_DEVICE void c4_forward(
gmem_out.store(kv_out, result);
}
template <typename Trait, typename BufFloat, typename InFloat>
SGL_DEVICE void c4_write_decode(BufFloat* kv_buf, const InFloat* kv_src) {
template <typename Trait, typename InFloat>
SGL_DEVICE void c4_write_decode(InFloat* kv_buf, const InFloat* kv_src) {
using namespace device;
using StorageIn = AlignedVector<InFloat, kTileElements>;
using StorageBuf = AlignedVector<BufFloat, kTileElements>;
const auto gmem_in = tile::Memory<StorageIn>::warp();
const auto gmem_buf = tile::Memory<StorageBuf>::warp();
const auto gmem = tile::Memory<StorageIn>::warp();
StorageIn data[4];
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
const auto d = gmem_in.load(kv_src + Trait::kHeadDim * i);
StorageBuf o;
data[i] = gmem.load(kv_src + Trait::kHeadDim * i);
}
#pragma unroll
for (int32_t t = 0; t < kTileElements; ++t)
o[t] = cast<BufFloat>(d[t]);
gmem_buf.store(kv_buf + Trait::kHeadDim * i, o);
for (int32_t i = 0; i < 4; ++i) {
gmem.store(kv_buf + Trait::kHeadDim * i, data[i]);
}
}
template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
C4_KERNEL void flash_c4_decode(const __grid_constant__ Compress4DecodeParams params) {
using namespace device;
using Trait = C4Trait<kHeadDim>;
@@ -243,7 +204,7 @@ C4_KERNEL void flash_c4_decode(const __grid_constant__ Compress4DecodeParams par
const auto plan = params.plan_d[global_bid];
const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
const auto kv_buffer = static_cast<BufFloat*>(params.kv_buffer) + split_offset;
const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
const auto score_bias = static_cast<const InFloat*>(params.score_bias) + split_offset;
const auto kv_src = kv_input + global_bid * Trait::kElementSize;
@@ -253,15 +214,14 @@ 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, BufFloat, InFloat>(kv_dst, kv_src);
c4_write_decode<Trait>(kv_dst, kv_src);
if (plan.seq_len % 4 == 0) {
const auto need_overlap = plan.seq_len > 4;
c4_forward<Trait, kUsePDL, BufFloat, InFloat, OutFloat>(
kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, 8);
c4_forward<Trait, kUsePDL>(kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, 8);
}
}
template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
C4_KERNEL void flash_c4_prefill(const __grid_constant__ Compress4PrefillParams params) {
using namespace device;
using Trait = C4Trait<kHeadDim>;
@@ -276,7 +236,7 @@ C4_KERNEL void flash_c4_prefill(const __grid_constant__ Compress4PrefillParams p
const auto plan = params.plan_c[global_pid];
const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
const auto kv_buffer = static_cast<BufFloat*>(params.kv_buffer) + split_offset;
const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
const auto score_bias = static_cast<const InFloat*>(params.score_bias) + split_offset;
if (plan.is_invalid()) return;
@@ -287,16 +247,14 @@ 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, BufFloat, InFloat, OutFloat>(
kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, plan.buffer_len);
c4_forward<Trait, kUsePDL>(kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, plan.buffer_len);
}
template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
WRITE_KERNEL void write_c4_prefill(const __grid_constant__ Compress4PrefillParams params) {
using namespace device;
using Trait = C4Trait<kHeadDim>;
using StorageIn = AlignedVector<InFloat, kTileElements>;
using StorageBuf = AlignedVector<BufFloat, kTileElements>;
const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
const uint32_t global_wid = global_tid / kWarpThreads; // warp id
@@ -309,37 +267,32 @@ WRITE_KERNEL void write_c4_prefill(const __grid_constant__ Compress4PrefillParam
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<BufFloat*>(params.kv_buffer) + split_offset;
const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
if (plan.is_invalid()) return;
// each warp will handle a contiguous region
const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize;
const auto kv_buf = kv_buffer + plan.write_loc * Trait::kElementSize;
const auto gmem_in = tile::Memory<StorageIn>::warp();
const auto gmem_buf = tile::Memory<StorageBuf>::warp();
const auto gmem = tile::Memory<StorageIn>::warp();
PDLWaitPrimary<kUsePDL>();
StorageIn data[4];
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
data[i] = gmem_in.load(kv_src, i);
data[i] = gmem.load(kv_src, i);
}
PDLTriggerSecondary<kUsePDL>();
#pragma unroll
for (int32_t i = 0; i < 4; ++i) {
StorageBuf o;
#pragma unroll
for (int32_t t = 0; t < kTileElements; ++t)
o[t] = cast<BufFloat>(data[i][t]);
gmem_buf.store(kv_buf, o, i);
gmem.store(kv_buf, data[i], i);
}
}
template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
struct FlashCompress4Kernel {
static constexpr auto decode_kernel = flash_c4_decode<kHeadDim, BufFloat, InFloat, OutFloat, kUsePDL>;
static constexpr auto prefill_c_kernel = flash_c4_prefill<kHeadDim, BufFloat, InFloat, OutFloat, kUsePDL>;
static constexpr auto prefill_w_kernel = write_c4_prefill<kHeadDim, BufFloat, InFloat, OutFloat, kUsePDL>;
static constexpr auto decode_kernel = flash_c4_decode<kHeadDim, InFloat, OutFloat, kUsePDL>;
static constexpr auto prefill_c_kernel = flash_c4_prefill<kHeadDim, InFloat, OutFloat, kUsePDL>;
static constexpr auto prefill_w_kernel = write_c4_prefill<kHeadDim, InFloat, OutFloat, kUsePDL>;
static constexpr uint32_t kBlockSize = 128;
static constexpr uint32_t kTileDim = kTileElements * device::kWarpThreads;
static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
@@ -359,7 +312,7 @@ struct FlashCompress4Kernel {
device_.set_options<kDLGPU>();
TensorMatcher({-1, 4, Trait::kElementSize}) // kv score
.with_dtype<BufFloat>()
.with_dtype<InFloat>()
.with_device(device_)
.verify(kv_buffer);
TensorMatcher({N, Trait::kElementSize}) // kv score input
@@ -406,7 +359,7 @@ struct FlashCompress4Kernel {
device_.set_options<kDLGPU>();
TensorMatcher({-1, 4, Trait::kElementSize}) // kv score
.with_dtype<BufFloat>()
.with_dtype<InFloat>()
.with_device(device_)
.verify(kv_buffer);
TensorMatcher({N, Trait::kElementSize}) // kv score input (ragged)
python/sglang/jit_kernel/dsv4/compress.py
+3 -15
View File
@@ -44,14 +44,11 @@ def _jit_compress_norm_rope_module(
@cache_once
def _jit_compress_module(
head_dim: int,
dtype_buf: torch.dtype,
dtype_in: torch.dtype,
dtype_out: torch.dtype,
ratio: Literal[4, 128],
) -> Module:
args = make_cpp_args(
head_dim, dtype_buf, dtype_in, dtype_out, is_arch_support_pdl()
)
args = make_cpp_args(head_dim, dtype_in, dtype_out, is_arch_support_pdl())
kernel_class = f"FlashCompress{ratio}Kernel<{args}>"
return load_jit(
make_name(f"compress_{ratio}_v2"),
@@ -327,17 +324,8 @@ def compress_forward(
assert compress_ratio == 128 and head_dim == 512
module = _jit_compress_128_online_module(512)
else:
# kv_score_buffer (fp32 runtime state pool) may differ from input/ape, so
# the kernel keeps a BufFloat template and casts the buffer to fp32 at
# load. ape/weight are cast to bf16 at load (apply_ape_hotfix), matching
# kv_score_input's dtype, so they need no extra template params.
module = _jit_compress_module(
head_dim,
kv_score_buffer.dtype,
kv_score_input.dtype,
out.dtype,
compress_ratio,
)
dtype_in, dtype_out = kv_score_input.dtype, out.dtype
module = _jit_compress_module(head_dim, dtype_in, dtype_out, compress_ratio)
fn = module.decode if plan.is_decode else module.prefill
fn(kv_score_buffer, kv_score_input, out, ape, *plan[1:3])
return out
python/sglang/srt/layers/attention/dsv4/compressor.py
-4
View File
@@ -395,10 +395,6 @@ class Compressor(nn.Module):
ape = torch.cat([ape[0], ape[1]], dim=0)
self.ape.data.copy_(ape.view(self.ratio, -1))
if _use_aiter:
self.ape.data = self.ape.data.to(torch.bfloat16)
self.norm.weight.data = self.norm.weight.data.to(torch.bfloat16)
def get_state_pool(self, attn_backend: AttentionBackend) -> CompressStatePool:
token_to_kv_pool = attn_backend.token_to_kv_pool
assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool)