[HiCache & JIT Kernel] Refactoring HiCache Write-Back Kernel (#21631)

python/sglang/jit_kernel/csrc/kvcacheio/hicache.cuh

@@ -1,3 +1,5 @@
+#pragma once
+
 #include <sgl_kernel/tensor.h>
 #include <sgl_kernel/utils.h>
 

python/sglang/jit_kernel/csrc/kvcacheio/relayout.cuh

@@ -0,0 +1,90 @@
+#pragma once
+
+#include "hicache.cuh"
+#include <limits>
+
+namespace {
+
+struct HicacheRelayoutParams {
+  void* __restrict__ k_cache_dst;
+  void* __restrict__ v_cache_dst;
+  const void* __restrict__ indices_src;
+  const void* __restrict__ k_ptr_src;
+  const void* __restrict__ v_ptr_src;
+  uint32_t num_pages;
+  uint32_t num_layers;
+  uint32_t page_size;
+};
+
+template <typename IndexType, int64_t kElementSize, bool kIsMLA>
+__global__ void hicache_relayout_kernel(const __grid_constant__ HicacheRelayoutParams params) {
+  using namespace device;
+  using pack_t = uint4;
+  static_assert(kElementSize % 16 == 0, "hicache_relayout_kernel requires 16-byte aligned element size");
+  constexpr uint32_t kVecBytes = 16;
+  constexpr uint32_t kVecPerItem = kElementSize / kVecBytes;
+
+  const auto& [k_cache_dst, v_cache_dst, indices_src, k_ptr_src, v_ptr_src, num_pages, num_layers, page_size] = params;
+  const auto k_ptr_src_arr = static_cast<const void* const*>(k_ptr_src);
+  const auto v_ptr_src_arr = static_cast<const void* const*>(v_ptr_src);
+  const auto tid = static_cast<uint64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  const auto stride = static_cast<uint64_t>(gridDim.x) * blockDim.x;
+  const auto total_vecs = static_cast<uint64_t>(num_pages) * page_size * num_layers * kVecPerItem;
+
+  for (uint64_t linear_vec_id = tid; linear_vec_id < total_vecs; linear_vec_id += stride) {
+    const auto page_id =
+        static_cast<uint32_t>(linear_vec_id / (static_cast<uint64_t>(page_size) * num_layers * kVecPerItem));
+    const auto page_vec_id =
+        static_cast<uint32_t>(linear_vec_id % (static_cast<uint64_t>(page_size) * num_layers * kVecPerItem));
+    const auto token_in_page = page_vec_id / (num_layers * kVecPerItem);
+    const auto token_vec_id = page_vec_id % (num_layers * kVecPerItem);
+    const auto layer_id = token_vec_id / kVecPerItem;
+    const auto vec_id = token_vec_id % kVecPerItem;
+    const auto src_page = static_cast<uint32_t>(static_cast<const IndexType*>(indices_src)[page_id]);
+    const auto src_token = src_page + token_in_page;
+    const auto src_k = pointer::offset(
+        static_cast<const void*>(k_ptr_src_arr[layer_id]),
+        static_cast<int64_t>(src_token) * kElementSize + static_cast<int64_t>(vec_id) * kVecBytes);
+    const auto dst_k =
+        pointer::offset(static_cast<void*>(k_cache_dst), static_cast<int64_t>(linear_vec_id) * kVecBytes);
+    const auto vec_k = details::load_nc(reinterpret_cast<const pack_t*>(src_k));
+    details::store_nc(reinterpret_cast<pack_t*>(dst_k), vec_k);
+
+    if constexpr (!kIsMLA) {
+      const auto src_v = pointer::offset(
+          static_cast<const void*>(v_ptr_src_arr[layer_id]),
+          static_cast<int64_t>(src_token) * kElementSize + static_cast<int64_t>(vec_id) * kVecBytes);
+      const auto dst_v =
+          pointer::offset(static_cast<void*>(v_cache_dst), static_cast<int64_t>(linear_vec_id) * kVecBytes);
+      const auto vec_v = details::load_nc(reinterpret_cast<const pack_t*>(src_v));
+      details::store_nc(reinterpret_cast<pack_t*>(dst_v), vec_v);
+    }
+  }
+}
+
+template <int64_t kElementSize, bool kIsMLA>
+inline void launch_hicache_relayout_kernel(
+    const HicacheRelayoutParams& params,
+    int64_t num_pages,
+    int64_t num_layers,
+    int64_t page_size,
+    bool use_int32,
+    DLDevice device) {
+  using namespace host;
+
+  constexpr uint32_t kRelayoutBlockSize = 256;
+  constexpr uint32_t kVecPerItem = kElementSize / 16;
+  const auto total_vecs = static_cast<uint64_t>(num_pages) * page_size * num_layers * kVecPerItem;
+  const auto kernel = use_int32 ? hicache_relayout_kernel<int32_t, kElementSize, kIsMLA>
+                                : hicache_relayout_kernel<int64_t, kElementSize, kIsMLA>;
+  if (total_vecs == 0) {
+    return;
+  }
+
+  const auto grid = div_ceil(total_vecs, static_cast<uint64_t>(kRelayoutBlockSize));
+  RuntimeCheck(
+      grid <= std::numeric_limits<uint32_t>::max(), "HiCache staged relayout: CUDA grid size exceeds uint32 range");
+  LaunchKernel(static_cast<uint32_t>(grid), kRelayoutBlockSize, device)(kernel, params);
+}
+
+}  // namespace

python/sglang/jit_kernel/csrc/kvcacheio/staged_write_back.cuh

@@ -0,0 +1,328 @@
+#pragma once
+
+#include "hicache.cuh"
+#include "relayout.cuh"
+#include <dlfcn.h>
+#include <limits>
+#include <vector>
+
+namespace {
+
+#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12080
+#if CUDA_VERSION >= 13000
+using CudaMemcpyBatchPtr = const void*;
+using CudaMemcpyBatchAsyncFn = cudaError_t (*)(
+    CudaMemcpyBatchPtr*,
+    CudaMemcpyBatchPtr*,
+    const size_t*,
+    size_t,
+    cudaMemcpyAttributes*,
+    size_t*,
+    size_t,
+    cudaStream_t);
+#else
+using CudaMemcpyBatchPtr = void*;
+using CudaMemcpyBatchAsyncFn = cudaError_t (*)(
+    CudaMemcpyBatchPtr*,
+    CudaMemcpyBatchPtr*,
+    size_t*,
+    size_t,
+    cudaMemcpyAttributes*,
+    size_t*,
+    size_t,
+    size_t*,
+    cudaStream_t);
+#endif
+
+inline auto get_cuda_memcpy_batch_async() -> CudaMemcpyBatchAsyncFn {
+  static CudaMemcpyBatchAsyncFn cuda_memcpy_batch_async = []() {
+    void* symbol = dlsym(RTLD_DEFAULT, "cudaMemcpyBatchAsync");
+    return reinterpret_cast<CudaMemcpyBatchAsyncFn>(symbol);
+  }();
+  return cuda_memcpy_batch_async;
+}
+
+inline auto call_cuda_memcpy_batch_async(
+    CudaMemcpyBatchAsyncFn copy_fn,
+    CudaMemcpyBatchPtr* dsts,
+    CudaMemcpyBatchPtr* srcs,
+    size_t* sizes,
+    size_t count,
+    cudaMemcpyAttributes* attrs,
+    size_t* attrs_idxs,
+    size_t num_attrs,
+    cudaStream_t stream) -> cudaError_t {
+#if CUDA_VERSION >= 13000
+  return copy_fn(dsts, srcs, sizes, count, attrs, attrs_idxs, num_attrs, stream);
+#else
+  size_t fail_idx = std::numeric_limits<size_t>::max();
+  return copy_fn(dsts, srcs, sizes, count, attrs, attrs_idxs, num_attrs, &fail_idx, stream);
+#endif
+}
+#endif
+
+inline void copy_page_first_pages_fallback(
+    const std::vector<tvm::ffi::TensorView>& src_ptrs,
+    std::vector<tvm::ffi::TensorView> dst_ptrs,
+    const int64_t* dst_indices_ptr,
+    int64_t num_pages,
+    int64_t page_size,
+    cudaStream_t stream) {
+  using namespace host;
+
+  RuntimeCheck(src_ptrs.size() == dst_ptrs.size(), "Source and destination tensors must have the same count");
+  for (const auto tensor_id : irange(src_ptrs.size())) {
+    RuntimeCheck(
+        src_ptrs[tensor_id].dtype() == dst_ptrs[tensor_id].dtype(),
+        "Source and destination tensors must have the same dtype");
+    const int64_t elem_size = host::dtype_bytes(src_ptrs[tensor_id].dtype());
+    const int64_t src_stride0 = src_ptrs[tensor_id].stride(0);
+    const int64_t dst_stride0 = dst_ptrs[tensor_id].stride(0);
+    const size_t src_page_bytes = static_cast<size_t>(page_size * src_stride0 * elem_size);
+    const size_t dst_page_bytes = static_cast<size_t>(page_size * dst_stride0 * elem_size);
+    RuntimeCheck(src_page_bytes == dst_page_bytes, "Source and destination page spans must match");
+    for (const auto page_offset : irange(num_pages)) {
+      const char* src_ptr = static_cast<const char*>(src_ptrs[tensor_id].data_ptr()) +
+                            static_cast<size_t>(page_offset * page_size * src_stride0 * elem_size);
+      char* dst_ptr = static_cast<char*>(dst_ptrs[tensor_id].data_ptr()) +
+                      static_cast<size_t>(dst_indices_ptr[page_offset * page_size] * dst_stride0 * elem_size);
+      RuntimeDeviceCheck(cudaMemcpyAsync(dst_ptr, src_ptr, src_page_bytes, cudaMemcpyDeviceToHost, stream));
+    }
+  }
+}
+
+inline bool try_copy_page_first_pages_batch(
+    const std::vector<tvm::ffi::TensorView>& src_ptrs,
+    std::vector<tvm::ffi::TensorView> dst_ptrs,
+    const int64_t* dst_indices_ptr,
+    int64_t num_pages,
+    int64_t page_size,
+    int device_id,
+    cudaStream_t stream) {
+#if defined(USE_ROCM) || !defined(CUDA_VERSION) || (CUDA_VERSION < 12080)
+  return false;
+#else
+  host::RuntimeCheck(src_ptrs.size() == dst_ptrs.size(), "Source and destination tensors must have the same count");
+  constexpr size_t kLargeCopyThresholdBytes = 128 * 1024;
+  thread_local std::vector<CudaMemcpyBatchPtr> batch_srcs;
+  thread_local std::vector<CudaMemcpyBatchPtr> batch_dsts;
+  thread_local std::vector<size_t> batch_sizes;
+
+  int driver_version = 0;
+  cudaError_t driver_version_err = cudaDriverGetVersion(&driver_version);
+  if (driver_version_err != cudaSuccess || driver_version < 12080) {
+    return false;
+  }
+
+  auto copy_fn = get_cuda_memcpy_batch_async();
+  if (copy_fn == nullptr) {
+    return false;
+  }
+
+  const size_t num_copies = static_cast<size_t>(src_ptrs.size()) * static_cast<size_t>(num_pages);
+  batch_srcs.clear();
+  batch_dsts.clear();
+  batch_sizes.clear();
+  batch_srcs.reserve(num_copies);
+  batch_dsts.reserve(num_copies);
+  batch_sizes.reserve(num_copies);
+
+  size_t first_page_bytes = 0;
+  for (const auto tensor_id : host::irange(src_ptrs.size())) {
+    host::RuntimeCheck(
+        src_ptrs[tensor_id].dtype() == dst_ptrs[tensor_id].dtype(),
+        "Source and destination tensors must have the same dtype");
+    const int64_t elem_size = host::dtype_bytes(src_ptrs[tensor_id].dtype());
+    const int64_t src_stride0 = src_ptrs[tensor_id].stride(0);
+    const int64_t dst_stride0 = dst_ptrs[tensor_id].stride(0);
+    const size_t src_page_bytes = static_cast<size_t>(page_size * src_stride0 * elem_size);
+    const size_t dst_page_bytes = static_cast<size_t>(page_size * dst_stride0 * elem_size);
+    host::RuntimeCheck(src_page_bytes == dst_page_bytes, "Source and destination page spans must match");
+    if (tensor_id == 0) {
+      first_page_bytes = src_page_bytes;
+    }
+    for (const auto page_offset : host::irange(num_pages)) {
+      char* src_ptr = static_cast<char*>(src_ptrs[tensor_id].data_ptr()) +
+                      static_cast<size_t>(page_offset * page_size * src_stride0 * elem_size);
+      char* dst_ptr = static_cast<char*>(dst_ptrs[tensor_id].data_ptr()) +
+                      static_cast<size_t>(dst_indices_ptr[page_offset * page_size] * dst_stride0 * elem_size);
+      batch_srcs.push_back(src_ptr);
+      batch_dsts.push_back(dst_ptr);
+      batch_sizes.push_back(src_page_bytes);
+    }
+  }
+  if (first_page_bytes < kLargeCopyThresholdBytes) {
+    return false;
+  }
+
+  std::vector<size_t> attrs_idxs(1, 0);
+  cudaMemcpyAttributes attrs{};
+  attrs.srcAccessOrder = cudaMemcpySrcAccessOrderStream;
+  attrs.srcLocHint.type = cudaMemLocationTypeDevice;
+  attrs.srcLocHint.id = device_id;
+  attrs.dstLocHint.type = cudaMemLocationTypeHost;
+  attrs.dstLocHint.id = 0;
+  attrs.flags = 0;
+
+  cudaError_t err = call_cuda_memcpy_batch_async(
+      copy_fn,
+      batch_dsts.data(),
+      batch_srcs.data(),
+      batch_sizes.data(),
+      num_copies,
+      &attrs,
+      attrs_idxs.data(),
+      1,
+      stream);
+  if (err == cudaErrorNotSupported || err == cudaErrorCallRequiresNewerDriver || err == cudaErrorInvalidValue) {
+    (void)cudaGetLastError();
+    return false;
+  }
+  host::RuntimeCheck(err == cudaSuccess, "cudaMemcpyBatchAsync failed. error=", cudaGetErrorString(err));
+  return true;
+#endif
+}
+
+template <int64_t kElementSize, uint32_t kUnroll, uint32_t kBlockQuota, uint32_t kBlockSize>
+struct HiCacheStagedWriteBackKernel {
+ private:
+  template <bool kIsMLA>
+  static void run_staged_impl(
+      const tvm::ffi::TensorView k_cache_dst,
+      const tvm::ffi::TensorView v_cache_dst,
+      const tvm::ffi::TensorView dst_indices_cpu,
+      const tvm::ffi::TensorView staging_k,
+      const tvm::ffi::TensorView staging_v,
+      const tvm::ffi::TensorView page_indices_src,
+      const tvm::ffi::TensorView k_ptr_src,
+      const tvm::ffi::TensorView v_ptr_src,
+      const int64_t page_size) {
+    using namespace host;
+
+    auto T = SymbolicSize{"num_tokens"};
+    auto N = SymbolicSize{"num_layers"};
+    auto D = SymbolicSize{"element_dim"};
+    auto P = SymbolicSize{"num_pages"};
+    auto cache_dtype = SymbolicDType{};
+    auto indices_dtype = SymbolicDType{};
+    auto dst_indices_dtype = SymbolicDType{};
+    auto device_ = SymbolicDevice{};
+
+    TensorMatcher({T, N, D})  //
+        .with_dtype(cache_dtype)
+        .with_device<kDLCUDA>(device_)
+        .verify(staging_k);
+    if constexpr (!kIsMLA) {
+      TensorMatcher({T, N, D})  //
+          .with_dtype(cache_dtype)
+          .with_device<kDLCUDA>(device_)
+          .verify(staging_v);
+    }
+    TensorMatcher({-1, N, D})  //
+        .with_dtype(cache_dtype)
+        .with_device<kDLCPU, kDLCUDAHost>()
+        .verify(k_cache_dst);
+    if constexpr (!kIsMLA) {
+      TensorMatcher({-1, N, D})  //
+          .with_dtype(cache_dtype)
+          .with_device<kDLCPU, kDLCUDAHost>()
+          .verify(v_cache_dst);
+    }
+    TensorMatcher({N})  //
+        .with_dtype<uint64_t>()
+        .with_device<kDLCUDA>(device_)
+        .verify(k_ptr_src);
+    if constexpr (!kIsMLA) {
+      TensorMatcher({N})  //
+          .with_dtype<uint64_t>()
+          .with_device<kDLCUDA>(device_)
+          .verify(v_ptr_src);
+    }
+    TensorMatcher({P})  //
+        .with_dtype<int32_t, int64_t>(indices_dtype)
+        .with_device<kDLCUDA>(device_)
+        .verify(page_indices_src);
+    TensorMatcher({T})  //
+        .with_dtype<int64_t>(dst_indices_dtype)
+        .with_device<kDLCPU, kDLCUDAHost>()
+        .verify(dst_indices_cpu);
+
+    RuntimeCheck(page_size > 0, "HiCache staged relayout: page_size must be positive");
+    RuntimeCheck(T.unwrap() == P.unwrap() * page_size, "HiCache staged relayout: staging token count mismatch");
+    RuntimeCheck(
+        kElementSize == D.unwrap() * dtype_bytes(cache_dtype.unwrap()),
+        "HiCache staged relayout: element size mismatch");
+    RuntimeCheck(kElementSize % 16 == 0, "HiCache staged relayout: element size must be 16-byte aligned");
+
+    const auto params = HicacheRelayoutParams{
+        .k_cache_dst = staging_k.data_ptr(),
+        .v_cache_dst = kIsMLA ? nullptr : staging_v.data_ptr(),
+        .indices_src = page_indices_src.data_ptr(),
+        .k_ptr_src = k_ptr_src.data_ptr(),
+        .v_ptr_src = kIsMLA ? nullptr : v_ptr_src.data_ptr(),
+        .num_pages = static_cast<uint32_t>(P.unwrap()),
+        .num_layers = static_cast<uint32_t>(N.unwrap()),
+        .page_size = static_cast<uint32_t>(page_size),
+    };
+    const auto device = device_.unwrap();
+    const auto use_int32 = indices_dtype.unwrap().bits == 32;
+    launch_hicache_relayout_kernel<kElementSize, kIsMLA>(params, P.unwrap(), N.unwrap(), page_size, use_int32, device);
+
+    auto stream = LaunchKernel::resolve_device(device);
+    const int64_t* dst_indices_ptr = static_cast<const int64_t*>(dst_indices_cpu.data_ptr());
+    if constexpr (kIsMLA) {
+      if (!try_copy_page_first_pages_batch(
+              {staging_k}, {k_cache_dst}, dst_indices_ptr, P.unwrap(), page_size, device.device_id, stream)) {
+        copy_page_first_pages_fallback({staging_k}, {k_cache_dst}, dst_indices_ptr, P.unwrap(), page_size, stream);
+      }
+    } else {
+      if (!try_copy_page_first_pages_batch(
+              {staging_k, staging_v},
+              {k_cache_dst, v_cache_dst},
+              dst_indices_ptr,
+              P.unwrap(),
+              page_size,
+              device.device_id,
+              stream)) {
+        copy_page_first_pages_fallback(
+            {staging_k, staging_v}, {k_cache_dst, v_cache_dst}, dst_indices_ptr, P.unwrap(), page_size, stream);
+      }
+    }
+  }
+
+ public:
+  static void run_all_lf_pf_staged(
+      const tvm::ffi::TensorView k_cache_dst,
+      const tvm::ffi::TensorView v_cache_dst,
+      const tvm::ffi::TensorView dst_indices_cpu,
+      const tvm::ffi::TensorView staging_k,
+      const tvm::ffi::TensorView staging_v,
+      const tvm::ffi::TensorView page_indices_src,
+      const tvm::ffi::TensorView k_ptr_src,
+      const tvm::ffi::TensorView v_ptr_src,
+      const int64_t page_size) {
+    run_staged_impl<false>(
+        k_cache_dst,
+        v_cache_dst,
+        dst_indices_cpu,
+        staging_k,
+        staging_v,
+        page_indices_src,
+        k_ptr_src,
+        v_ptr_src,
+        page_size);
+  }
+
+  static void run_all_mla_lf_pf_staged(
+      const tvm::ffi::TensorView cache_dst,
+      const tvm::ffi::TensorView dst_indices_cpu,
+      const tvm::ffi::TensorView staging,
+      const tvm::ffi::TensorView page_indices_src,
+      const tvm::ffi::TensorView ptr_src,
+      const int64_t page_size) {
+    run_staged_impl<true>(
+        cache_dst, cache_dst, dst_indices_cpu, staging, staging, page_indices_src, ptr_src, ptr_src, page_size);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/hicache.py

@@ -24,12 +24,24 @@ def _jit_hicache_module(*, element_size: int, unroll: int, block_quota: int) ->
     return load_jit(
         "hicache",
         *args,
-        cuda_files=["hicache.cuh"],
+        cuda_files=[
+            "kvcacheio/hicache.cuh",
+            "kvcacheio/relayout.cuh",
+            "kvcacheio/staged_write_back.cuh",
+        ],
         cuda_wrappers=[
             ("launch_one", f"&HiCacheKernel<{args}>::run_one"),
             ("launch_all", f"&HiCacheKernel<{args}>::run_all"),
             ("launch_one_mla", f"&HiCacheKernel<{args}>::run_one_mla"),
             ("launch_all_mla", f"&HiCacheKernel<{args}>::run_all_mla"),
+            (
+                "launch_all_lf_pf_staged",
+                f"&HiCacheStagedWriteBackKernel<{args}>::run_all_lf_pf_staged",
+            ),
+            (
+                "launch_all_mla_lf_pf_staged",
+                f"&HiCacheStagedWriteBackKernel<{args}>::run_all_mla_lf_pf_staged",
+            ),
         ],
     )
 
@@ -203,3 +215,93 @@ def transfer_hicache_all_layer_mla(
         cache_src_stride_bytes,
         cache_dst_stride_bytes,
     )
+
+
+@debug_kernel_api
+def transfer_hicache_all_layer_staged_lf_pf(
+    k_ptr_src: torch.Tensor,
+    v_ptr_src: torch.Tensor,
+    src_indices: torch.Tensor,
+    dst_indices: torch.Tensor,
+    staging_k: torch.Tensor,
+    staging_v: torch.Tensor,
+    dst_k: torch.Tensor,
+    dst_v: torch.Tensor,
+    *,
+    page_size: int,
+    element_size: int | None = None,
+    unroll: int | None = None,
+    block_quota: int | None = None,
+) -> None:
+    element_dim = staging_k[0, 0].numel()
+    element_size = element_size or (element_dim * staging_k.element_size())
+    block_quota = block_quota or DEFAULT_BLOCK_QUOTA
+    unroll = unroll or _default_unroll(element_size)
+    src_page_indices = src_indices[::page_size].contiguous()
+    module = _jit_hicache_module(
+        element_size=element_size,
+        unroll=unroll,
+        block_quota=block_quota,
+    )
+    staging_page_capacity = staging_k.shape[0] // page_size
+    staging_k = staging_k.view(staging_k.shape[0], staging_k.shape[1], -1)
+    staging_v = staging_v.view(staging_v.shape[0], staging_v.shape[1], -1)
+    dst_k = dst_k.view(dst_k.shape[0], dst_k.shape[1], -1)
+    dst_v = dst_v.view(dst_v.shape[0], dst_v.shape[1], -1)
+    for page_begin in range(0, src_page_indices.numel(), staging_page_capacity):
+        chunk_pages = min(staging_page_capacity, src_page_indices.numel() - page_begin)
+        chunk_tokens = chunk_pages * page_size
+        module.launch_all_lf_pf_staged(
+            dst_k,
+            dst_v,
+            dst_indices[
+                page_begin * page_size : (page_begin + chunk_pages) * page_size
+            ],
+            staging_k[:chunk_tokens],
+            staging_v[:chunk_tokens],
+            src_page_indices[page_begin : page_begin + chunk_pages],
+            k_ptr_src,
+            v_ptr_src,
+            page_size,
+        )
+
+
+@debug_kernel_api
+def transfer_hicache_all_layer_mla_staged_lf_pf(
+    ptr_src: torch.Tensor,
+    src_indices: torch.Tensor,
+    dst_indices: torch.Tensor,
+    staging: torch.Tensor,
+    dst: torch.Tensor,
+    *,
+    page_size: int,
+    element_size: int | None = None,
+    unroll: int | None = None,
+    block_quota: int | None = None,
+) -> None:
+    element_dim = staging[0, 0].numel()
+    element_size = element_size or (element_dim * staging.element_size())
+    block_quota = block_quota or DEFAULT_BLOCK_QUOTA
+    unroll = unroll or _default_unroll(element_size)
+    src_page_indices = src_indices[::page_size].contiguous()
+    module = _jit_hicache_module(
+        element_size=element_size,
+        unroll=unroll,
+        block_quota=block_quota,
+    )
+    staging_page_capacity = staging.shape[0] // page_size
+    staging = staging.view(staging.shape[0], staging.shape[1], -1)
+    dst = dst.view(dst.shape[0], dst.shape[1], -1)
+    for page_begin in range(0, src_page_indices.numel(), staging_page_capacity):
+        chunk_pages = min(staging_page_capacity, src_page_indices.numel() - page_begin)
+        chunk_tokens = chunk_pages * page_size
+        module.launch_all_mla_lf_pf_staged(
+            dst,
+            dst_indices[
+                page_begin * page_size : (page_begin + chunk_pages) * page_size
+            ],
+            staging[:chunk_tokens],
+            src_page_indices[page_begin : page_begin + chunk_pages],
+            ptr_src,
+            page_size,
+        )

python/sglang/srt/managers/cache_controller.py

@@ -726,9 +726,15 @@ class HiCacheController:
             return
 
         op = CacheOperation.merge_ops(self.write_queue)
-        host_indices, device_indices = self.move_indices(
-            op.host_indices, op.device_indices
-        )
+        # For now, kernel write-back keeps host indices on CPU only for page_first.
+        # More layouts can use this path once their write-back kernels accept CPU
+        # destination indices.
+        if self.io_backend == "kernel" and self.mem_pool_host.layout == "page_first":
+            host_indices, device_indices = op.host_indices, op.device_indices
+        else:
+            host_indices, device_indices = self.move_indices(
+                op.host_indices, op.device_indices
+            )
         self.write_queue.clear()
 
         start_event = device_module.Event()

python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py

@@ -394,9 +394,17 @@ class HybridCacheController(BaseHiCacheController):
         if not self.write_queue:
             return
         op = CacheOperation.merge_ops(self.write_queue)
-        host_indices, device_indices, resolved_pool_transfers = (
-            self.move_hybrid_indices(op)
-        )
+        # For now, kernel write-back keeps host indices on CPU only for page_first.
+        # More layouts can use this path once their write-back kernels accept CPU
+        # destination indices.
+        if self.io_backend == "kernel" and self.mem_pool_host.layout == "page_first":
+            host_indices = op.host_indices
+            device_indices = op.device_indices
+            resolved_pool_transfers = op.pool_transfers
+        else:
+            host_indices, device_indices, resolved_pool_transfers = (
+                self.move_hybrid_indices(op)
+            )
         self.write_queue.clear()
         start_event = device_module.Event()
         finish_event = device_module.Event()

python/sglang/srt/mem_cache/memory_pool_host.py

@@ -24,6 +24,12 @@ from sglang.jit_kernel.hicache import (
 from sglang.jit_kernel.hicache import (
     transfer_hicache_all_layer_mla as jit_transfer_hicache_all_layer_mla,
 )
+from sglang.jit_kernel.hicache import (
+    transfer_hicache_all_layer_mla_staged_lf_pf as jit_transfer_hicache_all_layer_mla_staged_lf_pf,
+)
+from sglang.jit_kernel.hicache import (
+    transfer_hicache_all_layer_staged_lf_pf as jit_transfer_hicache_all_layer_staged_lf_pf,
+)
 from sglang.jit_kernel.hicache import (
     transfer_hicache_one_layer as jit_transfer_hicache_one_layer,
 )
@@ -70,6 +76,8 @@ logger = logging.getLogger(__name__)
 # Host RAM to leave free when sizing HiCache pools (OS, other processes).
 HICACHE_HOST_MEMORY_RESERVE_BYTES: int = 10 * (1024**3)
 
+_WRITE_BACK_STAGING_PAGE_CHUNK = 64
+
 
 def synchronized(func):
     @wraps(func)
@@ -428,6 +436,7 @@ class MHATokenToKVPoolHost(HostKVCache):
             dtype=torch.uint64,
             device=self.device_pool.device,
         )
+        self._init_write_back_staging_buffers()
 
     def get_size_per_token(self):
         self.head_num = self.device_pool.head_num
@@ -476,6 +485,28 @@ class MHATokenToKVPoolHost(HostKVCache):
         )
         return buffer
 
+    def _init_write_back_staging_buffers(self):
+        self.staging_page_capacity = 0
+        self.staging_token_capacity = 0
+        self.staging_k_buffer = None
+        self.staging_v_buffer = None
+        if self.layout != "page_first" or (_is_npu or _is_xpu or _is_mps):
+            return
+
+        self.staging_page_capacity = min(self.page_num, _WRITE_BACK_STAGING_PAGE_CHUNK)
+        self.staging_token_capacity = self.staging_page_capacity * self.page_size
+        self.staging_k_buffer = torch.empty(
+            (
+                self.staging_token_capacity,
+                self.layer_num,
+                self.head_num,
+                self.head_dim,
+            ),
+            dtype=self.dtype,
+            device=self.device_pool.device,
+        )
+        self.staging_v_buffer = torch.empty_like(self.staging_k_buffer)
+
     @property
     def k_buffer(self):
         return self.kv_buffer[0]
@@ -631,18 +662,16 @@ class MHATokenToKVPoolHost(HostKVCache):
                     )
             elif self.layout == "page_first":
                 if self.can_use_jit:
-                    # Use transposed data ptrs so the kernel writes to
-                    # [layer, page, item] view with stride layout_dim per token.
-                    jit_transfer_hicache_all_layer(
-                        k_ptr_dst=self.k_data_ptrs,
-                        v_ptr_dst=self.v_data_ptrs,
-                        indices_dst=host_indices,
+                    jit_transfer_hicache_all_layer_staged_lf_pf(
                         k_ptr_src=device_pool.k_data_ptrs,
                         v_ptr_src=device_pool.v_data_ptrs,
-                        indices_src=device_indices,
-                        kv_cache_src_stride_bytes=self.token_stride_size,
-                        kv_cache_dst_stride_bytes=self.layout_dim,
-                        element_size=self.element_dim * self.dtype.itemsize,
+                        src_indices=device_indices,
+                        dst_indices=host_indices,
+                        staging_k=self.staging_k_buffer,
+                        staging_v=self.staging_v_buffer,
+                        dst_k=self.k_buffer,
+                        dst_v=self.v_buffer,
+                        page_size=self.page_size,
                     )
                 else:
                     transfer_kv_all_layer_lf_pf(
@@ -1194,6 +1223,7 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
             dtype=torch.uint64,
             device=self.device_pool.device,
         )
+        self._init_write_back_staging_buffers()
 
     def get_contiguous_buf_infos(self):
         """Return (data_ptrs, data_lens, item_lens) in the same format as device pool,
@@ -1289,6 +1319,26 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
         )
         return buffer
 
+    def _init_write_back_staging_buffers(self):
+        self.staging_page_capacity = 0
+        self.staging_token_capacity = 0
+        self.staging_buffer = None
+        if self.layout != "page_first" or (_is_npu or _is_xpu or _is_mps):
+            return
+
+        self.staging_page_capacity = min(self.page_num, _WRITE_BACK_STAGING_PAGE_CHUNK)
+        self.staging_token_capacity = self.staging_page_capacity * self.page_size
+        self.staging_buffer = torch.empty(
+            (
+                self.staging_token_capacity,
+                self.layer_num,
+                1,
+                self.kv_cache_dim,
+            ),
+            dtype=self.dtype,
+            device=self.device_pool.device,
+        )
+
     def load_to_device_per_layer(
         self, device_pool, host_indices, device_indices, layer_id, io_backend
     ):
@@ -1398,14 +1448,13 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
                     )
             elif self.layout == "page_first":
                 if self.can_use_jit:
-                    jit_transfer_hicache_all_layer_mla(
-                        ptr_dst=self.data_ptrs,
-                        indices_dst=host_indices,
+                    jit_transfer_hicache_all_layer_mla_staged_lf_pf(
                         ptr_src=device_pool.data_ptrs,
-                        indices_src=device_indices,
-                        cache_src_stride_bytes=self.token_stride_size,
-                        cache_dst_stride_bytes=self.layout_dim,
-                        element_size=self.kv_cache_dim * self.dtype.itemsize,
+                        src_indices=device_indices,
+                        dst_indices=host_indices,
+                        staging=self.staging_buffer,
+                        dst=self.kv_buffer,
+                        page_size=self.page_size,
                     )
                 else:
                     transfer_kv_all_layer_mla_lf_pf(

python/sglang/srt/server_args.py

@@ -701,7 +701,7 @@ class ServerArgs:
     hicache_size: int = 0
     hicache_write_policy: str = "write_through"
     hicache_io_backend: str = "kernel"
-    hicache_mem_layout: str = "layer_first"
+    hicache_mem_layout: str = "page_first"
     hicache_storage_backend: Optional[str] = None
     hicache_storage_prefetch_policy: str = "timeout"
     hicache_storage_backend_extra_config: Optional[str] = None
@@ -3987,8 +3987,6 @@ class ServerArgs:
         Resolution order:
         1) Layout <-> I/O compatibility for direct conflicts.
         2) Storage <-> layout compatibility (may rewrite layout).
-        3) I/O <-> decode-attention compatibility (may rewrite I/O or decode backend).
-        4) Re-run step (1) if step (3) changed I/O backend.
         """
         # Skip all normalization when neither hicache nor decode-offload path is active.
         if not (
@@ -4003,13 +4001,6 @@ class ServerArgs:
         # Step 2: Storage-layout normalization without changing io backend.
         self._resolve_storage_layout_compatibility()
 
-        # Step 3: IO-decode backend compatibility (may change io backend).
-        io_changed = self._resolve_io_decode_attention_compatibility()
-
-        # Step 4: Re-normalize layout after io backend changes.
-        if io_changed:
-            self._resolve_layout_io_compatibility()
-
     def _resolve_layout_io_compatibility(self):
         if (
             self.hicache_mem_layout == "page_first_direct"
@@ -4050,41 +4041,6 @@ class ServerArgs:
             f"switching to {new_layout} layout for {self.hicache_io_backend} io backend"
         )
 
-    def _resolve_io_decode_attention_compatibility(self) -> bool:
-        if self.hicache_io_backend != "kernel":
-            return False
-
-        # Only patch settings when the effective decode backend is FA3.
-        effective_decode_backend = (
-            self.decode_attention_backend or self.attention_backend
-        )
-        if effective_decode_backend != "fa3":
-            return False
-
-        if self.decode_attention_backend is not None:
-            self.hicache_io_backend = "direct"
-            logger.warning(
-                "FlashAttention3 decode backend is not compatible with hierarchical cache. "
-                "Setting hicache_io_backend to vanilla I/O, which may lead to suboptimal performance with small page sizes."
-            )
-            return True
-
-        # If decode backend is implicit, pick a safe backend without changing io backend.
-        if not self.use_mla_backend():
-            # FlashInfer does not support attention sinks.
-            if (
-                is_flashinfer_available()
-                and not self.get_model_config().has_attention_sinks
-            ):
-                self.decode_attention_backend = "flashinfer"
-            else:
-                self.decode_attention_backend = "triton"
-        else:
-            self.decode_attention_backend = (
-                "flashinfer" if is_sm100_supported() else "triton"
-            )
-        return False
-
     def _handle_load_format(self):
         if (
             self.load_format == "auto" or self.load_format == "gguf"