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

python/sglang/jit_kernel/csrc/deepseek_v4/hisparse_transfer.cuh

@@ -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

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 <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) {
   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<int64_t>(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<device::hisparse::TransferDirection::HostToDevice>(
+      // 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<void*>(host_cache_k),
           /*dst_index=*/static_cast<int32_t>(dst_loc),

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",

python/sglang/jit_kernel/dsv4/hisparse.py

@@ -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)

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,

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<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) {
-  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;

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]

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)

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]
 

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
 

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)