[PD & HiSparse] Add DeepSeek V4 support for HiSparse direct Prefill-to-Decode DRAM (#24880)
This commit is contained in:
@@ -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
|
|
||||||
@@ -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",
|
||||||
|
|||||||
@@ -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)
|
|
||||||
@@ -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;
|
||||||
|
|||||||
@@ -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]
|
||||||
|
|||||||
@@ -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()
|
||||||
|
|||||||
Reference in New Issue
Block a user