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

[AMD][DSV4] Reland unified-KV pool sizing and SWA ring accounting, fully gated (#38192)

Browse Source 
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
This commit is contained in:
yuttian1
2026-09-07 13:13:04 -07:00
committed by GitHub
co-authored by hnyls2002 Liangsheng Yin
parent 6287ebf43a
commit 570087ceda
20 changed files with 763 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh
+26 -15
View File
@@ -50,6 +50,7 @@ struct Prefill0Params {
/// \brief Trailing tokens the write plan keeps resident in the compress state ring.
/// Derived from the ring in `plan_compress_prefill`; see the bound there.
int32_t mtp_pad;
bool use_req_ring;
};
struct Prefill1Params {
@@ -67,6 +68,7 @@ struct Prefill1Params {
int32_t swa_page_size;
int32_t ring_size;
int32_t compress_ratio;
bool use_req_ring;
};
struct DecodeParams {
@@ -80,6 +82,7 @@ struct DecodeParams {
int32_t swa_page_size;
int32_t ring_size;
int32_t compress_ratio;
bool use_req_ring;
};
struct Prefill1ParamsLegacy {
@@ -203,7 +206,7 @@ __global__ __launch_bounds__(1024, 1) //
const int32_t last_c_pos = (sl / cr) * cr;
const int32_t first_w_pos = min(last_c_pos - (is_overlap ? cr : 0), sl - params.mtp_pad);
bool do_write = position >= first_w_pos;
if (!do_write && is_overlap) do_write = (position % sps) >= (sps - cr);
if (!do_write && is_overlap && !params.use_req_ring) do_write = (position % sps) >= (sps - cr);
if (do_write) {
const uint32_t out_idx = atomicAdd(&counter_w, 1u);
params.plan_w[out_idx] = pack_w(ragged_id, batch_id, position + 1);
@@ -236,7 +239,7 @@ __global__ __launch_bounds__(1024, 1) //
}
bool do_write = position >= first_w_pos;
if (!do_write && is_overlap) do_write = (position % sps) >= (sps - cr);
if (!do_write && is_overlap && !params.use_req_ring) do_write = (position % sps) >= (sps - cr);
if (do_write) {
const uint32_t out_idx = atomicAdd(&counter_w, 1u);
params.plan_w[out_idx] = pack_w(ragged_id, static_cast<uint32_t>(batch_id), position + 1);
@@ -270,7 +273,7 @@ __global__ void plan_compress_prefill_kernel_1(const Prefill1Params params) {
const auto ring_offset = swa_loc % params.ring_size;
return swa_page * params.ring_size + ring_offset;
};
const auto compute_c128_loc = [&](int64_t rid, int32_t position) {
const auto compute_req_ring_loc = [&](int64_t rid, int32_t position) {
return static_cast<int32_t>(rid * params.ring_size + position % params.ring_size);
};
@@ -283,9 +286,9 @@ __global__ void plan_compress_prefill_kernel_1(const Prefill1Params params) {
const auto position_1 = static_cast<int32_t>(plan_c.seq_len - 1);
// only used for c4, harmless for c128
const auto position_0 = max(position_1 - params.compress_ratio, 0);
if (params.compress_ratio == 128) {
plan_c.read_page_0 = compute_c128_loc(rid, position_0) / 128;
plan_c.read_page_1 = compute_c128_loc(rid, position_1) / 128;
if (params.compress_ratio == 128 || params.use_req_ring) {
plan_c.read_page_0 = compute_req_ring_loc(rid, position_0) / params.compress_ratio;
plan_c.read_page_1 = compute_req_ring_loc(rid, position_1) / params.compress_ratio;
} else {
const auto raw_loc_0 = mapping[position_0];
const auto raw_loc_1 = mapping[position_1];
@@ -307,8 +310,8 @@ __global__ void plan_compress_prefill_kernel_1(const Prefill1Params params) {
// `seq_len` (`write_loc`) may not be aligned here
const auto position = static_cast<int32_t>(plan_w.write_loc - 1);
plan_w.ragged_id = ragged_id;
if (params.compress_ratio == 128) {
plan_w.write_loc = compute_c128_loc(rid, position);
if (params.compress_ratio == 128 || params.use_req_ring) {
plan_w.write_loc = compute_req_ring_loc(rid, position);
} else {
const auto raw_loc = mapping[position];
plan_w.write_loc = compute_loc(params.f2s_ptr[raw_loc]);
@@ -329,7 +332,7 @@ __global__ void plan_compress_decode_kernel(const DecodeParams params) {
const auto ring_offset = swa_loc % params.ring_size;
return swa_page * params.ring_size + ring_offset;
};
const auto compute_c128_loc = [&](int64_t rid, int32_t position) {
const auto compute_req_ring_loc = [&](int64_t rid, int32_t position) {
return static_cast<int32_t>(rid * params.ring_size + position % params.ring_size);
};
const auto seq_len = static_cast<int32_t>(params.seq_ptr[idx]);
@@ -338,10 +341,10 @@ __global__ void plan_compress_decode_kernel(const DecodeParams params) {
int32_t write_loc;
int32_t read_page_0;
int32_t read_page_1;
if (params.compress_ratio == 128) {
write_loc = compute_c128_loc(rid, position_1);
read_page_0 = compute_c128_loc(rid, position_0) / 128;
read_page_1 = compute_c128_loc(rid, position_1) / 128;
if (params.compress_ratio == 128 || params.use_req_ring) {
write_loc = compute_req_ring_loc(rid, position_1);
read_page_0 = compute_req_ring_loc(rid, position_0) / params.compress_ratio;
read_page_1 = compute_req_ring_loc(rid, position_1) / params.compress_ratio;
} else {
const auto raw_loc_0 = mapping[position_0];
const auto raw_loc_1 = mapping[position_1];
@@ -461,6 +464,7 @@ inline PrefillPlan plan_compress_prefill(
const int32_t compress_ratio,
const int32_t swa_page_size,
const int32_t ring_size,
const bool use_req_ring,
const bool use_cuda_graph) {
auto B = SymbolicSize{"batch_size"};
auto N = SymbolicSize{"num_q_tokens"};
@@ -503,6 +507,7 @@ inline PrefillPlan plan_compress_prefill(
const auto batch_size = static_cast<uint32_t>(B.unwrap());
constexpr auto kMaxTokens = static_cast<uint32_t>(std::numeric_limits<uint16_t>::max());
RuntimeCheck(compress_ratio == 4 || compress_ratio == 128);
RuntimeCheck(!use_req_ring || compress_ratio == 4);
RuntimeCheck(batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens);
// `swa_page_size` >= `ring_size` >= `compress_ratio`
RuntimeCheck(swa_page_size % ring_size == 0 && ring_size % compress_ratio == 0);
@@ -537,6 +542,7 @@ inline PrefillPlan plan_compress_prefill(
.compress_ratio = compress_ratio,
.swa_page_size = swa_page_size,
.mtp_pad = mtp_pad,
.use_req_ring = use_req_ring,
};
LaunchKernel(1, kMaxPrefillBatchSize, device)(plan_compress_prefill_kernel0, params0);
// kernel_1 sees the already-padded buffers, so num_c == num_w == num_padded == num_q_tokens.
@@ -555,6 +561,7 @@ inline PrefillPlan plan_compress_prefill(
.swa_page_size = swa_page_size,
.ring_size = ring_size,
.compress_ratio = compress_ratio,
.use_req_ring = use_req_ring,
};
const auto block_size_1 = 256;
const auto num_blocks_1 = div_ceil(params1.num_work, block_size_1);
@@ -582,7 +589,7 @@ inline PrefillPlan plan_compress_prefill(
RuntimeCheck(0 < extend_len && extend_len <= seq_len);
const auto should_write = [=](int32_t position) {
if (position >= first_w_pos) return true;
return is_overlap && position % swa_page_size >= (swa_page_size - compress_ratio);
return is_overlap && !use_req_ring && position % swa_page_size >= (swa_page_size - compress_ratio);
};
for (const auto j : irange(extend_len)) {
const int32_t position = prefix_len + j;
@@ -631,6 +638,7 @@ inline PrefillPlan plan_compress_prefill(
.swa_page_size = swa_page_size,
.ring_size = ring_size,
.compress_ratio = compress_ratio,
.use_req_ring = use_req_ring,
};
const auto block_size = 256;
const auto num_blocks = div_ceil(params.num_work, block_size);
@@ -645,7 +653,8 @@ inline tvm::ffi::Tensor plan_compress_decode(
const tvm::ffi::TensorView seq_lens, // CPU/GPU
const int32_t compress_ratio,
const int32_t swa_page_size,
const int32_t ring_size) {
const int32_t ring_size,
const bool use_req_ring) {
auto B = SymbolicSize{"batch_size"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLGPU>();
@@ -667,6 +676,7 @@ inline tvm::ffi::Tensor plan_compress_decode(
.with_device(device_)
.verify(seq_lens);
RuntimeCheck(!use_req_ring || compress_ratio == 4);
const auto batch_size = static_cast<uint32_t>(B.unwrap());
const auto device = device_.unwrap();
auto D = ffi::empty({batch_size, sizeof(PlanD)}, kDLUInt8, device);
@@ -681,6 +691,7 @@ inline tvm::ffi::Tensor plan_compress_decode(
.swa_page_size = swa_page_size,
.ring_size = ring_size,
.compress_ratio = compress_ratio,
.use_req_ring = use_req_ring,
};
const auto block_size = 256;
const auto num_blocks = div_ceil(batch_size, block_size);
python/sglang/kernels/ops/attention/dsv4/attn.py
+4 -1
View File
@@ -100,6 +100,7 @@ def create_paged_compress_data_kernel(
stride_out_1_1: tl.constexpr,
compress_ratio: tl.constexpr,
is_overlap: tl.constexpr,
use_req_ring: tl.constexpr,
swa_page_size: tl.constexpr,
ring_size: tl.constexpr,
BLOCK: tl.constexpr,
@@ -133,7 +134,7 @@ def create_paged_compress_data_kernel(
else:
pos = write_overlap_pos
pos = tl.maximum(pos, 0)
if compress_ratio == 128:
if compress_ratio == 128 or use_req_ring:
state_loc = rid * ring_size + (pos % ring_size)
else:
loc = tl.load(
@@ -182,6 +183,7 @@ def triton_create_paged_compress_data(
extend_seq_lens: torch.Tensor,
req_to_token: torch.Tensor,
full_to_swa_index_mapping: torch.Tensor,
use_req_ring: bool = False,
block: int = 128,
) -> Tuple[torch.Tensor, torch.Tensor]:
batch_size = req_pool_indices.shape[0]
@@ -205,6 +207,7 @@ def triton_create_paged_compress_data(
stride_out_1_1=out_1.stride(1), # type: ignore
compress_ratio=compress_ratio, # type: ignore
is_overlap=1 if is_overlap else 0, # type: ignore
use_req_ring=1 if use_req_ring else 0, # type: ignore
swa_page_size=swa_page_size, # type: ignore
ring_size=ring_size, # type: ignore
BLOCK=block, # type: ignore
python/sglang/kernels/ops/attention/dsv4/compress.py
+16 -3
View File
@@ -162,6 +162,7 @@ class CompressorDecodePlan(NamedTuple):
seq_lens: torch.Tensor,
swa_page_size: int,
ring_size: int,
use_req_ring: bool = False,
) -> CompressorDecodePlan:
if _is_xpu:
fn = plan_compress_decode
@@ -169,7 +170,7 @@ class CompressorDecodePlan(NamedTuple):
module = _jit_compress_plan_module()
fn = module.plan_decode
plan_d = fn(
args = (
req_pool_indices,
req_to_token,
full_to_state,
@@ -178,6 +179,10 @@ class CompressorDecodePlan(NamedTuple):
int(swa_page_size),
int(ring_size),
)
assert not (_is_xpu and use_req_ring), (
"use_req_ring is not supported by the XPU compress plan builder"
)
plan_d = fn(*args) if _is_xpu else fn(*args, bool(use_req_ring))
return CompressorDecodePlan(compress_ratio, torch.from_dlpack(plan_d))
@staticmethod
@@ -247,6 +252,7 @@ class CompressorPrefillPlan(NamedTuple):
ring_size: int,
num_q_tokens: int,
use_cuda_graph: bool = False,
use_req_ring: bool = False,
) -> CompressorPrefillPlan:
is_gpu_input = seq_lens.device.type in ["cuda", "xpu"]
pin_buffer = torch.empty(
@@ -274,7 +280,7 @@ class CompressorPrefillPlan(NamedTuple):
module = _jit_compress_plan_module()
fn = module.plan_prefill
plan_c, plan_w = fn(
args = (
req_pool_indices,
req_to_token,
full_to_state,
@@ -285,7 +291,14 @@ class CompressorPrefillPlan(NamedTuple):
int(compress_ratio),
int(swa_page_size),
int(ring_size),
bool(use_cuda_graph),
)
assert not (_is_xpu and use_req_ring), (
"use_req_ring is not supported by the XPU compress plan builder"
)
plan_c, plan_w = (
fn(*args, bool(use_cuda_graph))
if _is_xpu
else fn(*args, bool(use_req_ring), bool(use_cuda_graph))
)
return CompressorPrefillPlan(
compress_ratio,