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)

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

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