perf(kvflash): async page-out D2H and page-in H2D via pinned copy stream

[cheese-cakee]

Summary

KvFlashPager currently blocks the decode thread for every page-out (D2H, 1.26 ms p50) and page-in (H2D, 0.63 ms p50) via synchronous ggml_backend_tensor_get/set. This PR replaces those with cudaMemcpyAsync on a dedicated page_stream_, making page-out fire-and-forget and page-in paid at most once per alloc_span() call.

Changes

• server/src/common/kvflash_pager.h
  • Include gpu_runtime_compat.h when DFLASH27B_BACKEND_CUDA, DFLASH27B_BACKEND_HIP, or GGML_USE_HIP is set; define KVFLASH_HAS_ASYNC_DMA
  • ChunkState::host_data: std::vector<uint8_t> → void * (cudaMallocHost-pinned) under KVFLASH_HAS_ASYNC_DMA
  • Add cudaStream_t page_stream_ and bool has_pending_page_in_ to KvFlashPager
  • Add destructor + private cleanup_() that syncs and frees pinned buffers then destroys the stream
  • attach(): calls cleanup_() before re-init, creates stream with cudaStreamCreate
  • reset(): syncs stream + frees all pinned host_data before clearing chunks_
  • zero_free_blocks(): syncs stream after async memsets
  • copy_chunk(): cudaMemcpyAsync(DeviceToHost / HostToDevice, page_stream_) replaces ggml_backend_tensor_get/set
  • zero_block(): cudaMemsetAsync(..., page_stream_) replaces ggml_backend_tensor_set with zero buffer
  • alloc_span() / reselect(): cudaStreamSynchronize(page_stream_) only when has_pending_page_in_ is set; page-outs left async

How it works

All transfers share page_stream_, giving this ordering guarantee within each alloc_span() / reselect() call:

D2H (evict chunk C to host)
  → cudaMemsetAsync (zero freed GPU slot)
    → H2D (recall chunk C' from host)

No cross-stream events needed. The CPU thread is unblocked after issuing the transfers; it syncs only if a page-in was queued (H2D must land before the attention forward reads the recalled rows).

Performance / Numbers

Measured baselines from server/DESIGN.md and test_kvflash profile (E):

Operation	Before (synchronous)	After (async)
page_out D2H p50	1.26 ms — blocks decode thread	queued async; ~0 ms on critical path
page_in H2D p50	0.63 ms — blocks decode thread	one cudaStreamSynchronize per step with recall
N evictions + M recalls	sum of individual transfers	single sync at end covers all

Numbers not yet validated on RTX 3090 (no local GPU). Benchmark via test_kvflash --skip-profile section E, or by instrumenting alloc_span() with std::chrono.

Limitations / Known follow-up

• No cross-stream sync between page_stream_ and ggml's compute stream: relies on ggml_backend_graph_compute already synchronising before alloc_span() is called, which holds for the current single-threaded inference loop. If a future caller issues an async graph compute and calls alloc_span() concurrently the ordering must be revisited.
• Page-out D2H is fire-and-forget; if reset() is called before the DMA completes, reset() syncs page_stream_ first — safe, but adds latency at reset time.
• cudaStreamCreate return value is not checked; a failed create leaves page_stream_ = nullptr, in which case cudaMemcpyAsync/cudaMemsetAsync fall back to the default (synchronous) stream — correct but unintentionally blocking.
• CUDA graph for KVFlash spec-verify path will need cudaEvent_t page_done_event_ so the compute graph can cudaStreamWaitEvent instead of a hard sync; that is left for the graph PR.

Verification / Test plan

• Build: cmake -B server/build -S server -DCMAKE_BUILD_TYPE=Release && cmake --build server/build --target test_kvflash -j
• Correctness: server/build/test_kvflash <qwen35.gguf> — suites A–D (relocation, paging, reselect/recall) must pass bit-exact vs baseline
• Performance: suite E decode ms/step and page-event microbenchmark at pool=1K/4K / logical=128K
• HIP: gfx1151 ROCm build (same gpu_runtime_compat.h shim used by qwen35_dflash_target.cpp since qwen35: include gpu_runtime_compat.h for HIP builds #395)

[cubic-dev-ai]

2 issues found across 1 file

<sub>Reply with feedback, questions, or to request a fix.

Re-trigger cubic</sub>

[howard0su]

any side impact to introduce this option? If this is always better, maybe we can default on this.

[cheese-cakee]

Small clarification first — KVFLASH_HAS_ASYNC_DMA isn't a flag you can toggle, it's a compile-time #ifdef that auto-enables on every CUDA/HIP build:

#if defined(DFLASH27B_BACKEND_CUDA) || defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP)
#include "gpu_runtime_compat.h"
#define KVFLASH_HAS_ASYNC_DMA 1
#endif

The sync path is just the CPU fallback. So it's already "on" everywhere a GPU exists.

On side impacts, the main one is pinned host memory. cudaMallocHost allocations can't be swapped and compete with CUDA's own pinned pool, so over-allocating hurts the whole system. The ordering assumption (sync before alloc_span) also only holds because the current compute loop is single-threaded, would need cross-stream events if that changes.