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

benchmark/hicache/bench_hicache_write_back.py

@@ -0,0 +1,391 @@
+from __future__ import annotations
+
+import csv
+import platform
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+import triton.testing
+from sgl_kernel.kvcacheio import (
+    transfer_kv_all_layer_lf_pf,
+    transfer_kv_all_layer_mla_lf_pf,
+)
+
+from sglang.jit_kernel.hicache import (
+    can_use_hicache_jit_kernel,
+    transfer_hicache_all_layer_mla_staged_lf_pf,
+    transfer_hicache_all_layer_staged_lf_pf,
+)
+
+ROOT = Path(__file__).resolve().parents[2]
+OUT_DIR = ROOT / "benchmark" / "hicache" / "profiles"
+CSV_PATH = OUT_DIR / "writeback_compare.csv"
+MD_PATH = ROOT / "benchmark" / "hicache" / "HICACHE_WRITEBACK_COMPARISON.md"
+
+DTYPE = torch.bfloat16
+PAGE_SIZE = 64
+NUM_LAYERS_LIST = [16, 24, 32, 40, 48, 56, 64, 72, 80]
+BATCH_PAGES = [4, 8, 16, 32, 64]
+ELEMENT_DIMS = [256, 512, 1024, 2048]
+TOTAL_PAGES = max(128, max(BATCH_PAGES) * 2)
+WARMUP = 5
+REP = 25
+
+
+@dataclass(frozen=True)
+class BenchRow:
+    mode: str
+    num_layers: int
+    element_dim: int
+    batch_pages: int
+    item_bytes: int
+    baseline_gib_s: float
+    staged_gib_s: float
+    speedup: float
+
+
+def _make_page_starts(total_pages: int, num_pages: int, page_size: int) -> torch.Tensor:
+    order = torch.randperm(total_pages, device="cuda", dtype=torch.int64)
+    return order[:num_pages] * page_size
+
+
+def _make_token_indices(
+    page_starts: torch.Tensor, page_size: int, *, device: str
+) -> torch.Tensor:
+    arange = torch.arange(page_size, device=device, dtype=torch.int64)
+    return (page_starts.to(device=device)[:, None] + arange).reshape(-1)
+
+
+def _to_gib_per_s(moved_bytes: int, ms: float) -> float:
+    return moved_bytes / (ms * 1e-3) / (1024**3)
+
+
+def _bench(fn) -> float:
+    ms = triton.testing.do_bench(fn, warmup=WARMUP, rep=REP)
+    return float(ms)
+
+
+def _validate_mha_correctness(
+    baseline,
+    staged,
+    dst_k_baseline: torch.Tensor,
+    dst_v_baseline: torch.Tensor,
+    dst_k_staged: torch.Tensor,
+    dst_v_staged: torch.Tensor,
+) -> None:
+    dst_k_baseline.zero_()
+    dst_v_baseline.zero_()
+    dst_k_staged.zero_()
+    dst_v_staged.zero_()
+    baseline()
+    staged()
+    torch.cuda.synchronize()
+    torch.testing.assert_close(dst_k_staged, dst_k_baseline)
+    torch.testing.assert_close(dst_v_staged, dst_v_baseline)
+
+
+def _validate_mla_correctness(
+    baseline,
+    staged,
+    dst_baseline: torch.Tensor,
+    dst_staged: torch.Tensor,
+) -> None:
+    dst_baseline.zero_()
+    dst_staged.zero_()
+    baseline()
+    staged()
+    torch.cuda.synchronize()
+    torch.testing.assert_close(dst_staged, dst_baseline)
+
+
+def _bench_mha(num_layers: int, element_dim: int, batch_pages: int) -> BenchRow:
+    item_bytes = element_dim * torch.tensor([], dtype=DTYPE).element_size()
+    assert can_use_hicache_jit_kernel(element_size=item_bytes)
+
+    total_tokens = TOTAL_PAGES * PAGE_SIZE
+    chunk_tokens = batch_pages * PAGE_SIZE
+    src_page_starts = _make_page_starts(TOTAL_PAGES, batch_pages, PAGE_SIZE)
+    dst_page_starts = _make_page_starts(TOTAL_PAGES, batch_pages, PAGE_SIZE)
+    src_indices = _make_token_indices(src_page_starts, PAGE_SIZE, device="cuda")
+    dst_indices = _make_token_indices(dst_page_starts, PAGE_SIZE, device="cuda")
+
+    src_k_layers = [
+        torch.randn(total_tokens, element_dim, dtype=DTYPE, device="cuda")
+        for _ in range(num_layers)
+    ]
+    src_v_layers = [
+        torch.randn(total_tokens, element_dim, dtype=DTYPE, device="cuda")
+        for _ in range(num_layers)
+    ]
+    src_k_ptrs = torch.tensor(
+        [x.data_ptr() for x in src_k_layers], dtype=torch.uint64, device="cuda"
+    )
+    src_v_ptrs = torch.tensor(
+        [x.data_ptr() for x in src_v_layers], dtype=torch.uint64, device="cuda"
+    )
+
+    dst_k_baseline = torch.empty(
+        total_tokens, num_layers, element_dim, dtype=DTYPE, pin_memory=True
+    )
+    dst_v_baseline = torch.empty_like(dst_k_baseline, pin_memory=True)
+    dst_k_staged = torch.empty_like(dst_k_baseline, pin_memory=True)
+    dst_v_staged = torch.empty_like(dst_v_baseline, pin_memory=True)
+    staging_k = torch.empty(
+        chunk_tokens, num_layers, element_dim, dtype=DTYPE, device="cuda"
+    )
+    staging_v = torch.empty_like(staging_k)
+
+    baseline = lambda: transfer_kv_all_layer_lf_pf(
+        src_k_layers=src_k_ptrs,
+        dst_k=dst_k_baseline,
+        src_v_layers=src_v_ptrs,
+        dst_v=dst_v_baseline,
+        src_indices=src_indices,
+        dst_indices=dst_indices,
+        item_size=item_bytes,
+        dst_layout_dim=item_bytes * num_layers,
+        num_layers=num_layers,
+    )
+    staged = lambda: transfer_hicache_all_layer_staged_lf_pf(
+        k_ptr_src=src_k_ptrs,
+        v_ptr_src=src_v_ptrs,
+        src_indices=src_indices,
+        dst_indices=dst_indices,
+        staging_k=staging_k,
+        staging_v=staging_v,
+        dst_k=dst_k_staged,
+        dst_v=dst_v_staged,
+        page_size=PAGE_SIZE,
+    )
+
+    _validate_mha_correctness(
+        baseline,
+        staged,
+        dst_k_baseline,
+        dst_v_baseline,
+        dst_k_staged,
+        dst_v_staged,
+    )
+
+    moved_bytes = chunk_tokens * num_layers * item_bytes * 2
+    baseline_ms = _bench(baseline)
+    staged_ms = _bench(staged)
+    baseline_bw = _to_gib_per_s(moved_bytes, baseline_ms)
+    staged_bw = _to_gib_per_s(moved_bytes, staged_ms)
+    return BenchRow(
+        mode="MHA",
+        num_layers=num_layers,
+        element_dim=element_dim,
+        batch_pages=batch_pages,
+        item_bytes=item_bytes,
+        baseline_gib_s=baseline_bw,
+        staged_gib_s=staged_bw,
+        speedup=staged_bw / baseline_bw,
+    )
+
+
+def _bench_mla(num_layers: int, element_dim: int, batch_pages: int) -> BenchRow:
+    item_bytes = element_dim * torch.tensor([], dtype=DTYPE).element_size()
+    assert can_use_hicache_jit_kernel(element_size=item_bytes)
+
+    total_tokens = TOTAL_PAGES * PAGE_SIZE
+    chunk_tokens = batch_pages * PAGE_SIZE
+    src_page_starts = _make_page_starts(TOTAL_PAGES, batch_pages, PAGE_SIZE)
+    dst_page_starts = _make_page_starts(TOTAL_PAGES, batch_pages, PAGE_SIZE)
+    src_indices = _make_token_indices(src_page_starts, PAGE_SIZE, device="cuda")
+    dst_indices = _make_token_indices(dst_page_starts, PAGE_SIZE, device="cuda")
+
+    src_layers = [
+        torch.randn(total_tokens, element_dim, dtype=DTYPE, device="cuda")
+        for _ in range(num_layers)
+    ]
+    src_ptrs = torch.tensor(
+        [x.data_ptr() for x in src_layers], dtype=torch.uint64, device="cuda"
+    )
+    dst_baseline = torch.empty(
+        total_tokens, num_layers, element_dim, dtype=DTYPE, pin_memory=True
+    )
+    dst_staged = torch.empty_like(dst_baseline, pin_memory=True)
+    staging = torch.empty(
+        chunk_tokens, num_layers, element_dim, dtype=DTYPE, device="cuda"
+    )
+
+    baseline = lambda: transfer_kv_all_layer_mla_lf_pf(
+        src_layers=src_ptrs,
+        dst=dst_baseline,
+        src_indices=src_indices,
+        dst_indices=dst_indices,
+        item_size=item_bytes,
+        dst_layout_dim=item_bytes * num_layers,
+        num_layers=num_layers,
+    )
+    staged = lambda: transfer_hicache_all_layer_mla_staged_lf_pf(
+        ptr_src=src_ptrs,
+        src_indices=src_indices,
+        dst_indices=dst_indices,
+        staging=staging,
+        dst=dst_staged,
+        page_size=PAGE_SIZE,
+    )
+
+    _validate_mla_correctness(
+        baseline,
+        staged,
+        dst_baseline,
+        dst_staged,
+    )
+
+    moved_bytes = chunk_tokens * num_layers * item_bytes
+    baseline_ms = _bench(baseline)
+    staged_ms = _bench(staged)
+    baseline_bw = _to_gib_per_s(moved_bytes, baseline_ms)
+    staged_bw = _to_gib_per_s(moved_bytes, staged_ms)
+    return BenchRow(
+        mode="MLA",
+        num_layers=num_layers,
+        element_dim=element_dim,
+        batch_pages=batch_pages,
+        item_bytes=item_bytes,
+        baseline_gib_s=baseline_bw,
+        staged_gib_s=staged_bw,
+        speedup=staged_bw / baseline_bw,
+    )
+
+
+def _write_csv(rows: list[BenchRow]) -> None:
+    OUT_DIR.mkdir(parents=True, exist_ok=True)
+    with CSV_PATH.open("w", newline="") as f:
+        writer = csv.writer(f)
+        writer.writerow(
+            [
+                "mode",
+                "num_layers",
+                "element_dim",
+                "item_bytes",
+                "batch_pages",
+                "baseline_gib_s",
+                "staged_gib_s",
+                "speedup",
+            ]
+        )
+        for row in rows:
+            writer.writerow(
+                [
+                    row.mode,
+                    row.num_layers,
+                    row.element_dim,
+                    row.item_bytes,
+                    row.batch_pages,
+                    f"{row.baseline_gib_s:.6f}",
+                    f"{row.staged_gib_s:.6f}",
+                    f"{row.speedup:.4f}",
+                ]
+            )
+
+
+def _table(rows: list[BenchRow], mode: str) -> str:
+    lines = [
+        "| num_layers | element_dim | item_bytes | batch_pages | `sgl_kernel *_lf_pf` GiB/s | `jit *_staged_lf_pf` GiB/s | speedup |",
+        "| ---: | ---: | ---: | ---: | ---: | ---: | ---: |",
+    ]
+    for row in rows:
+        if row.mode != mode:
+            continue
+        lines.append(
+            f"| {row.num_layers} | {row.element_dim} | {row.item_bytes} | {row.batch_pages} | "
+            f"{row.baseline_gib_s:.3f} | {row.staged_gib_s:.3f} | {row.speedup:.2f}x |"
+        )
+    return "\n".join(lines)
+
+
+def _summarize(rows: list[BenchRow], mode: str) -> tuple[BenchRow, BenchRow]:
+    filtered = [row for row in rows if row.mode == mode]
+    best = max(filtered, key=lambda row: row.speedup)
+    worst = min(filtered, key=lambda row: row.speedup)
+    return best, worst
+
+
+def _write_report_en(rows: list[BenchRow]) -> None:
+    gpu_name = torch.cuda.get_device_name(0)
+    mha_best, mha_worst = _summarize(rows, "MHA")
+    mla_best, mla_worst = _summarize(rows, "MLA")
+    md = f"""# HiCache Writeback Benchmark Comparison
+
+Environment:
+
+- GPU: `{gpu_name}`
+- Platform: `{platform.platform()}`
+- DType: `{DTYPE}`
+- `page_size`: `{PAGE_SIZE}`
+- `num_layers` sweep: `{NUM_LAYERS_LIST}`
+- `total_pages`: `{TOTAL_PAGES}`
+- `batch_pages` sweep: `{BATCH_PAGES}`
+- `element_dim` sweep: `{ELEMENT_DIMS}`
+- `src_indices` / `dst_indices`: CUDA token indices
+- Timing: `triton.testing.do_bench(warmup={WARMUP}, rep={REP})`
+
+Comparison target:
+
+- MHA: `sgl_kernel.transfer_kv_all_layer_lf_pf` vs `sglang.jit_kernel.hicache.transfer_hicache_all_layer_staged_lf_pf`
+- MLA: `sgl_kernel.transfer_kv_all_layer_mla_lf_pf` vs `sglang.jit_kernel.hicache.transfer_hicache_all_layer_mla_staged_lf_pf`
+
+Metric:
+
+- Effective writeback bandwidth in `GiB/s`
+- `speedup = staged_jit / sgl_kernel_lf_pf`
+
+## MHA
+
+{_table(rows, "MHA")}
+
+## MLA
+
+{_table(rows, "MLA")}
+
+## Takeaways
+
+- The staged JIT path is compared directly against the installed `sgl_kernel` LF->PF kernels on the same host-destination writeback workload.
+- Best MHA staged case: `num_layers={mha_best.num_layers}`, `element_dim={mha_best.element_dim}`, `batch_pages={mha_best.batch_pages}`, `speedup={mha_best.speedup:.2f}x`.
+- Weakest MHA staged case: `num_layers={mha_worst.num_layers}`, `element_dim={mha_worst.element_dim}`, `batch_pages={mha_worst.batch_pages}`, `speedup={mha_worst.speedup:.2f}x`.
+- Best MLA staged case: `num_layers={mla_best.num_layers}`, `element_dim={mla_best.element_dim}`, `batch_pages={mla_best.batch_pages}`, `speedup={mla_best.speedup:.2f}x`.
+- Weakest MLA staged case: `num_layers={mla_worst.num_layers}`, `element_dim={mla_worst.element_dim}`, `batch_pages={mla_worst.batch_pages}`, `speedup={mla_worst.speedup:.2f}x`.
+- MHA bandwidth uses combined payload bytes for `K + V`.
+- MLA bandwidth uses single-buffer payload bytes.
+- Raw data is also available in `benchmark/hicache/profiles/writeback_compare.csv`.
+"""
+    MD_PATH.write_text(md)
+
+
+def main() -> None:
+    torch.manual_seed(0)
+    rows: list[BenchRow] = []
+    for num_layers in NUM_LAYERS_LIST:
+        for element_dim in ELEMENT_DIMS:
+            for batch_pages in BATCH_PAGES:
+                rows.append(_bench_mha(num_layers, element_dim, batch_pages))
+                rows.append(_bench_mla(num_layers, element_dim, batch_pages))
+    _write_csv(rows)
+    _write_report_en(rows)
+    for row in rows:
+        print(
+            row.mode,
+            "num_layers=",
+            row.num_layers,
+            "element_dim=",
+            row.element_dim,
+            "batch_pages=",
+            row.batch_pages,
+            "baseline_gib_s=",
+            f"{row.baseline_gib_s:.3f}",
+            "staged_gib_s=",
+            f"{row.staged_gib_s:.3f}",
+            "speedup=",
+            f"{row.speedup:.2f}x",
+        )
+    print(f"wrote {CSV_PATH}")
+    print(f"wrote {MD_PATH}")
+
+
+if __name__ == "__main__":
+    main()

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"

