mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-09 22:59:44 +08:00
88e70d1975
The open-source engine has no KV reuse across requests, so every turn of a
multi-turn conversation prefills its whole prompt. This adds deepinfer/hicache:
- Pinned host pools organized in pages (one page = one 128-token KV block)
plus slots for the SWA state that the production KV store persists
(select_state_slice_to_store / load_state_from_store_to_prefill_cache).
Pages are indexed by a hash chain over the prompt, LRU evicted, written
through at the end of every prefill (engine hook pe_do_save_kvcache).
- Transfers use one small CUDA kernel copying pages between GPU and pinned
host memory through UVA pointers on the engine stream ("kernel" I/O
method, ~180 GiB/s on GB300); completion is tracked with CUDA events.
- A hit is capped so that at least yoco keep_tokens remain to recompute
(bounded replay rebuilds the decoder SWA state from that suffix), then
snapped to the nearest 128-aligned position with a stored state; states
are stored at the last few boundaries of a prompt so short follow-ups hit.
The sequence is pre-allocated, filled from host and scheduled with
num_cached_tokens set (HostPrefixCache.try_resume).
- Engine: read-only Sequence.block_table binding (to copy a finished
prefill's blocks out) and its stub.
Validated on 4x GB300 with DeepSeek-V4.1-Flash through an OpenAI-compatible
server (separate PR): greedy outputs byte-identical with and without the
cache over 8 multi-turn requests; InferenceX AgentX replay at concurrency 1
served 224 requests per hour vs 99, TTFT p50 0.59 s vs 8.9 s, with 89-96%
of prompt tokens served from the cache. tests/hicache/test_host_cache.py
covers the cache logic against a fake engine.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TNVW4sqoaPqbNtSrgpA2zK
46 lines
2.4 KiB
Plaintext
46 lines
2.4 KiB
Plaintext
// Page copy between two pools addressed by UVA pointers: either pool may live in GPU memory or in
|
|
// pinned host memory (the "kernel" transfer method of SGLang's HiCache). One thread block (row of the
|
|
// grid) per page; pages are copied as 16-byte vectors on the current CUDA stream.
|
|
#include <torch/extension.h>
|
|
#include <c10/cuda/CUDAException.h>
|
|
#include <c10/cuda/CUDAStream.h>
|
|
#include <cuda_runtime.h>
|
|
|
|
#include <algorithm>
|
|
#include <cstdint>
|
|
|
|
namespace {
|
|
|
|
__global__ void copy_pages_kernel(const uint8_t *__restrict__ src_base, uint8_t *__restrict__ dst_base,
|
|
const int32_t *__restrict__ src_ids, const int32_t *__restrict__ dst_ids, int64_t page_bytes) {
|
|
const int64_t n16 = page_bytes >> 4;
|
|
const int4 *src = reinterpret_cast<const int4 *>(src_base + static_cast<int64_t>(src_ids[blockIdx.x]) * page_bytes);
|
|
int4 *dst = reinterpret_cast<int4 *>(dst_base + static_cast<int64_t>(dst_ids[blockIdx.x]) * page_bytes);
|
|
const int64_t stride = static_cast<int64_t>(gridDim.y) * blockDim.x;
|
|
for (int64_t i = static_cast<int64_t>(blockIdx.y) * blockDim.x + threadIdx.x; i < n16; i += stride) {
|
|
dst[i] = src[i];
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
void copy_pages(int64_t src_base, int64_t dst_base, torch::Tensor src_ids, torch::Tensor dst_ids, int64_t page_bytes) {
|
|
TORCH_CHECK(page_bytes > 0 && page_bytes % 16 == 0, "page_bytes must be a positive multiple of 16, got ", page_bytes);
|
|
TORCH_CHECK(src_ids.is_cuda() && dst_ids.is_cuda(), "page ids must be CUDA tensors");
|
|
TORCH_CHECK(src_ids.scalar_type() == torch::kInt32 && dst_ids.scalar_type() == torch::kInt32, "page ids must be int32");
|
|
TORCH_CHECK(src_ids.is_contiguous() && dst_ids.is_contiguous(), "page ids must be contiguous");
|
|
TORCH_CHECK(src_ids.numel() == dst_ids.numel(), "src/dst page id count mismatch");
|
|
const int64_t n = src_ids.numel();
|
|
if (n == 0) {
|
|
return;
|
|
}
|
|
constexpr int threads = 256;
|
|
const int64_t n16 = page_bytes / 16;
|
|
const unsigned blocks_y = static_cast<unsigned>(std::min<int64_t>((n16 + threads - 1) / threads, 16));
|
|
const dim3 grid(static_cast<unsigned>(n), blocks_y);
|
|
auto stream = c10::cuda::getCurrentCUDAStream();
|
|
copy_pages_kernel<<<grid, threads, 0, stream.stream()>>>(reinterpret_cast<const uint8_t *>(src_base),
|
|
reinterpret_cast<uint8_t *>(dst_base), src_ids.data_ptr<int32_t>(), dst_ids.data_ptr<int32_t>(), page_bytes);
|
|
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
|
}
|