[Feature] KV Cache NPU-to-SSD Offload Plan

[superxf]

Goal

Support KV cache offload for the serving path, with the final target being direct NPU memory <-> on-card SSD transfer. Before the NPU <-> SSD transfer API is ready, PR #31 implements a conservative CPU offload path as Phase 0 so the scheduler, block metadata, transfer jobs, and load-back flow can be validated first.

flowchart LR
    subgraph host["Host Side"]
        cpu["CPU / Host Memory"]
    end

    subgraph card["NPU Card Side"]
        npu["NPU KV Cache"]
        ssd["On-card SSD"]
        npu <--> ssd
    end

    npu -. "Phase 0 fallback" .-> cpu
    npu --> ssd

Loading

The main process owns request scheduling and KV block metadata. The worker process owns the real NPU KV tensors and executes transfer jobs. The scheduler therefore only creates load/store jobs and updates metadata after completion; the transfer backend performs the actual data movement.

Current PR Scope

PR #31 adds the CPU offload version of this flow:

• Split logical KV block identity from NPU physical page id.
• Add KV block residency states: NPU, CPU, MOVING_TO_CPU, MOVING_TO_NPU, and future SSD states.
• Add transfer job abstractions and a CPU transfer backend.
• Add a CLI switch through --max-cpu-offload-blocks; 0 disables CPU offload.
• Integrate load/store job handling into the serving worker and async engine.
• Keep the policy conservative: no background offload, no proactive eviction.

This PR is intended to validate the serving scheduler and state machine. Later, the CPU backend can be replaced by an NPU <-> SSD backend while keeping the same scheduler-facing contract.

Scheduling Strategy

The serving scheduler follows the existing run_prefill / run_decode path, not the L3 generate path.

Phase 1: Running Queue

1. Iterate over currently running requests.
2. Compute how many new tokens can be scheduled under the token budget.
3. Allocate extra KV blocks if the request needs more pages.
4. If allocation succeeds, resolve all logical block ids to resident NPU physical page ids and schedule the request.
5. If allocation fails, preempt the lowest-priority running request.

Preemption is the only trigger for offload in the current conservative design.

Preemption With CPU Offload Enabled

When --max-cpu-offload-blocks > 0, the scheduler tries to offload the victim request's KV blocks to CPU:

1. Pick only victim-owned resident NPU blocks.
2. Do not offload shared prefix-cache blocks, because other requests may still depend on them being resident on NPU.
3. Create an NPU -> CPU store job.
4. Mark the victim as PREEMPTED and move it to the front of the waiting queue.
5. The current scheduling round stops after the preemption, matching the conservative vLLM-style behavior.

If CPU offload capacity is exhausted, or no exclusive resident block can be offloaded, the scheduler falls back to recomputation:

1. Release the victim request's block references.
2. Clear cached_block_ids, allocated_block_ids, and num_computed_tokens.
3. Move the victim back to the waiting queue.

This keeps scheduling correct even when the offload medium is unavailable or full.

Phase 2: Waiting Queue

Waiting requests are scheduled only when the current round did not preempt a running request.

For each waiting request:

1. Try prefix-cache lookup if the request has no existing blocks.
2. Compute the token budget and required new KV blocks.
3. Check whether existing blocks are resident on NPU.
4. If some blocks are on CPU, create CPU -> NPU load jobs and defer this request to a later scheduling round.
5. Only after required old blocks are resident does the scheduler allocate new KV blocks.
6. Once all blocks are resident, resolve logical block ids to physical page ids and schedule prefill/decode.

This ordering avoids allocating new blocks before the scheduler knows whether load-back can reserve NPU pages.

Data Flow

Store On Preemption

1. Scheduler selects a preemption victim.
2. KvCacheManager builds a store job for exclusive resident blocks.
3. Blocks are marked MOVING_TO_CPU in Phase 0, later MOVING_TO_SSD for the final design.
4. Worker submits the transfer job.
5. On completion, metadata is updated to CPU or SSD, and the old NPU physical page can be reused.

Load On Resume

1. Scheduler sees a waiting request with non-resident blocks.
2. It reserves target NPU physical pages.
3. It creates CPU/SSD -> NPU load jobs.
4. The request is deferred until the job completes.
5. After completion, logical block ids resolve to the latest NPU physical page ids.
6. run_prefill / run_decode receives only resident physical page ids.

Final SSD Direction

The final target is to keep the same scheduler policy and replace the CPU transfer backend with direct NPU <-> on-card SSD transfer:

• CPU states become SSD states for production offload.
• CPULoadStoreSpec is replaced by SSDLoadStoreSpec.
• The worker-side backend uses the future NPU <-> SSD transfer API.
• Scheduler policy remains conservative at first: offload only on preemption caused by KV pressure.

Because the SSD is directly attached to the NPU, the final SSD path is expected to be faster than CPU offload. The CPU path is mainly a correctness and integration milestone.

Risks

• Reusing an NPU page while a store job is still in flight can corrupt offloaded KV data.
• Offloading shared prefix-cache blocks can break other requests that still rely on those blocks.
• If logical block id and physical page id diverge, block_table / slot_mapping can point to the wrong NPU page.
• CPU/SSD capacity exhaustion must not crash the scheduler; it should fall back to recomputation.
• Excessive proactive offload can cause IO thrashing, so background offload is intentionally not part of the current policy.
• Main-process metadata and worker-side transfer completion must stay synchronized.

Validation Plan

• Unit-test block residency transitions, load/store job rollback, and logical-to-physical page resolution.
• Unit-test preemption fallback when CPU offload capacity is full.
• Unit-test that shared prefix-cache blocks are not offloaded by a victim preemption.
• Run serving-path NPU smoke tests for prefill/decode with offload enabled.
• Run an accuracy comparison with and without offload for the same prompt and sampling configuration.