From b5bcd76a41b88d58d00932a260715185a2c15a3c Mon Sep 17 00:00:00 2001 From: huangtingwei <141888744+huangtingwei9988@users.noreply.github.com> Date: Tue, 16 Jun 2026 10:44:27 +0800 Subject: [PATCH] [HiCache & JIT Kernel] Refactoring HiCache Write-Back Kernel (#21631) --- benchmark/hicache/bench_hicache_write_back.py | 391 ++++++++++++++++++ .../csrc/{ => kvcacheio}/hicache.cuh | 2 + .../jit_kernel/csrc/kvcacheio/relayout.cuh | 90 ++++ .../csrc/kvcacheio/staged_write_back.cuh | 328 +++++++++++++++ python/sglang/jit_kernel/hicache.py | 104 ++++- .../sglang/srt/managers/cache_controller.py | 12 +- .../hybrid_cache/hybrid_cache_controller.py | 14 +- .../sglang/srt/mem_cache/memory_pool_host.py | 83 +++- python/sglang/srt/server_args.py | 46 +-- test/registered/jit/test_hicache.py | 193 ++++++++- .../unit/mem_cache/test_hiradix_cache_unit.py | 1 + .../unit/server_args/test_server_args.py | 24 +- 12 files changed, 1207 insertions(+), 81 deletions(-) create mode 100644 benchmark/hicache/bench_hicache_write_back.py rename python/sglang/jit_kernel/csrc/{ => kvcacheio}/hicache.cuh (99%) create mode 100644 python/sglang/jit_kernel/csrc/kvcacheio/relayout.cuh create mode 100644 python/sglang/jit_kernel/csrc/kvcacheio/staged_write_back.cuh diff --git a/benchmark/hicache/bench_hicache_write_back.py b/benchmark/hicache/bench_hicache_write_back.py new file mode 100644 index 000000000..6dd6e7d76 --- /dev/null +++ b/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() diff --git a/python/sglang/jit_kernel/csrc/hicache.cuh b/python/sglang/jit_kernel/csrc/kvcacheio/hicache.cuh similarity index 99% rename from python/sglang/jit_kernel/csrc/hicache.cuh rename to python/sglang/jit_kernel/csrc/kvcacheio/hicache.cuh index ae2970611..7abb5ab19 100644 --- a/python/sglang/jit_kernel/csrc/hicache.cuh +++ b/python/sglang/jit_kernel/csrc/kvcacheio/hicache.cuh @@ -1,3 +1,5 @@ +#pragma once + #include #include diff --git a/python/sglang/jit_kernel/csrc/kvcacheio/relayout.cuh b/python/sglang/jit_kernel/csrc/kvcacheio/relayout.cuh new file mode 100644 index 000000000..888293f5c --- /dev/null +++ b/python/sglang/jit_kernel/csrc/kvcacheio/relayout.cuh @@ -0,0 +1,90 @@ +#pragma once + +#include "hicache.cuh" +#include + +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 +__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(k_ptr_src); + const auto v_ptr_src_arr = static_cast(v_ptr_src); + const auto tid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const auto stride = static_cast(gridDim.x) * blockDim.x; + const auto total_vecs = static_cast(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(linear_vec_id / (static_cast(page_size) * num_layers * kVecPerItem)); + const auto page_vec_id = + static_cast(linear_vec_id % (static_cast(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(static_cast(indices_src)[page_id]); + const auto src_token = src_page + token_in_page; + const auto src_k = pointer::offset( + static_cast(k_ptr_src_arr[layer_id]), + static_cast(src_token) * kElementSize + static_cast(vec_id) * kVecBytes); + const auto dst_k = + pointer::offset(static_cast(k_cache_dst), static_cast(linear_vec_id) * kVecBytes); + const auto vec_k = details::load_nc(reinterpret_cast(src_k)); + details::store_nc(reinterpret_cast(dst_k), vec_k); + + if constexpr (!kIsMLA) { + const auto src_v = pointer::offset( + static_cast(v_ptr_src_arr[layer_id]), + static_cast(src_token) * kElementSize + static_cast(vec_id) * kVecBytes); + const auto dst_v = + pointer::offset(static_cast(v_cache_dst), static_cast(linear_vec_id) * kVecBytes); + const auto vec_v = details::load_nc(reinterpret_cast(src_v)); + details::store_nc(reinterpret_cast(dst_v), vec_v); + } + } +} + +template +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(num_pages) * page_size * num_layers * kVecPerItem; + const auto kernel = use_int32 ? hicache_relayout_kernel + : hicache_relayout_kernel; + if (total_vecs == 0) { + return; + } + + const auto grid = div_ceil(total_vecs, static_cast(kRelayoutBlockSize)); + RuntimeCheck( + grid <= std::numeric_limits::max(), "HiCache staged relayout: CUDA grid size exceeds uint32 range"); + LaunchKernel(static_cast(grid), kRelayoutBlockSize, device)(kernel, params); +} + +} // namespace diff --git a/python/sglang/jit_kernel/csrc/kvcacheio/staged_write_back.cuh b/python/sglang/jit_kernel/csrc/kvcacheio/staged_write_back.cuh new file mode 100644 index 000000000..0e0a08715 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/kvcacheio/staged_write_back.cuh @@ -0,0 +1,328 @@ +#pragma once + +#include "hicache.cuh" +#include "relayout.cuh" +#include +#include +#include + +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(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::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& src_ptrs, + std::vector 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(page_size * src_stride0 * elem_size); + const size_t dst_page_bytes = static_cast(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(src_ptrs[tensor_id].data_ptr()) + + static_cast(page_offset * page_size * src_stride0 * elem_size); + char* dst_ptr = static_cast(dst_ptrs[tensor_id].data_ptr()) + + static_cast(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& src_ptrs, + std::vector 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 batch_srcs; + thread_local std::vector batch_dsts; + thread_local std::vector 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(src_ptrs.size()) * static_cast(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(page_size * src_stride0 * elem_size); + const size_t dst_page_bytes = static_cast(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(src_ptrs[tensor_id].data_ptr()) + + static_cast(page_offset * page_size * src_stride0 * elem_size); + char* dst_ptr = static_cast(dst_ptrs[tensor_id].data_ptr()) + + static_cast(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 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 +struct HiCacheStagedWriteBackKernel { + private: + template + 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(device_) + .verify(staging_k); + if constexpr (!kIsMLA) { + TensorMatcher({T, N, D}) // + .with_dtype(cache_dtype) + .with_device(device_) + .verify(staging_v); + } + TensorMatcher({-1, N, D}) // + .with_dtype(cache_dtype) + .with_device() + .verify(k_cache_dst); + if constexpr (!kIsMLA) { + TensorMatcher({-1, N, D}) // + .with_dtype(cache_dtype) + .with_device() + .verify(v_cache_dst); + } + TensorMatcher({N}) // + .with_dtype() + .with_device(device_) + .verify(k_ptr_src); + if constexpr (!kIsMLA) { + TensorMatcher({N}) // + .with_dtype() + .with_device(device_) + .verify(v_ptr_src); + } + TensorMatcher({P}) // + .with_dtype(indices_dtype) + .with_device(device_) + .verify(page_indices_src); + TensorMatcher({T}) // + .with_dtype(dst_indices_dtype) + .with_device() + .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(P.unwrap()), + .num_layers = static_cast(N.unwrap()), + .page_size = static_cast(page_size), + }; + const auto device = device_.unwrap(); + const auto use_int32 = indices_dtype.unwrap().bits == 32; + launch_hicache_relayout_kernel(params, P.unwrap(), N.unwrap(), page_size, use_int32, device); + + auto stream = LaunchKernel::resolve_device(device); + const int64_t* dst_indices_ptr = static_cast(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( + 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( + cache_dst, cache_dst, dst_indices_cpu, staging, staging, page_indices_src, ptr_src, ptr_src, page_size); + } +}; + +} // namespace diff --git a/python/sglang/jit_kernel/hicache.py b/python/sglang/jit_kernel/hicache.py index 0e5ed5802..b4b80bdb7 100644 --- a/python/sglang/jit_kernel/hicache.py +++ b/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, + ) diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 787b8d75e..0a7f2f297 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/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() diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index e69160368..6a3c617e2 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/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() diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 1e6596011..2f7e551a4 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/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( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 3235fa052..e4edefa98 100644 --- a/python/sglang/srt/server_args.py +++ b/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" diff --git a/test/registered/jit/test_hicache.py b/test/registered/jit/test_hicache.py index 527aeefe7..0d5ae5ba4 100644 --- a/test/registered/jit/test_hicache.py +++ b/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"])) diff --git a/test/registered/unit/mem_cache/test_hiradix_cache_unit.py b/test/registered/unit/mem_cache/test_hiradix_cache_unit.py index e7cd78329..a83c59dfd 100644 --- a/test/registered/unit/mem_cache/test_hiradix_cache_unit.py +++ b/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) diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index b4db117b3..abd4b9065 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/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):