test/registered/jit/test_hicache.py

@@ -31,17 +31,21 @@ POOL_SIZE = PAGE_SIZE * 8
 MHA_ELEMENT_DIMS = [128, 256, 512, 1024]
 MLA_ELEMENT_DIMS = [576]
 LAYOUTS = ["layer_first", "page_first"]
+STAGED_WRITE_BACK_PAGE_COUNTS = [1, 63, 64, 65, 67, 128, 129]
 
 
 def _token_indices_for_pages(
-    pages: torch.Tensor, page_size: int = PAGE_SIZE, device: str = DEVICE
+    pages: torch.Tensor,
+    page_size: int = PAGE_SIZE,
+    device: str = DEVICE,
+    dtype: torch.dtype = torch.int64,
 ) -> torch.Tensor:
     parts = [
         torch.arange(
             int(page) * page_size,
             (int(page) + 1) * page_size,
             device=device,
-            dtype=torch.int64,
+            dtype=dtype,
         )
         for page in pages.tolist()
     ]
@@ -71,6 +75,12 @@ def _copy_tensor_with_offset(tensor: torch.Tensor, offset: int) -> None:
     tensor.copy_(data + offset)
 
 
+def _assert_page_filled(tensor: torch.Tensor, page: int, value: float) -> None:
+    page_slice = tensor[page * PAGE_SIZE : (page + 1) * PAGE_SIZE]
+    expected = torch.full_like(page_slice, value)
+    assert torch.equal(page_slice.cpu(), expected.cpu())
+
+
 def _run_transfer_roundtrip_mha(layout: str, element_dim: int) -> None:
     device_pool = MHATokenToKVPool(
         size=POOL_SIZE,
@@ -99,9 +109,14 @@ def _run_transfer_roundtrip_mha(layout: str, element_dim: int) -> None:
     host_pages = torch.tensor([0, 1, 2], device=DEVICE, dtype=torch.int64)
     device_indices = _token_indices_for_pages(device_pages)
     host_indices = _token_indices_for_pages(host_pages)
+    host_indices_backup = (
+        _token_indices_for_pages(host_pages, device="cpu")
+        if layout == "page_first"
+        else host_indices
+    )
 
     host_pool.backup_from_device_all_layer(
-        device_pool, host_indices, device_indices, "kernel"
+        device_pool, host_indices_backup, device_indices, "kernel"
     )
     torch.cuda.synchronize()
 
@@ -187,9 +202,14 @@ def _run_transfer_roundtrip_mla(layout: str, element_dim: int) -> None:
     host_pages = torch.tensor([0, 1, 2], device=DEVICE, dtype=torch.int64)
     device_indices = _token_indices_for_pages(device_pages)
     host_indices = _token_indices_for_pages(host_pages)
+    host_indices_backup = (
+        _token_indices_for_pages(host_pages, device="cpu")
+        if layout == "page_first"
+        else host_indices
+    )
 
     host_pool.backup_from_device_all_layer(
-        device_pool, host_indices, device_indices, "kernel"
+        device_pool, host_indices_backup, device_indices, "kernel"
     )
     torch.cuda.synchronize()
 
@@ -231,6 +251,153 @@ def _run_transfer_roundtrip_mla(layout: str, element_dim: int) -> None:
             )
 
 
