Skip to content

perf(connector): explore layerwise async KV cache load/store #18

Description

@GentleCold

Motivation

Currently start_load_kv uses full synchronous loading: submit all layer GDS reads at once and block until completion. This avoids the risk of wait_for_layer_load hooks not being called under vLLM FULL CUDA graph mode, but introduces two problems:

  1. High memory peak: all layer staging buffers must exist simultaneously
  2. Long blocking time: the decode step must wait for all layers to load before computation starts

If a layerwise async pipeline can be implemented (submit GDS read per layer → copy to KV cache per layer), it could overlap IO with subsequent layer loading, reducing both blocking time and memory footprint.

Proposal

On top of Issue #17's performance diagnosis, explore the feasibility of layerwise async load/store:

  1. Load path: change start_load_kv from "full blocking" to "per-layer submit + intra-layer sync" — each GDS read completes and immediately copies to the corresponding layer's KV cache slot, reducing staging buffer lifetime
  2. Store path: try per-layer GDS write submission during save_kv (if kvikio supports it), compare latency against full submission
  3. Compatibility check: confirm whether layerwise loading breaks deterministic CUDA graph execution; if not feasible, at least evaluate gains in non-graph mode
  4. Metrics: measure layerwise vs full-load p50/p99 latency and peak memory difference

Out of scope

  • Modify vLLM CUDA graph behavior to support wait_for_layer_load (read-only constraint)
  • Introduce new thread/process models (reuse existing daser-io thread)

Metadata

Metadata

Assignees

Labels

enhancementNew feature or request

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions