[PD & HiSparse] Add DeepSeek V4 support for HiSparse direct Prefill-to-Decode DRAM (#24880)

This commit is contained in:
huangtingwei
2026-06-05 15:39:48 +08:00
committed by GitHub
parent 66b932154f
commit 00fefef16b
12 changed files with 478 additions and 309 deletions
@@ -1,82 +0,0 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/deepseek_v4/kvcacheio.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <cstdint>
namespace {
/// NOTE: for offload to cpu kernel, we use persistent kernel
inline constexpr uint32_t kBlockSize = 1024;
inline constexpr uint32_t kBlockQuota = 4;
#define OFFLOAD_KERNEL __global__ __launch_bounds__(kBlockSize, 1)
struct OffloadParams {
void** gpu_caches;
void** cpu_caches;
const int64_t* gpu_indices;
const int64_t* cpu_indices;
uint32_t num_items;
uint32_t num_layers;
};
OFFLOAD_KERNEL void offload_to_cpu(const __grid_constant__ OffloadParams params) {
using namespace device::hisparse;
const auto [gpu_caches, cpu_caches, gpu_indices, cpu_indices, num_items, num_layers] = params;
const auto global_tid = blockIdx.x * blockDim.x + threadIdx.x;
constexpr auto kNumWarps = (kBlockSize / 32) * kBlockQuota;
for (auto i = global_tid / 32; i < num_items; i += kNumWarps) {
const int32_t gpu_index = gpu_indices[i];
const int32_t cpu_index = cpu_indices[i];
for (auto j = 0u; j < num_layers; ++j) {
const auto gpu_cache = gpu_caches[j];
const auto cpu_cache = cpu_caches[j];
transfer_item<TransferDirection::DeviceToHost>(
/*dst_cache=*/cpu_cache,
/*src_cache=*/gpu_cache,
/*dst_index=*/cpu_index,
/*src_index=*/gpu_index);
}
}
}
[[maybe_unused]]
void hisparse_transfer(
tvm::ffi::TensorView gpu_ptrs,
tvm::ffi::TensorView cpu_ptrs,
tvm::ffi::TensorView gpu_indices,
tvm::ffi::TensorView cpu_indices) {
using namespace host;
auto N = SymbolicSize{"num_items"};
auto L = SymbolicSize{"num_layers"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({L}) // 1D cache pointers
.with_dtype<uint64_t>()
.with_device(device_)
.verify(gpu_ptrs)
.verify(cpu_ptrs);
TensorMatcher({N}) // 1D indices
.with_dtype<int64_t>()
.with_device(device_)
.verify(gpu_indices)
.verify(cpu_indices);
const auto params = OffloadParams{
.gpu_caches = static_cast<void**>(gpu_ptrs.data_ptr()),
.cpu_caches = static_cast<void**>(cpu_ptrs.data_ptr()),
.gpu_indices = static_cast<const int64_t*>(gpu_indices.data_ptr()),
.cpu_indices = static_cast<const int64_t*>(cpu_indices.data_ptr()),
.num_items = static_cast<uint32_t>(N.unwrap()),
.num_layers = static_cast<uint32_t>(L.unwrap()),
};
LaunchKernel(kBlockQuota, kBlockSize, device_.unwrap())(offload_to_cpu, params);
}
} // namespace
+61 -4
View File
@@ -52,6 +52,62 @@ transfer_item_warp(int32_t lane_id, const void* src_addr, void* dst_addr, int64_
} }
} }
template <int BLOCK_SIZE>
__global__ __launch_bounds__(BLOCK_SIZE, 1) void transfer_cache_dsv4_mla_kernel(
void** src_caches,
void** dst_caches,
const int64_t* src_indices,
const int64_t* dst_indices,
uint32_t num_items,
uint32_t num_layers) {
const int global_tid = blockIdx.x * blockDim.x + threadIdx.x;
constexpr int NUM_WARPS = BLOCK_SIZE / WARP_SIZE;
const int total_warps = gridDim.x * NUM_WARPS;
for (uint32_t i = global_tid / WARP_SIZE; i < num_items; i += total_warps) {
const int32_t src_index = static_cast<int32_t>(src_indices[i]);
const int32_t dst_index = static_cast<int32_t>(dst_indices[i]);
for (uint32_t layer_id = 0; layer_id < num_layers; ++layer_id) {
device::hisparse::transfer_item(
/*dst_cache=*/dst_caches[layer_id],
/*src_cache=*/src_caches[layer_id],
/*dst_index=*/dst_index,
/*src_index=*/src_index);
}
}
}
template <int BLOCK_SIZE>
void transfer_cache_dsv4_mla(
tvm::ffi::TensorView src_ptrs,
tvm::ffi::TensorView dst_ptrs,
tvm::ffi::TensorView src_indices,
tvm::ffi::TensorView dst_indices) {
using namespace host;
auto N = SymbolicSize{"num_items"};
auto L = SymbolicSize{"num_layers"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({L}).with_dtype<uint64_t>().with_device(device).verify(src_ptrs).verify(dst_ptrs);
TensorMatcher({N}).with_dtype<int64_t>().with_device(device).verify(src_indices).verify(dst_indices);
const auto num_items = static_cast<uint32_t>(N.unwrap());
if (num_items == 0) {
return;
}
const auto num_layers = static_cast<uint32_t>(L.unwrap());
const int num_warps = BLOCK_SIZE / WARP_SIZE;
const int grid = (num_items + num_warps - 1) / num_warps;
LaunchKernel(grid, BLOCK_SIZE, device.unwrap())(
transfer_cache_dsv4_mla_kernel<BLOCK_SIZE>,
static_cast<void**>(src_ptrs.data_ptr()),
static_cast<void**>(dst_ptrs.data_ptr()),
static_cast<const int64_t*>(src_indices.data_ptr()),
static_cast<const int64_t*>(dst_indices.data_ptr()),
num_items,
num_layers);
}
__device__ __forceinline__ int warp_inclusive_scan(int* s_data, int lane_id, int offset, int count, int accumulator) { __device__ __forceinline__ int warp_inclusive_scan(int* s_data, int lane_id, int offset, int count, int accumulator) {
int idx = lane_id + offset; int idx = lane_id + offset;
int val = (idx < count) ? s_data[idx] : 0; int val = (idx < count) ? s_data[idx] : 0;
@@ -89,7 +145,7 @@ struct SmemLayout {
// //
// IsDsv4Layout selects the miss-copy addressing: // IsDsv4Layout selects the miss-copy addressing:
// false -> generic byte-stride: device + host both linear, stride = item_size_bytes // false -> generic byte-stride: device + host both linear, stride = item_size_bytes
// true -> DSv4 page-padded device + linear host (kvcacheio.cuh hardcoded constants) // true -> DSv4 page-padded device + page-padded host (kvcacheio.cuh constants)
template < template <
int BLOCK_SIZE, int BLOCK_SIZE,
int NUM_TOP_K, int NUM_TOP_K,
@@ -377,9 +433,10 @@ __global__ void load_cache_to_device_buffer_kernel(
const int64_t dst_loc = static_cast<int64_t>(req_device_buffer_locs[evict_slot]); const int64_t dst_loc = static_cast<int64_t>(req_device_buffer_locs[evict_slot]);
if constexpr (IsDsv4Layout) { if constexpr (IsDsv4Layout) {
// DSv4 path: page-padded device layout + linear host layout, K-only. // DSv4 path: page-padded device layout + page-padded host layout, K-only.
// Uses kvcacheio.cuh's hardcoded constants (kGPUPageSize=64, kCPUItemBytes=584). // The host cache is pinned DRAM but uses the same row layout as the GPU C4
device::hisparse::transfer_item<device::hisparse::TransferDirection::HostToDevice>( // cache, so use the page-padded address calculation for both ends.
device::hisparse::transfer_item(
/*dst_cache=*/device_buffer_k, /*dst_cache=*/device_buffer_k,
/*src_cache=*/const_cast<void*>(host_cache_k), /*src_cache=*/const_cast<void*>(host_cache_k),
/*dst_index=*/static_cast<int32_t>(dst_loc), /*dst_index=*/static_cast<int32_t>(dst_loc),
@@ -18,7 +18,6 @@ from .elementwise import (
fused_rope_inplace, fused_rope_inplace,
) )
from .gemm import linear_bf16_fp32 from .gemm import linear_bf16_fp32
from .hisparse import hisparse_offload_to_host
from .moe import ( from .moe import (
hash_topk, hash_topk,
mask_topk_ids, mask_topk_ids,
@@ -44,7 +43,6 @@ __all__ = [
"fused_k_norm_rope_flashmla", "fused_k_norm_rope_flashmla",
"make_name", "make_name",
"linear_bf16_fp32", "linear_bf16_fp32",
"hisparse_offload_to_host",
"get_paged_mqa_logits_metadata", "get_paged_mqa_logits_metadata",
"triton_create_paged_compress_data", "triton_create_paged_compress_data",
"topk_transform_512", "topk_transform_512",
-27
View File
@@ -1,27 +0,0 @@
import torch
from sglang.jit_kernel.utils import (
cache_once,
load_jit,
)
from .utils import make_name
@cache_once
def _jit_hisparse_transfer_module():
return load_jit(
make_name("hisparse_transfer"),
cuda_files=["deepseek_v4/hisparse_transfer.cuh"],
cuda_wrappers=[("hisparse_transfer", "hisparse_transfer")],
)
def hisparse_offload_to_host(
gpu_ptrs: torch.Tensor,
cpu_ptrs: torch.Tensor,
gpu_indices: torch.Tensor,
cpu_indices: torch.Tensor,
) -> None:
module = _jit_hisparse_transfer_module()
module.hisparse_transfer(gpu_ptrs, cpu_ptrs, gpu_indices, cpu_indices)
+34 -1
View File
@@ -39,6 +39,39 @@ def _jit_sparse_module(
) )
@functools.cache
def _jit_dsv4_transfer_module(block_size: int) -> Module:
template_args = make_cpp_args(block_size)
return load_jit(
"sparse_cache_dsv4_transfer",
block_size,
cuda_files=["hisparse.cuh"],
cuda_wrappers=[
(
"transfer_cache_dsv4_mla",
f"transfer_cache_dsv4_mla<{template_args}>",
)
],
)
def transfer_cache_dsv4_mla(
src_ptrs: torch.Tensor,
dst_ptrs: torch.Tensor,
src_indices: torch.Tensor,
dst_indices: torch.Tensor,
block_size: int = 1024,
) -> None:
"""Transfer DSv4 C4 tokens between page-padded C4 buffers."""
module = _jit_dsv4_transfer_module(block_size)
module.transfer_cache_dsv4_mla(
src_ptrs,
dst_ptrs,
src_indices,
dst_indices,
)
def _load_cache_to_device_buffer_mla( def _load_cache_to_device_buffer_mla(
*, *,
is_dsv4_layout: bool, is_dsv4_layout: bool,
@@ -156,7 +189,7 @@ def load_cache_to_device_buffer_dsv4_mla(
block_size: int = 256, block_size: int = 256,
num_real_reqs: torch.Tensor | None = None, num_real_reqs: torch.Tensor | None = None,
) -> None: ) -> None:
"""DSv4 hisparse swap-in: page-padded device + linear host (kvcacheio.cuh layout).""" """DSv4 hisparse swap-in: page-padded device + page-padded host C4 layout."""
_load_cache_to_device_buffer_mla( _load_cache_to_device_buffer_mla(
is_dsv4_layout=True, is_dsv4_layout=True,
top_k_tokens=top_k_tokens, top_k_tokens=top_k_tokens,
@@ -8,58 +8,38 @@
namespace device::hisparse { namespace device::hisparse {
/// NOTE: We call nope+rope as a "value" here. /// NOTE: We call nope+rope as a "value" here.
/// GPU Cache layout: /// Paged C4 cache layout:
/// VALUE 0, VALUE 1, ..., VALUE 63, /// VALUE 0, VALUE 1, ..., VALUE 63,
/// SCALE 0, SCALE 1, ..., SCALE 63, /// SCALE 0, SCALE 1, ..., SCALE 63,
/// [Padding to align to 576 bytes] /// [Padding to align to 576 bytes]
/// CPU Cache follow a trivial linear layout without any padding. inline constexpr int64_t kPageSize = 64;
inline constexpr int64_t kGPUPageSize = 64; inline constexpr int64_t kPageBits = 6; // log2(kPageSize)
inline constexpr int64_t kGPUPageBits = 6; // log2(kGPUPageSize)
inline constexpr int64_t kValueBytes = 576; inline constexpr int64_t kValueBytes = 576;
inline constexpr int64_t kScaleBytes = 8; inline constexpr int64_t kScaleBytes = 8;
/// NOTE: FlashMLA requires each page to be aligned to 576 bytes /// NOTE: FlashMLA requires each page to be aligned to 576 bytes
inline constexpr int64_t kCPUItemBytes = kValueBytes + kScaleBytes; inline constexpr int64_t kItemBytes = kValueBytes + kScaleBytes;
inline constexpr int64_t kGPUPageBytes = host::div_ceil(kCPUItemBytes * kGPUPageSize, 576) * 576; inline constexpr int64_t kPageBytes = host::div_ceil(kItemBytes * kPageSize, 576) * 576;
inline constexpr int64_t kGPUScaleOffset = kValueBytes * kGPUPageSize; inline constexpr int64_t kScaleOffset = kValueBytes * kPageSize;
struct PointerInfo { struct PointerInfo {
int64_t* value_ptr; int64_t* value_ptr;
int64_t* scale_ptr; int64_t* scale_ptr;
}; };
SGL_DEVICE PointerInfo get_pointer_gpu(void* cache, int32_t index) { SGL_DEVICE PointerInfo get_pointer_paged(void* cache, int32_t index) {
using namespace device; using namespace device;
static_assert(1 << kGPUPageBits == kGPUPageSize); static_assert(1 << kPageBits == kPageSize);
const int32_t page_num = index >> kGPUPageBits; const int32_t page_num = index >> kPageBits;
const int32_t page_offset = index & (kGPUPageSize - 1); const int32_t page_offset = index & (kPageSize - 1);
const auto page_ptr = pointer::offset(cache, page_num * kGPUPageBytes); const auto page_ptr = pointer::offset(cache, page_num * kPageBytes);
const auto value_ptr = pointer::offset(page_ptr, page_offset * kValueBytes); const auto value_ptr = pointer::offset(page_ptr, page_offset * kValueBytes);
const auto scale_ptr = pointer::offset(page_ptr, kGPUScaleOffset + page_offset * kScaleBytes); const auto scale_ptr = pointer::offset(page_ptr, kScaleOffset + page_offset * kScaleBytes);
return {static_cast<int64_t*>(value_ptr), static_cast<int64_t*>(scale_ptr)}; return {static_cast<int64_t*>(value_ptr), static_cast<int64_t*>(scale_ptr)};
} }
SGL_DEVICE PointerInfo get_pointer_cpu(void* cache, int32_t index) {
using namespace device;
const auto value_ptr = pointer::offset(cache, index * kCPUItemBytes);
const auto scale_ptr = pointer::offset(value_ptr, kValueBytes);
return {static_cast<int64_t*>(value_ptr), static_cast<int64_t*>(scale_ptr)};
}
enum class TransferDirection {
DeviceToDevice = 0,
DeviceToHost = 1,
HostToDevice = 2,
};
template <TransferDirection direction>
SGL_DEVICE void transfer_item(void* dst_cache, void* src_cache, const int32_t dst_index, const int32_t src_index) { SGL_DEVICE void transfer_item(void* dst_cache, void* src_cache, const int32_t dst_index, const int32_t src_index) {
constexpr bool is_dst_device = (direction != TransferDirection::DeviceToHost); const auto [dst_value_ptr, dst_scale_ptr] = get_pointer_paged(dst_cache, dst_index);
constexpr bool is_src_device = (direction != TransferDirection::HostToDevice); const auto [src_value_ptr, src_scale_ptr] = get_pointer_paged(src_cache, src_index);
constexpr auto dst_fn = is_dst_device ? get_pointer_gpu : get_pointer_cpu;
constexpr auto src_fn = is_src_device ? get_pointer_gpu : get_pointer_cpu;
const auto [dst_value_ptr, dst_scale_ptr] = dst_fn(dst_cache, dst_index);
const auto [src_value_ptr, src_scale_ptr] = src_fn(src_cache, src_index);
int64_t local_items[2]; int64_t local_items[2];
const int64_t* tail_src_ptr; const int64_t* tail_src_ptr;
+128 -1
View File
@@ -3,7 +3,11 @@ import sys
import pytest import pytest
import torch import torch
from sglang.jit_kernel.hisparse import load_cache_to_device_buffer_mla from sglang.jit_kernel.hisparse import (
load_cache_to_device_buffer_dsv4_mla,
load_cache_to_device_buffer_mla,
transfer_cache_dsv4_mla,
)
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.ci.ci_register import register_cuda_ci
@@ -26,6 +30,12 @@ PADDED_BUFFER_SIZE = HOT_BUFFER_SIZE + 1
HOST_CACHE_SIZE = 16 HOST_CACHE_SIZE = 16
DEVICE_CACHE_SIZE = 16 DEVICE_CACHE_SIZE = 16
ITEM_SIZE_BYTES = KV_DIM * torch.empty((), dtype=DTYPE).element_size() ITEM_SIZE_BYTES = KV_DIM * torch.empty((), dtype=DTYPE).element_size()
DSV4_PAGE_SIZE = 64
DSV4_VALUE_BYTES = 576
DSV4_SCALE_BYTES = 8
DSV4_ITEM_BYTES = DSV4_VALUE_BYTES + DSV4_SCALE_BYTES
DSV4_PAGE_BYTES = ((DSV4_ITEM_BYTES * DSV4_PAGE_SIZE + 575) // 576) * 576
DSV4_SCALE_OFFSET = DSV4_VALUE_BYTES * DSV4_PAGE_SIZE
def _host_cache() -> torch.Tensor: def _host_cache() -> torch.Tensor:
@@ -36,6 +46,46 @@ def _host_cache() -> torch.Tensor:
return host_cache return host_cache
def _dsv4_token_pattern(seed: int) -> tuple[torch.Tensor, torch.Tensor]:
value = (
(torch.arange(DSV4_VALUE_BYTES, dtype=torch.int16) + seed)
.remainder(256)
.to(torch.uint8)
)
scale = (
(torch.arange(DSV4_SCALE_BYTES, dtype=torch.int16) + seed + 17)
.remainder(256)
.to(torch.uint8)
)
return value, scale
def _write_dsv4_token(cache: torch.Tensor, loc: int, seed: int) -> None:
page = loc // DSV4_PAGE_SIZE
offset = loc % DSV4_PAGE_SIZE
value, scale = _dsv4_token_pattern(seed)
cache[page, offset * DSV4_VALUE_BYTES : (offset + 1) * DSV4_VALUE_BYTES].copy_(
value.to(cache.device)
)
scale_start = DSV4_SCALE_OFFSET + offset * DSV4_SCALE_BYTES
cache[page, scale_start : scale_start + DSV4_SCALE_BYTES].copy_(
scale.to(cache.device)
)
def _read_dsv4_token(cache: torch.Tensor, loc: int) -> torch.Tensor:
page = loc // DSV4_PAGE_SIZE
offset = loc % DSV4_PAGE_SIZE
value = cache[page, offset * DSV4_VALUE_BYTES : (offset + 1) * DSV4_VALUE_BYTES]
scale_start = DSV4_SCALE_OFFSET + offset * DSV4_SCALE_BYTES
scale = cache[page, scale_start : scale_start + DSV4_SCALE_BYTES]
return torch.cat([value, scale])
def _dsv4_ptrs(cache: torch.Tensor) -> torch.Tensor:
return torch.tensor([cache.data_ptr()], dtype=torch.uint64, device=DEVICE)
def _run_kernel( def _run_kernel(
*, *,
top_k_tokens: torch.Tensor, top_k_tokens: torch.Tensor,
@@ -132,6 +182,83 @@ def _make_state(
} }
@pytest.mark.skipif(is_hip(), reason="DSV4 paged-layout HiSparse test is CUDA-only.")
def test_transfer_cache_dsv4_mla_copies_paged_token() -> None:
src_cache = torch.zeros((2, DSV4_PAGE_BYTES), dtype=torch.uint8, device=DEVICE)
dst_cache = torch.zeros(
(2, DSV4_PAGE_BYTES), dtype=torch.uint8, device="cpu", pin_memory=True
)
src_loc = DSV4_PAGE_SIZE + 6
dst_loc = DSV4_PAGE_SIZE + 1
_write_dsv4_token(src_cache, src_loc, seed=41)
transfer_cache_dsv4_mla(
src_ptrs=_dsv4_ptrs(src_cache),
dst_ptrs=_dsv4_ptrs(dst_cache),
src_indices=torch.tensor([src_loc], dtype=torch.int64, device=DEVICE),
dst_indices=torch.tensor([dst_loc], dtype=torch.int64, device=DEVICE),
)
torch.cuda.synchronize()
assert torch.equal(
_read_dsv4_token(dst_cache, dst_loc).to(DEVICE),
_read_dsv4_token(src_cache, src_loc),
)
@pytest.mark.skipif(is_hip(), reason="DSV4 paged-layout HiSparse test is CUDA-only.")
def test_dsv4_swap_in_reads_paged_host_layout() -> None:
host_cache = torch.zeros(
(2, DSV4_PAGE_BYTES), dtype=torch.uint8, device="cpu", pin_memory=True
)
device_buffer = torch.zeros((2, DSV4_PAGE_BYTES), dtype=torch.uint8, device=DEVICE)
host_loc = DSV4_PAGE_SIZE + 1
swap_loc = DSV4_PAGE_SIZE + 12
_write_dsv4_token(host_cache, host_loc, seed=41)
top_k_tokens = torch.tensor([[3]], dtype=torch.int32, device=DEVICE)
device_buffer_tokens = torch.full(
(1, PADDED_BUFFER_SIZE), -1, dtype=torch.int32, device=DEVICE
)
host_cache_locs = torch.zeros((1, 8), dtype=torch.int64, device=DEVICE)
host_cache_locs[0, 3] = host_loc
device_buffer_locs = torch.tensor(
[[swap_loc, swap_loc + 1, swap_loc + 2, swap_loc + 3, swap_loc + 4]],
dtype=torch.int32,
device=DEVICE,
)
lru_slots = torch.arange(HOT_BUFFER_SIZE, dtype=torch.int16, device=DEVICE).view(
1, -1
)
out = torch.full_like(top_k_tokens, -1)
load_cache_to_device_buffer_dsv4_mla(
top_k_tokens=top_k_tokens,
device_buffer_tokens=device_buffer_tokens,
host_cache_locs=host_cache_locs,
device_buffer_locs=device_buffer_locs,
host_cache=host_cache,
device_buffer=device_buffer,
top_k_device_locs=out,
req_pool_indices=torch.tensor([0], dtype=torch.int64, device=DEVICE),
seq_lens=torch.tensor([8], dtype=torch.int32, device=DEVICE),
lru_slots=lru_slots,
item_size_bytes=DSV4_ITEM_BYTES,
num_top_k=1,
hot_buffer_size=HOT_BUFFER_SIZE,
page_size=1,
block_size=256,
num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE),
)
torch.cuda.synchronize()
assert out.item() == swap_loc
assert torch.equal(
_read_dsv4_token(device_buffer, swap_loc),
_read_dsv4_token(host_cache, host_loc).to(DEVICE),
)
def _long_case(): def _long_case():
# One-request baseline used by the stateful cases below: # One-request baseline used by the stateful cases below:
# req 0 LRU slots : [0, 1, 2, 3] # req 0 LRU slots : [0, 1, 2, 3]
+21 -2
View File
@@ -397,6 +397,16 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
kv_data_ptrs, kv_data_lens, kv_item_lens = ( kv_data_ptrs, kv_data_lens, kv_item_lens = (
transfer_kv_pool.get_contiguous_buf_infos() transfer_kv_pool.get_contiguous_buf_infos()
) )
if self.scheduler.enable_hisparse and isinstance(
self.token_to_kv_pool, DeepSeekV4TokenToKVPool
):
device_kv_data_ptrs, device_kv_data_lens, device_kv_item_lens = (
self.token_to_kv_pool.get_contiguous_buf_infos()
)
c4_layer_num = self.scheduler.hisparse_coordinator.mem_pool_host.layer_num
kv_data_ptrs += device_kv_data_ptrs[c4_layer_num:]
kv_data_lens += device_kv_data_lens[c4_layer_num:]
kv_item_lens += device_kv_item_lens[c4_layer_num:]
if self.draft_token_to_kv_pool is not None: if self.draft_token_to_kv_pool is not None:
# We should also transfer draft model kv cache. The indices are # We should also transfer draft model kv cache. The indices are
# always shared with a target model. # always shared with a target model.
@@ -932,7 +942,15 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
decode_req.req.cache_protected_len = total_prefix_len decode_req.req.cache_protected_len = total_prefix_len
page_size = self.token_to_kv_pool_allocator.page_size page_size = self.token_to_kv_pool_allocator.page_size
kv_transfer_page_size = page_size
if self.scheduler.enable_hisparse: if self.scheduler.enable_hisparse:
# Direct-to-host sends host/C4 rows; keep allocator.page_size
# logical and use the compressed page size only for these indices.
kv_transfer_page_size = getattr(
self.token_to_kv_pool_allocator,
"hisparse_page_size",
page_size,
)
# Must cast to int32 for ZMQ serialization -- from_zmq reads np.int32. # Must cast to int32 for ZMQ serialization -- from_zmq reads np.int32.
kv_indices = ( kv_indices = (
dst_kv_indices[: origin_input_len - prefix_len] dst_kv_indices[: origin_input_len - prefix_len]
@@ -1003,7 +1021,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
self.req_to_metadata_buffer_idx_allocator.alloc() self.req_to_metadata_buffer_idx_allocator.alloc()
) )
assert decode_req.metadata_buffer_index is not None assert decode_req.metadata_buffer_index is not None
page_indices = kv_to_page_indices(kv_indices, page_size) page_indices = kv_to_page_indices(kv_indices, kv_transfer_page_size)
decode_req.kv_receiver.send_metadata( decode_req.kv_receiver.send_metadata(
page_indices, page_indices,
decode_req.metadata_buffer_index, decode_req.metadata_buffer_index,
@@ -1299,13 +1317,14 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
last_loc=torch.tensor([-1], dtype=torch.int64, device=device), last_loc=torch.tensor([-1], dtype=torch.int64, device=device),
extend_num_tokens=fill_len, extend_num_tokens=fill_len,
) )
# Allocate host indices for the RDMA transfer target. # Allocate host indices for the RDMA transfer target.
host_indices = coordinator.mem_pool_host.alloc_paged_token_slots( host_indices = coordinator.mem_pool_host.alloc_paged_token_slots(
coordinator.req_to_host_pool, coordinator.req_to_host_pool,
coordinator.req_to_host_pool_allocated_len, coordinator.req_to_host_pool_allocated_len,
req.req_pool_idx, req.req_pool_idx,
0, 0,
fill_len, coordinator.host_token_len(fill_len),
) )
elif self.token_to_kv_pool_allocator.page_size == 1: elif self.token_to_kv_pool_allocator.page_size == 1:
kv_loc = self.token_to_kv_pool_allocator.alloc(delta_len) kv_loc = self.token_to_kv_pool_allocator.alloc(delta_len)
@@ -8,11 +8,13 @@ import torch
from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.schedule_batch import Req
from sglang.srt.mem_cache.hisparse_memory_pool import ( from sglang.srt.mem_cache.hisparse_memory_pool import (
DeepSeekV4HiSparseTokenToKVPoolAllocator, DeepSeekV4HiSparseTokenToKVPoolAllocator,
DeepSeekV4SingleKVPoolHost,
HiSparseDSATokenToKVPool, HiSparseDSATokenToKVPool,
HiSparseTokenToKVPoolAllocator, HiSparseTokenToKVPoolAllocator,
) )
from sglang.srt.mem_cache.memory_pool_host import MLATokenToKVPoolHost from sglang.srt.mem_cache.memory_pool_host import (
DeepSeekV4PagedHostPool,
MLATokenToKVPoolHost,
)
from sglang.srt.utils import get_device_module from sglang.srt.utils import get_device_module
device_module = get_device_module() device_module = get_device_module()
@@ -65,15 +67,23 @@ class HiSparseCoordinator:
) )
if self.is_dsv4_hisparse: if self.is_dsv4_hisparse:
self.mem_pool_device = self.token_to_kv_pool_allocator.hisparse_kvcache self.mem_pool_device = self.token_to_kv_pool_allocator.hisparse_kvcache
host_size = self.token_to_kv_pool_allocator.size_full // self.compress_ratio page_size = self.mem_pool_device.page_size
self.mem_pool_host = DeepSeekV4SingleKVPoolHost( num_host_pages = (
self.mem_pool_device, self.token_to_kv_pool_allocator.size_full // self.compress_ratio
host_size, + page_size
page_size=self.mem_pool_device.page_size, - 1
) // page_size
self.mem_pool_host = DeepSeekV4PagedHostPool(
pool_name="dsv4_hisparse_c4",
device_buffers=self.mem_pool_device.kv_buffer,
item_bytes=self.mem_pool_device.bytes_per_page_padded,
num_host_pages=num_host_pages,
slot_page_size=page_size,
layout="layer_first",
) )
self.item_size_bytes = ( self.item_size_bytes = (
self.mem_pool_host.kv_cache_total_dim self.mem_pool_device.kv_cache_total_dim
* self.mem_pool_host.dtype.itemsize * self.mem_pool_device.store_dtype.itemsize
) )
else: else:
assert isinstance( assert isinstance(
@@ -246,15 +256,10 @@ class HiSparseCoordinator:
buffer. In the staging path this is correct (prefill filled the buffer), buffer. In the staging path this is correct (prefill filled the buffer),
but here the buffer is empty. but here the buffer is empty.
""" """
if self.is_dsv4_hisparse:
# TODO(dsv4): wire PD direct-to-host. Needs (a) load_to_device_per_layer
raise NotImplementedError(
"PD direct-to-host admission is not supported for dsv4 hisparse yet."
)
self.alloc_device_buffer(req) self.alloc_device_buffer(req)
if req.kv_allocated_len <= self.device_buffer_size: host_len = self.host_token_len(req.kv_allocated_len)
if host_len <= self.device_buffer_size:
# Short sequences (seq_len <= device_buffer_size): the kernel fast path # Short sequences (seq_len <= device_buffer_size): the kernel fast path
# returns device_buffer_locs directly without any host loading, so we # returns device_buffer_locs directly without any host loading, so we
# must preload all tokens from host pool into the device buffer # must preload all tokens from host pool into the device buffer
@@ -271,9 +276,14 @@ class HiSparseCoordinator:
self._skip_first_backup[req.req_pool_idx] = True self._skip_first_backup[req.req_pool_idx] = True
logger.debug("HiSparse: admitting request %s directly", req.rid) logger.debug("HiSparse: admitting request %s directly", req.rid)
def host_token_len(self, kv_allocated_len: int) -> int:
if self.is_dsv4_hisparse:
return kv_allocated_len // self.compress_ratio
return kv_allocated_len
def _preload_to_device_buffer(self, req: Req) -> None: def _preload_to_device_buffer(self, req: Req) -> None:
"""Preload all tokens from host pool into the device buffer.""" """Preload all tokens from host pool into the device buffer."""
n = req.kv_allocated_len n = self.host_token_len(req.kv_allocated_len)
host_indices = self.req_to_host_pool[req.req_pool_idx, :n] host_indices = self.req_to_host_pool[req.req_pool_idx, :n]
device_locs = self.req_to_device_buffer[req.req_pool_idx, :n] device_locs = self.req_to_device_buffer[req.req_pool_idx, :n]
@@ -4,7 +4,6 @@ import logging
import weakref import weakref
from typing import Optional from typing import Optional
import psutil
import torch import torch
from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.radix_attention import RadixAttention
@@ -17,7 +16,6 @@ from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
HiSparseC4DevicePool, HiSparseC4DevicePool,
) )
from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
from sglang.srt.mem_cache.memory_pool_host import HiSparseHostPoolMixin
from sglang.srt.utils import is_cuda, is_hip from sglang.srt.utils import is_cuda, is_hip
from sglang.srt.utils.common import get_num_new_pages from sglang.srt.utils.common import get_num_new_pages
@@ -384,121 +382,6 @@ class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
) )
class DeepSeekV4SingleKVPoolHost(HiSparseHostPoolMixin):
def __init__(
self,
device_pool: HiSparseC4DevicePool,
host_size: int,
page_size: int,
pin_memory: bool = True,
device: str = "cpu",
):
assert host_size > 0, "Host size must be specified and greater than 0"
self.device_pool = device_pool
self.size = host_size
self.page_size = page_size
self.num_pages = (self.size + self.page_size - 1) // self.page_size
self.size = self.num_pages * self.page_size
self.pin_memory = pin_memory
self.device = device
self.dtype = device_pool.store_dtype
self.layer_num = device_pool.layer_num
self.kv_cache_total_dim = device_pool.kv_cache_total_dim
self.kv_buffer = self.init_kv_buffer()
self.data_refs = [self.kv_buffer[i] for i in range(self.layer_num)]
self.data_ptrs = torch.tensor(
[x.data_ptr() for x in self.data_refs],
dtype=torch.uint64,
device=self.device_pool.device,
)
self.clear()
def clear(self):
self.free_slots = torch.arange(
1, self.size + 1, dtype=torch.int64, device="cpu"
)
def init_kv_buffer(self):
dims = (self.layer_num, self.size + self.page_size, self.kv_cache_total_dim)
requested_bytes = (
self.layer_num
* (self.size + self.page_size)
* self.kv_cache_total_dim
* self.dtype.itemsize
)
host_mem = psutil.virtual_memory()
# preserve at least 10GB for other usage
ten_gb = 10 * (1024**3)
available_bytes = host_mem.available - ten_gb
if requested_bytes > available_bytes:
raise ValueError(
f"Not enough host memory available. Requesting "
f"{requested_bytes / 1e9:.2f} GB but only have "
f"{available_bytes / 1e9:.2f} GB free. Please reduce the "
f"size of the hierarchical cache."
)
else:
logger.info(
f"Allocating {requested_bytes / 1e9:.2f} GB host memory for hierarchical KV cache."
)
host_pool = torch.empty(dims, dtype=self.dtype, device=self.device)
assert self.pin_memory, "DeepSeekV4SingleKVPoolHost requires pin_memory=True"
if self.pin_memory:
torch.cuda.cudart().cudaHostRegister(
host_pool.data_ptr(), host_pool.numel() * host_pool.element_size(), 0
)
return host_pool
def backup_from_device_all_layer(
self, device_pool, host_indices, device_indices, io_backend="kernel"
):
if io_backend != "kernel":
raise ValueError(f"Unsupported IO backend: {io_backend}")
from sglang.jit_kernel.dsv4 import hisparse_offload_to_host
if host_indices.device != device_indices.device:
host_indices = host_indices.to(device=device_indices.device)
host_indices_i64 = (
host_indices.to(torch.int64)
if host_indices.dtype != torch.int64
else host_indices
)
device_indices_i64 = (
device_indices.to(torch.int64)
if device_indices.dtype != torch.int64
else device_indices
)
hisparse_offload_to_host(
gpu_ptrs=device_pool.data_ptrs,
cpu_ptrs=self.data_ptrs,
gpu_indices=device_indices_i64,
cpu_indices=host_indices_i64,
)
def available_size(self):
return len(self.free_slots)
def alloc(self, need_size: int) -> Optional[torch.Tensor]:
if need_size > self.available_size():
return None
select_index = self.free_slots[:need_size]
self.free_slots = self.free_slots[need_size:]
return select_index
def free(self, indices: torch.Tensor) -> int:
self.free_slots = torch.cat([self.free_slots, indices.cpu()])
return len(indices)
class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
def __init__( def __init__(
@@ -517,13 +400,16 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
self.dtype = self.hisparse_kvcache.dtype self.dtype = self.hisparse_kvcache.dtype
self.device = self.hisparse_kvcache.device self.device = self.hisparse_kvcache.device
self.page_size = self.hisparse_kvcache.page_size # Keep the public page_size as the logical DSV4 full/SWA page size.
# C4 HiSparse allocation/device-buffer code must use the compressed page size.
self.page_size = logical_attn_allocator.page_size
self.hisparse_page_size = self.hisparse_kvcache.page_size
self.logical_attn_allocator = logical_attn_allocator self.logical_attn_allocator = logical_attn_allocator
self._kvcache = logical_attn_allocator._kvcache self._kvcache = logical_attn_allocator._kvcache
self.hisparse_attn_allocator = PagedTokenToKVPoolAllocator( self.hisparse_attn_allocator = PagedTokenToKVPoolAllocator(
self._size_hisparse, self._size_hisparse,
self.page_size, self.hisparse_page_size,
self.dtype, self.dtype,
self.device, self.device,
self.hisparse_kvcache, self.hisparse_kvcache,
@@ -533,7 +419,7 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
self.full_to_hisparse_device_index_mapping = torch.cat( self.full_to_hisparse_device_index_mapping = torch.cat(
[ [
torch.zeros( torch.zeros(
self._kvcache.c4_logical_size + self.page_size, self._kvcache.c4_logical_size + self.hisparse_page_size,
dtype=torch.int64, dtype=torch.int64,
device=self.device, device=self.device,
), ),
@@ -606,12 +492,32 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
"use alloc_extend or alloc_decode instead." "use alloc_extend or alloc_decode instead."
) )
def alloc_logical_only(
self,
prefix_lens: torch.Tensor,
prefix_lens_cpu: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
last_loc: torch.Tensor,
extend_num_tokens: int,
):
"""Allocate decode logical indices without allocating C4 hisparse device pages."""
return self.logical_attn_allocator.alloc_extend(
prefix_lens,
prefix_lens_cpu,
seq_lens,
seq_lens_cpu,
last_loc,
extend_num_tokens,
)
def alloc_device_buffer(self, allocated_indices, need_size: int): def alloc_device_buffer(self, allocated_indices, need_size: int):
assert need_size % self.page_size == 0 assert need_size % self.hisparse_page_size == 0
hisparse_indices = self.full_to_hisparse_device_index_mapping[allocated_indices] hisparse_indices = self.full_to_hisparse_device_index_mapping[allocated_indices]
self.full_to_hisparse_device_index_mapping[allocated_indices] = 0 self.full_to_hisparse_device_index_mapping[allocated_indices] = 0
hisparse_indices = hisparse_indices[hisparse_indices > 0]
device_buffer_size = need_size - self.page_size device_buffer_size = need_size - self.hisparse_page_size
P = len(hisparse_indices) P = len(hisparse_indices)
if P > device_buffer_size + 1: if P > device_buffer_size + 1:
newest_src = hisparse_indices[P - 1].clone() newest_src = hisparse_indices[P - 1].clone()
@@ -623,14 +529,16 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
buffer_indices = hisparse_indices[:need_size] buffer_indices = hisparse_indices[:need_size]
surplus = hisparse_indices[need_size:] surplus = hisparse_indices[need_size:]
if surplus.numel() > 0: if surplus.numel() > 0:
buffer_pages = torch.unique(buffer_indices // self.page_size) buffer_pages = torch.unique(buffer_indices // self.hisparse_page_size)
surplus_pages = torch.unique(surplus // self.page_size) surplus_pages = torch.unique(surplus // self.hisparse_page_size)
pure_surplus = surplus_pages[~torch.isin(surplus_pages, buffer_pages)] pure_surplus = surplus_pages[~torch.isin(surplus_pages, buffer_pages)]
if pure_surplus.numel() > 0: if pure_surplus.numel() > 0:
self.hisparse_attn_allocator.is_not_in_free_group = True self.hisparse_attn_allocator.is_not_in_free_group = True
self.hisparse_attn_allocator.free(pure_surplus * self.page_size) self.hisparse_attn_allocator.free(
pure_surplus * self.hisparse_page_size
)
else: else:
page_residual_length = len(hisparse_indices) % self.page_size page_residual_length = len(hisparse_indices) % self.hisparse_page_size
if page_residual_length != 0: if page_residual_length != 0:
hisparse_indices = torch.cat( hisparse_indices = torch.cat(
[ [
@@ -638,7 +546,7 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
torch.arange( torch.arange(
hisparse_indices[-1] + 1, hisparse_indices[-1] + 1,
hisparse_indices[-1] hisparse_indices[-1]
+ self.page_size + self.hisparse_page_size
- page_residual_length - page_residual_length
+ 1, + 1,
device=self.device, device=self.device,
@@ -682,7 +590,7 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
) )
num_new_pages_hisparse = get_num_new_pages( num_new_pages_hisparse = get_num_new_pages(
seq_lens=seq_lens_cpu // self.compress_ratio, seq_lens=seq_lens_cpu // self.compress_ratio,
page_size=self.page_size, page_size=self.hisparse_page_size,
prefix_lens=prefix_lens_cpu // self.compress_ratio, prefix_lens=prefix_lens_cpu // self.compress_ratio,
) )
if ( if (
@@ -692,7 +600,7 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
return None return None
if ( if (
num_new_pages_hisparse num_new_pages_hisparse
> self.hisparse_attn_allocator.available_size() // self.page_size > self.hisparse_attn_allocator.available_size() // self.hisparse_page_size
): ):
return None return None
@@ -30,6 +30,7 @@ from sglang.jit_kernel.hicache import (
from sglang.jit_kernel.hicache import ( from sglang.jit_kernel.hicache import (
transfer_hicache_one_layer_mla as jit_transfer_hicache_one_layer_mla, transfer_hicache_one_layer_mla as jit_transfer_hicache_one_layer_mla,
) )
from sglang.jit_kernel.hisparse import transfer_cache_dsv4_mla
from sglang.srt.mem_cache.memory_pool import ( from sglang.srt.mem_cache.memory_pool import (
DSATokenToKVPool, DSATokenToKVPool,
KVCache, KVCache,
@@ -1874,7 +1875,7 @@ class LogicalHostPool:
return 0 return 0
class DeepSeekV4PagedHostPool(HostKVCache): class DeepSeekV4PagedHostPool(HiSparseHostPoolMixin, HostKVCache):
"""Host mirror for a DeepSeek V4 paged KV/indexer sub-pool.""" """Host mirror for a DeepSeek V4 paged KV/indexer sub-pool."""
def __init__( def __init__(
@@ -1979,14 +1980,28 @@ class DeepSeekV4PagedHostPool(HostKVCache):
) )
self.clear() self.clear()
def get_contiguous_buf_infos(self):
"""Return per-layer page-row buffers for PD direct-to-host transfer."""
data_ptrs = [int(self.data_ptrs[i].item()) for i in range(self.layer_num)]
data_lens = [self.kv_buffer[i].nbytes for i in range(self.layer_num)]
item_lens = [self.item_bytes * self.dtype.itemsize] * self.layer_num
return data_ptrs, data_lens, item_lens
def _to_page_indices(self, indices: torch.Tensor) -> torch.Tensor: def _to_page_indices(self, indices: torch.Tensor) -> torch.Tensor:
if indices.numel() % self.slot_page_size != 0:
raise ValueError(
f"{self.pool_name} transfer indices must be page-aligned, "
f"got numel={indices.numel()}, slot_page_size={self.slot_page_size}"
)
return indices.reshape(-1, self.slot_page_size)[:, 0] // self.slot_page_size return indices.reshape(-1, self.slot_page_size)[:, 0] // self.slot_page_size
def _has_transfer_indices(
self, host_indices: torch.Tensor | None, device_indices: torch.Tensor | None
) -> bool:
if host_indices is None or device_indices is None:
return False
if host_indices.numel() != device_indices.numel():
raise ValueError(
f"{self.pool_name} transfer index size mismatch: "
f"host={host_indices.numel()}, device={device_indices.numel()}"
)
return host_indices.numel() > 0
def get_size_per_token(self): def get_size_per_token(self):
return self.item_bytes return self.item_bytes
@@ -2026,12 +2041,26 @@ class DeepSeekV4PagedHostPool(HostKVCache):
def backup_from_device_all_layer( def backup_from_device_all_layer(
self, device_pool, host_indices, device_indices, io_backend self, device_pool, host_indices, device_indices, io_backend
): ):
if host_indices is None or device_indices is None: if not self._has_transfer_indices(host_indices, device_indices):
return
if (
host_indices.numel() % self.slot_page_size != 0
or device_indices.numel() % self.slot_page_size != 0
):
# Whole C4 pages can use the normal HiCache page-row copy below.
# Token-granular DSV4 C4 copy needs this helper because a token is
# not one contiguous byte range in the paged row:
# [value0..value63][scale0..scale63].
transfer_cache_dsv4_mla(
src_ptrs=self.device_ptrs,
dst_ptrs=self.data_ptrs,
src_indices=device_indices.to(dtype=torch.int64),
dst_indices=host_indices.to(dtype=torch.int64),
)
return return
host_rows = self._to_page_indices(host_indices) host_rows = self._to_page_indices(host_indices)
device_rows = self._to_page_indices(device_indices) device_rows = self._to_page_indices(device_indices)
if io_backend == "kernel" and self.layout == "layer_first": if io_backend == "kernel" and self.layout == "layer_first":
assert self.data_ptrs is not None
transfer_kv_all_layer_mla( transfer_kv_all_layer_mla(
src_layers=self.device_ptrs, src_layers=self.device_ptrs,
dst_layers=self.data_ptrs, dst_layers=self.data_ptrs,
@@ -2074,7 +2103,20 @@ class DeepSeekV4PagedHostPool(HostKVCache):
def load_to_device_per_layer( def load_to_device_per_layer(
self, device_pool, host_indices, device_indices, layer_id, io_backend self, device_pool, host_indices, device_indices, layer_id, io_backend
): ):
if host_indices is None or device_indices is None: if not self._has_transfer_indices(host_indices, device_indices):
return
if (
host_indices.numel() % self.slot_page_size != 0
or device_indices.numel() % self.slot_page_size != 0
):
# Same DSV4 C4 layout issue as backup: this is token-granular
# preload, so it cannot use the normal HiCache page-row copy.
transfer_cache_dsv4_mla(
src_ptrs=self.data_ptrs[layer_id : layer_id + 1],
dst_ptrs=self.device_ptrs[layer_id : layer_id + 1],
src_indices=host_indices.to(dtype=torch.int64),
dst_indices=device_indices.to(dtype=torch.int64),
)
return return
host_rows = self._to_page_indices(host_indices) host_rows = self._to_page_indices(host_indices)
device_rows = self._to_page_indices(device_indices) device_rows = self._to_page_indices(device_indices)
@@ -11,11 +11,15 @@ from sglang.test.test_utils import (
try_cached_model, try_cached_model,
) )
register_cuda_ci(est_time=250, stage="base-c", runner_config="deepep-8-gpu-h200") register_cuda_ci(est_time=500, stage="base-c", runner_config="deepep-8-gpu-h200")
DSV4_FLASH_MODEL = "sgl-project/DeepSeek-V4-Flash-FP8" DSV4_FLASH_MODEL = "sgl-project/DeepSeek-V4-Flash-FP8"
DEEPEP_CONFIG = '{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":96}}' DEEPEP_CONFIG = '{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":96}}'
DSV4_FLASH_LOADER_CONFIG = '{"enable_multithread_load": true, "num_threads": 64}'
DSV4_HISPARSE_CONFIG = (
'{"top_k":512,"device_buffer_size":4096,"host_to_device_ratio":2}'
)
DSV4_FLASH_ENV = { DSV4_FLASH_ENV = {
"SGLANG_DSV4_FP4_EXPERTS": "0", "SGLANG_DSV4_FP4_EXPERTS": "0",
@@ -123,5 +127,105 @@ class TestDisaggregationDSV4(PDDisaggregationServerBase, GSM8KMixin):
) )
class TestDisaggregationDSV4HiSparseMooncake(PDDisaggregationServerBase, GSM8KMixin):
gsm8k_accuracy_thres = 0.93
gsm8k_num_questions = 200
gsm8k_num_shots = 20
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.model = try_cached_model(DSV4_FLASH_MODEL)
cls.start_prefill()
cls.start_decode()
cls.wait_server_ready(cls.prefill_url + "/health", process=cls.process_prefill)
cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode)
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
4,
"--page-size",
256,
"--chunked-prefill-size",
8192,
"--max-running-requests",
16,
"--mem-fraction-static",
0.9,
"--skip-server-warmup",
"--reasoning-parser",
"deepseek-v4",
"--tool-call-parser",
"deepseekv4",
"--model-loader-extra-config",
DSV4_FLASH_LOADER_CONFIG,
"--watchdog-timeout",
"900",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
env=DSV4_FLASH_ENV,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
4,
"--base-gpu-id",
4,
"--page-size",
256,
"--chunked-prefill-size",
8192,
"--max-running-requests",
16,
"--mem-fraction-static",
0.9,
"--skip-server-warmup",
"--reasoning-parser",
"deepseek-v4",
"--tool-call-parser",
"deepseekv4",
"--model-loader-extra-config",
DSV4_FLASH_LOADER_CONFIG,
"--enable-hisparse",
"--hisparse-config",
DSV4_HISPARSE_CONFIG,
"--watchdog-timeout",
"900",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
env=DSV4_FLASH_ENV,
)
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()