+def _run_page_first_staged_write_back_mha(
+    layout: str, element_dim: int, page_count: int
+) -> None:
+    pool_size = PAGE_SIZE * (page_count + 8)
+    device_pool = MHATokenToKVPool(
+        size=pool_size,
+        page_size=PAGE_SIZE,
+        head_num=element_dim // 128,
+        head_dim=128,
+        dtype=torch.bfloat16,
+        layer_num=NUM_LAYERS,
+        device=DEVICE,
+        enable_memory_saver=False,
+    )
+    host_pool = _pinned_host_pool(
+        MHATokenToKVPoolHost,
+        device_pool=device_pool,
+        layout=layout,
+    )
+    assert host_pool.can_use_jit
+    assert host_pool.staging_page_capacity > 0
+    if page_count > 64:
+        assert host_pool.staging_page_capacity < page_count
+
+    for layer_id in range(NUM_LAYERS):
+        _copy_tensor_with_offset(device_pool.k_buffer[layer_id], layer_id)
+        _copy_tensor_with_offset(device_pool.v_buffer[layer_id], layer_id + 100)
+    host_pool.k_buffer.fill_(-7)
+    host_pool.v_buffer.fill_(-11)
+
+    device_pages = torch.arange(
+        2,
+        2 + page_count,
+        device=DEVICE,
+        dtype=torch.int64,
+    )
+    host_pages = torch.arange(
+        page_count,
+        0,
+        -1,
+        dtype=torch.int64,
+    )
+    src_index_dtype = torch.int32 if page_count == 64 else torch.int64
+    device_indices = _token_indices_for_pages(device_pages, dtype=src_index_dtype)
+    host_indices = _token_indices_for_pages(host_pages, device="cpu")
+    assert not host_indices.is_cuda
+
+    host_pool.backup_from_device_all_layer(
+        device_pool, host_indices, device_indices, "kernel"
+    )
+    torch.cuda.synchronize()
+
+    for layer_id in range(NUM_LAYERS):
+        for host_page, device_page in zip(host_pages.tolist(), device_pages.tolist()):
+            host_start = host_page * PAGE_SIZE
+            device_start = device_page * PAGE_SIZE
+            assert torch.equal(
+                host_pool.k_data_refs[layer_id][
+                    host_start : host_start + PAGE_SIZE
+                ].cpu(),
+                device_pool.k_buffer[layer_id][
+                    device_start : device_start + PAGE_SIZE
+                ].cpu(),
+            )
+            assert torch.equal(
+                host_pool.v_data_refs[layer_id][
+                    host_start : host_start + PAGE_SIZE
+                ].cpu(),
+                device_pool.v_buffer[layer_id][
+                    device_start : device_start + PAGE_SIZE
+                ].cpu(),
+            )
+
+    for layer_id in range(NUM_LAYERS):
+        for untouched_page in [0, page_count + 1]:
+            _assert_page_filled(host_pool.k_data_refs[layer_id], untouched_page, -7)
+            _assert_page_filled(host_pool.v_data_refs[layer_id], untouched_page, -11)
+
+
+def _run_page_first_staged_write_back_mla(
+    layout: str, element_dim: int, page_count: int
+) -> None:
+    pool_size = PAGE_SIZE * (page_count + 8)
+    device_pool = MLATokenToKVPool(
+        size=pool_size,
+        page_size=PAGE_SIZE,
+        kv_lora_rank=element_dim - 64,
+        qk_rope_head_dim=64,
+        dtype=torch.bfloat16,
+        layer_num=NUM_LAYERS,
+        device=DEVICE,
+        enable_memory_saver=False,
+    )
+    host_pool = _pinned_host_pool(
+        MLATokenToKVPoolHost,
+        device_pool=device_pool,
+        layout=layout,
+    )
+    assert host_pool.can_use_jit
+    assert host_pool.staging_page_capacity > 0
+    if page_count > 64:
+        assert host_pool.staging_page_capacity < page_count
+
+    for layer_id in range(NUM_LAYERS):
+        _copy_tensor_with_offset(device_pool.kv_buffer[layer_id], layer_id)
+    host_pool.kv_buffer.fill_(-13)
+
+    device_pages = torch.arange(
+        2,
+        2 + page_count,
+        device=DEVICE,
+        dtype=torch.int64,
+    )
+    host_pages = torch.arange(
+        page_count,
+        0,
+        -1,
+        dtype=torch.int64,
+    )
+    src_index_dtype = torch.int32 if page_count == 64 else torch.int64
+    device_indices = _token_indices_for_pages(device_pages, dtype=src_index_dtype)
+    host_indices = _token_indices_for_pages(host_pages, device="cpu")
+    assert not host_indices.is_cuda
+
+    host_pool.backup_from_device_all_layer(
+        device_pool, host_indices, device_indices, "kernel"
+    )
+    torch.cuda.synchronize()
+
+    for layer_id in range(NUM_LAYERS):
+        for host_page, device_page in zip(host_pages.tolist(), device_pages.tolist()):
+            host_start = host_page * PAGE_SIZE
+            device_start = device_page * PAGE_SIZE
+            assert torch.equal(
+                host_pool.data_refs[layer_id][
+                    host_start : host_start + PAGE_SIZE
+                ].cpu(),
+                device_pool.kv_buffer[layer_id][
+                    device_start : device_start + PAGE_SIZE
+                ].cpu(),
+            )
+
+    for layer_id in range(NUM_LAYERS):
+        for untouched_page in [0, page_count + 1]:
+            _assert_page_filled(host_pool.data_refs[layer_id], untouched_page, -13)
+
+
 @pytest.mark.parametrize("layout", LAYOUTS)
 @pytest.mark.parametrize("element_dim", MHA_ELEMENT_DIMS)
 def test_hicache_transfer_mha(layout: str, element_dim: int) -> None:
