From 00fefef16b028bf3e42483112871e0927722edea Mon Sep 17 00:00:00 2001 From: huangtingwei <141888744+huangtingwei9988@users.noreply.github.com> Date: Fri, 5 Jun 2026 15:39:48 +0800 Subject: [PATCH] [PD & HiSparse] Add DeepSeek V4 support for HiSparse direct Prefill-to-Decode DRAM (#24880) --- .../csrc/deepseek_v4/hisparse_transfer.cuh | 82 --------- python/sglang/jit_kernel/csrc/hisparse.cuh | 65 ++++++- python/sglang/jit_kernel/dsv4/__init__.py | 2 - python/sglang/jit_kernel/dsv4/hisparse.py | 27 --- python/sglang/jit_kernel/hisparse.py | 35 +++- .../sgl_kernel/deepseek_v4/kvcacheio.cuh | 48 ++--- .../sglang/jit_kernel/tests/test_hisparse.py | 129 +++++++++++++- python/sglang/srt/disaggregation/decode.py | 23 ++- .../srt/managers/hisparse_coordinator.py | 44 +++-- .../srt/mem_cache/hisparse_memory_pool.py | 166 ++++-------------- .../sglang/srt/mem_cache/memory_pool_host.py | 60 ++++++- .../test_disaggregation_dsv4.py | 106 ++++++++++- 12 files changed, 478 insertions(+), 309 deletions(-) delete mode 100644 python/sglang/jit_kernel/csrc/deepseek_v4/hisparse_transfer.cuh delete mode 100644 python/sglang/jit_kernel/dsv4/hisparse.py diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/hisparse_transfer.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/hisparse_transfer.cuh deleted file mode 100644 index aefec2437..000000000 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/hisparse_transfer.cuh +++ /dev/null @@ -1,82 +0,0 @@ -#include -#include - -#include - -#include - -#include -#include - -#include - -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( - /*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(); - TensorMatcher({L}) // 1D cache pointers - .with_dtype() - .with_device(device_) - .verify(gpu_ptrs) - .verify(cpu_ptrs); - TensorMatcher({N}) // 1D indices - .with_dtype() - .with_device(device_) - .verify(gpu_indices) - .verify(cpu_indices); - const auto params = OffloadParams{ - .gpu_caches = static_cast(gpu_ptrs.data_ptr()), - .cpu_caches = static_cast(cpu_ptrs.data_ptr()), - .gpu_indices = static_cast(gpu_indices.data_ptr()), - .cpu_indices = static_cast(cpu_indices.data_ptr()), - .num_items = static_cast(N.unwrap()), - .num_layers = static_cast(L.unwrap()), - }; - LaunchKernel(kBlockQuota, kBlockSize, device_.unwrap())(offload_to_cpu, params); -} - -} // namespace diff --git a/python/sglang/jit_kernel/csrc/hisparse.cuh b/python/sglang/jit_kernel/csrc/hisparse.cuh index 15da350a4..7a41cd8d8 100644 --- a/python/sglang/jit_kernel/csrc/hisparse.cuh +++ b/python/sglang/jit_kernel/csrc/hisparse.cuh @@ -52,6 +52,62 @@ transfer_item_warp(int32_t lane_id, const void* src_addr, void* dst_addr, int64_ } } +template +__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(src_indices[i]); + const int32_t dst_index = static_cast(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 +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(); + TensorMatcher({L}).with_dtype().with_device(device).verify(src_ptrs).verify(dst_ptrs); + TensorMatcher({N}).with_dtype().with_device(device).verify(src_indices).verify(dst_indices); + + const auto num_items = static_cast(N.unwrap()); + if (num_items == 0) { + return; + } + const auto num_layers = static_cast(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, + static_cast(src_ptrs.data_ptr()), + static_cast(dst_ptrs.data_ptr()), + static_cast(src_indices.data_ptr()), + static_cast(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) { int idx = lane_id + offset; int val = (idx < count) ? s_data[idx] : 0; @@ -89,7 +145,7 @@ struct SmemLayout { // // IsDsv4Layout selects the miss-copy addressing: // 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 < int BLOCK_SIZE, int NUM_TOP_K, @@ -377,9 +433,10 @@ __global__ void load_cache_to_device_buffer_kernel( const int64_t dst_loc = static_cast(req_device_buffer_locs[evict_slot]); if constexpr (IsDsv4Layout) { - // DSv4 path: page-padded device layout + linear host layout, K-only. - // Uses kvcacheio.cuh's hardcoded constants (kGPUPageSize=64, kCPUItemBytes=584). - device::hisparse::transfer_item( + // DSv4 path: page-padded device layout + page-padded host layout, K-only. + // The host cache is pinned DRAM but uses the same row layout as the GPU C4 + // cache, so use the page-padded address calculation for both ends. + device::hisparse::transfer_item( /*dst_cache=*/device_buffer_k, /*src_cache=*/const_cast(host_cache_k), /*dst_index=*/static_cast(dst_loc), diff --git a/python/sglang/jit_kernel/dsv4/__init__.py b/python/sglang/jit_kernel/dsv4/__init__.py index 3e7ea2b8c..e344599b0 100644 --- a/python/sglang/jit_kernel/dsv4/__init__.py +++ b/python/sglang/jit_kernel/dsv4/__init__.py @@ -18,7 +18,6 @@ from .elementwise import ( fused_rope_inplace, ) from .gemm import linear_bf16_fp32 -from .hisparse import hisparse_offload_to_host from .moe import ( hash_topk, mask_topk_ids, @@ -44,7 +43,6 @@ __all__ = [ "fused_k_norm_rope_flashmla", "make_name", "linear_bf16_fp32", - "hisparse_offload_to_host", "get_paged_mqa_logits_metadata", "triton_create_paged_compress_data", "topk_transform_512", diff --git a/python/sglang/jit_kernel/dsv4/hisparse.py b/python/sglang/jit_kernel/dsv4/hisparse.py deleted file mode 100644 index c1fdda2b0..000000000 --- a/python/sglang/jit_kernel/dsv4/hisparse.py +++ /dev/null @@ -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) diff --git a/python/sglang/jit_kernel/hisparse.py b/python/sglang/jit_kernel/hisparse.py index a6930ecb7..b88143da1 100644 --- a/python/sglang/jit_kernel/hisparse.py +++ b/python/sglang/jit_kernel/hisparse.py @@ -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( *, is_dsv4_layout: bool, @@ -156,7 +189,7 @@ def load_cache_to_device_buffer_dsv4_mla( block_size: int = 256, num_real_reqs: torch.Tensor | 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( is_dsv4_layout=True, top_k_tokens=top_k_tokens, diff --git a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/kvcacheio.cuh b/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/kvcacheio.cuh index 0a3acc477..ed3097b40 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/kvcacheio.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/kvcacheio.cuh @@ -8,58 +8,38 @@ namespace device::hisparse { /// NOTE: We call nope+rope as a "value" here. -/// GPU Cache layout: +/// Paged C4 cache layout: /// VALUE 0, VALUE 1, ..., VALUE 63, /// SCALE 0, SCALE 1, ..., SCALE 63, /// [Padding to align to 576 bytes] -/// CPU Cache follow a trivial linear layout without any padding. -inline constexpr int64_t kGPUPageSize = 64; -inline constexpr int64_t kGPUPageBits = 6; // log2(kGPUPageSize) +inline constexpr int64_t kPageSize = 64; +inline constexpr int64_t kPageBits = 6; // log2(kPageSize) inline constexpr int64_t kValueBytes = 576; inline constexpr int64_t kScaleBytes = 8; /// NOTE: FlashMLA requires each page to be aligned to 576 bytes -inline constexpr int64_t kCPUItemBytes = kValueBytes + kScaleBytes; -inline constexpr int64_t kGPUPageBytes = host::div_ceil(kCPUItemBytes * kGPUPageSize, 576) * 576; -inline constexpr int64_t kGPUScaleOffset = kValueBytes * kGPUPageSize; +inline constexpr int64_t kItemBytes = kValueBytes + kScaleBytes; +inline constexpr int64_t kPageBytes = host::div_ceil(kItemBytes * kPageSize, 576) * 576; +inline constexpr int64_t kScaleOffset = kValueBytes * kPageSize; struct PointerInfo { int64_t* value_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; - static_assert(1 << kGPUPageBits == kGPUPageSize); - const int32_t page_num = index >> kGPUPageBits; - const int32_t page_offset = index & (kGPUPageSize - 1); - const auto page_ptr = pointer::offset(cache, page_num * kGPUPageBytes); + static_assert(1 << kPageBits == kPageSize); + const int32_t page_num = index >> kPageBits; + const int32_t page_offset = index & (kPageSize - 1); + const auto page_ptr = pointer::offset(cache, page_num * kPageBytes); 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(value_ptr), static_cast(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(value_ptr), static_cast(scale_ptr)}; -} - -enum class TransferDirection { - DeviceToDevice = 0, - DeviceToHost = 1, - HostToDevice = 2, -}; - -template 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); - constexpr bool is_src_device = (direction != TransferDirection::HostToDevice); - 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); + const auto [dst_value_ptr, dst_scale_ptr] = get_pointer_paged(dst_cache, dst_index); + const auto [src_value_ptr, src_scale_ptr] = get_pointer_paged(src_cache, src_index); int64_t local_items[2]; const int64_t* tail_src_ptr; diff --git a/python/sglang/jit_kernel/tests/test_hisparse.py b/python/sglang/jit_kernel/tests/test_hisparse.py index c4b729f5e..ec4c779d5 100644 --- a/python/sglang/jit_kernel/tests/test_hisparse.py +++ b/python/sglang/jit_kernel/tests/test_hisparse.py @@ -3,7 +3,11 @@ import sys import pytest 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.test.ci.ci_register import register_cuda_ci @@ -26,6 +30,12 @@ PADDED_BUFFER_SIZE = HOT_BUFFER_SIZE + 1 HOST_CACHE_SIZE = 16 DEVICE_CACHE_SIZE = 16 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: @@ -36,6 +46,46 @@ def _host_cache() -> torch.Tensor: 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( *, 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(): # One-request baseline used by the stateful cases below: # req 0 LRU slots : [0, 1, 2, 3] diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 022446bf4..41f92d51c 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -397,6 +397,16 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): kv_data_ptrs, kv_data_lens, kv_item_lens = ( 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: # We should also transfer draft model kv cache. The indices are # always shared with a target model. @@ -932,7 +942,15 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): decode_req.req.cache_protected_len = total_prefix_len page_size = self.token_to_kv_pool_allocator.page_size + kv_transfer_page_size = page_size 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. kv_indices = ( dst_kv_indices[: origin_input_len - prefix_len] @@ -1003,7 +1021,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): self.req_to_metadata_buffer_idx_allocator.alloc() ) 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( page_indices, decode_req.metadata_buffer_index, @@ -1299,13 +1317,14 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): last_loc=torch.tensor([-1], dtype=torch.int64, device=device), extend_num_tokens=fill_len, ) + # Allocate host indices for the RDMA transfer target. host_indices = coordinator.mem_pool_host.alloc_paged_token_slots( coordinator.req_to_host_pool, coordinator.req_to_host_pool_allocated_len, req.req_pool_idx, 0, - fill_len, + coordinator.host_token_len(fill_len), ) elif self.token_to_kv_pool_allocator.page_size == 1: kv_loc = self.token_to_kv_pool_allocator.alloc(delta_len) diff --git a/python/sglang/srt/managers/hisparse_coordinator.py b/python/sglang/srt/managers/hisparse_coordinator.py index e03005546..1720affb1 100644 --- a/python/sglang/srt/managers/hisparse_coordinator.py +++ b/python/sglang/srt/managers/hisparse_coordinator.py @@ -8,11 +8,13 @@ import torch from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.hisparse_memory_pool import ( DeepSeekV4HiSparseTokenToKVPoolAllocator, - DeepSeekV4SingleKVPoolHost, HiSparseDSATokenToKVPool, 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 device_module = get_device_module() @@ -65,15 +67,23 @@ class HiSparseCoordinator: ) if self.is_dsv4_hisparse: 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 - self.mem_pool_host = DeepSeekV4SingleKVPoolHost( - self.mem_pool_device, - host_size, - page_size=self.mem_pool_device.page_size, + page_size = self.mem_pool_device.page_size + num_host_pages = ( + self.token_to_kv_pool_allocator.size_full // self.compress_ratio + + 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.mem_pool_host.kv_cache_total_dim - * self.mem_pool_host.dtype.itemsize + self.mem_pool_device.kv_cache_total_dim + * self.mem_pool_device.store_dtype.itemsize ) else: assert isinstance( @@ -246,15 +256,10 @@ class HiSparseCoordinator: buffer. In the staging path this is correct (prefill filled the buffer), 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) - 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 # returns device_buffer_locs directly without any host loading, so we # 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 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: """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] device_locs = self.req_to_device_buffer[req.req_pool_idx, :n] diff --git a/python/sglang/srt/mem_cache/hisparse_memory_pool.py b/python/sglang/srt/mem_cache/hisparse_memory_pool.py index d80d73697..c0dcd02b8 100644 --- a/python/sglang/srt/mem_cache/hisparse_memory_pool.py +++ b/python/sglang/srt/mem_cache/hisparse_memory_pool.py @@ -4,7 +4,6 @@ import logging import weakref from typing import Optional -import psutil import torch from sglang.srt.layers.radix_attention import RadixAttention @@ -17,7 +16,6 @@ from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( HiSparseC4DevicePool, ) 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.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): def __init__( @@ -517,13 +400,16 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.dtype = self.hisparse_kvcache.dtype 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._kvcache = logical_attn_allocator._kvcache self.hisparse_attn_allocator = PagedTokenToKVPoolAllocator( self._size_hisparse, - self.page_size, + self.hisparse_page_size, self.dtype, self.device, self.hisparse_kvcache, @@ -533,7 +419,7 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.full_to_hisparse_device_index_mapping = torch.cat( [ torch.zeros( - self._kvcache.c4_logical_size + self.page_size, + self._kvcache.c4_logical_size + self.hisparse_page_size, dtype=torch.int64, device=self.device, ), @@ -606,12 +492,32 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): "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): - 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] 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) if P > device_buffer_size + 1: newest_src = hisparse_indices[P - 1].clone() @@ -623,14 +529,16 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): buffer_indices = hisparse_indices[:need_size] surplus = hisparse_indices[need_size:] if surplus.numel() > 0: - buffer_pages = torch.unique(buffer_indices // self.page_size) - surplus_pages = torch.unique(surplus // self.page_size) + buffer_pages = torch.unique(buffer_indices // self.hisparse_page_size) + surplus_pages = torch.unique(surplus // self.hisparse_page_size) pure_surplus = surplus_pages[~torch.isin(surplus_pages, buffer_pages)] if pure_surplus.numel() > 0: 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: - page_residual_length = len(hisparse_indices) % self.page_size + page_residual_length = len(hisparse_indices) % self.hisparse_page_size if page_residual_length != 0: hisparse_indices = torch.cat( [ @@ -638,7 +546,7 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): torch.arange( hisparse_indices[-1] + 1, hisparse_indices[-1] - + self.page_size + + self.hisparse_page_size - page_residual_length + 1, device=self.device, @@ -682,7 +590,7 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): ) num_new_pages_hisparse = get_num_new_pages( 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, ) if ( @@ -692,7 +600,7 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): return None if ( 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 diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 69e65f4e4..92e5e47b8 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -30,6 +30,7 @@ from sglang.jit_kernel.hicache import ( from sglang.jit_kernel.hicache import ( 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 ( DSATokenToKVPool, KVCache, @@ -1874,7 +1875,7 @@ class LogicalHostPool: return 0 -class DeepSeekV4PagedHostPool(HostKVCache): +class DeepSeekV4PagedHostPool(HiSparseHostPoolMixin, HostKVCache): """Host mirror for a DeepSeek V4 paged KV/indexer sub-pool.""" def __init__( @@ -1979,14 +1980,28 @@ class DeepSeekV4PagedHostPool(HostKVCache): ) 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: - 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 + 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): return self.item_bytes @@ -2026,12 +2041,26 @@ class DeepSeekV4PagedHostPool(HostKVCache): def backup_from_device_all_layer( 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 host_rows = self._to_page_indices(host_indices) device_rows = self._to_page_indices(device_indices) if io_backend == "kernel" and self.layout == "layer_first": - assert self.data_ptrs is not None transfer_kv_all_layer_mla( src_layers=self.device_ptrs, dst_layers=self.data_ptrs, @@ -2074,7 +2103,20 @@ class DeepSeekV4PagedHostPool(HostKVCache): def load_to_device_per_layer( 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 host_rows = self._to_page_indices(host_indices) device_rows = self._to_page_indices(device_indices) diff --git a/test/registered/disaggregation/test_disaggregation_dsv4.py b/test/registered/disaggregation/test_disaggregation_dsv4.py index 362ba3581..91231cead 100644 --- a/test/registered/disaggregation/test_disaggregation_dsv4.py +++ b/test/registered/disaggregation/test_disaggregation_dsv4.py @@ -11,11 +11,15 @@ from sglang.test.test_utils import ( 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" 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 = { "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__": unittest.main()