@@ -243,5 +410,23 @@ def test_hicache_transfer_mla(layout: str, element_dim: int) -> None:
     _run_transfer_roundtrip_mla(layout, element_dim)
 
 
+@pytest.mark.parametrize("layout", ["page_first"])
+@pytest.mark.parametrize("element_dim", MHA_ELEMENT_DIMS)
+@pytest.mark.parametrize("page_count", STAGED_WRITE_BACK_PAGE_COUNTS)
+def test_hicache_page_first_staged_write_back_mha(
+    layout: str, element_dim: int, page_count: int
+) -> None:
+    _run_page_first_staged_write_back_mha(layout, element_dim, page_count)
+
+
+@pytest.mark.parametrize("layout", ["page_first"])
+@pytest.mark.parametrize("element_dim", MLA_ELEMENT_DIMS)
+@pytest.mark.parametrize("page_count", STAGED_WRITE_BACK_PAGE_COUNTS)
+def test_hicache_page_first_staged_write_back_mla(
+    layout: str, element_dim: int, page_count: int
+) -> None:
+    _run_page_first_staged_write_back_mla(layout, element_dim, page_count)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/unit/mem_cache/test_hiradix_cache_unit.py

@@ -37,6 +37,7 @@ class TestHiRadixCacheKVEvents(CustomTestCase):
             model_path="dummy",
             page_size=PAGE_SIZE,
             hicache_io_backend="direct",
+            hicache_mem_layout="layer_first",
             hicache_write_policy="write_through",
         )
         set_global_server_args_for_scheduler(server_args)

test/registered/unit/server_args/test_server_args.py

@@ -620,6 +620,14 @@ class TestHiCacheArgs(unittest.TestCase):
 
     def test_hicache_io_backend_and_mem_layout_compatibility(self):
         cases = [
+            {
+                "name": "default_kernel_page_first",
+                "overrides": {
+                    "enable_hierarchical_cache": True,
+                },
+                "expected_io_backend": "kernel",
+                "expected_mem_layout": "page_first",
+            },
             {
                 "name": "kernel_with_page_first_direct",
                 "overrides": {
@@ -660,8 +668,9 @@ class TestHiCacheArgs(unittest.TestCase):
                     "attention_backend": "triton",
                     "decode_attention_backend": "fa3",
                 },
-                "expected_io_backend": "direct",
-                "expected_mem_layout": "page_first_direct",
+                "expected_io_backend": "kernel",
+                "expected_mem_layout": "page_first",
+                "expected_decode_backend": "fa3",
             },
         ]
 
@@ -673,13 +682,10 @@ class TestHiCacheArgs(unittest.TestCase):
                     args,
                     expected_io_backend=case["expected_io_backend"],
                     expected_mem_layout=case["expected_mem_layout"],
+                    expected_decode_backend=case.get("expected_decode_backend"),
                 )
 
-    @patch.object(ServerArgs, "use_mla_backend", return_value=False)
-    @patch("sglang.srt.server_args.is_flashinfer_available", return_value=False)
-    def test_decode_attention_backend_with_implicit_fa3(
-        self, _mock_flashinfer, _mock_use_mla_backend
-    ):
+    def test_hicache_kernel_keeps_implicit_fa3_decode_backend(self):
         args = self._make_args(
             enable_hierarchical_cache=True,
             hicache_io_backend="kernel",
@@ -689,7 +695,9 @@ class TestHiCacheArgs(unittest.TestCase):
 
         args._handle_hicache()
 
-        self.assertEqual(args.decode_attention_backend, "triton")
+        self.assertEqual(args.hicache_io_backend, "kernel")
+        self.assertEqual(args.hicache_mem_layout, "page_first")
+        self.assertIsNone(args.decode_attention_backend)
 
 
 class TestNgramExternalSamArgs(CustomTestCase):