KV offload with a disk tier: concurrent restores can miss the whole prefix when the CPU pool has few recurrent-state slots #62

Summary

With TieringOffloadingSpec and a disk (fs) secondary tier, a request whose prefix is only on disk can get no prefix hit at all, even though every page of that prefix is on disk. This happens when more requests restore from disk at the same time than the CPU pool has free recurrent-state slots. The request is then computed from scratch. Output is not affected (it is a normal cold prefill), but the disk tier gives no benefit for that request.

It affects the grouped CPU pool layout, which hybrid models with a recurrent (GDN) group use when all prefix-cacheable groups have the same block span. This includes the default MTP4 configuration. In released versions DFlash2 uses the shared layout (its draft and target block spans differ) and is not affected. CPU-only offload (no secondary tier) is not affected.

Trigger

All of these must be true:

• TieringOffloadingSpec with at least one secondary tier (for example fs).
• Several requests reach the scheduler at about the same time, and their prefixes are on disk but not in the CPU pool (for example after a restart, or after CPU eviction).
• The number of such requests is greater than the number of requests whose recurrent states the CPU pool can hold. Each recurrent group gets its own pool, sized by mamba_state_slots from mamba_state_slots_reference_tokens (default in the docs: 65536). With MTP4 this is about 0.73 requests per GiB of cpu_bytes_to_use. We measured 4 state slots per group (2 requests) for a 2 GiB pool and 72 (36 requests) for a 48 GiB pool.

The disk tier example in docs/kv-offload.md uses a 2 GiB CPU pool. With it, 3 requests that restore from disk at the same time are enough for one of them to miss. With pools below about 11 GiB, 8 concurrent restores (--max-num-seqs 8) can trigger it. The 48 GiB quick-start pool does not trigger it at --max-num-seqs 8.

Cause

A block found only on the secondary tier is promoted to the CPU pool first. TieringOffloadingManager._initiate_promotion (vllm/v1/kv_offload/tiering/manager.py) reserves a CPU slot for each promoted block, and the slot stays reserved until the read finishes. When a recurrent group's pool has no free slot because other requests' promotions hold them, prepare_write returns None and lookup returns False ("not available"). In the connector (OffloadingConnectorScheduler._lookup in vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py), a recurrent group is looked up as a one-block sliding window. When none of its state keys is available, the lookup returns 0 tokens at once, even if the lookups of the other groups are still deferred. The request is then scheduled as a full prefill. Without the recurrent state at a block boundary, none of the prefix can be reused.

On the next scheduler step the same lookup would succeed, because finished promotions free their slots. The miss is decided on the first lookup only.

The small state pool comes from the sizing in mamba_state_slots (vllm/v1/kv_offload/cpu/spec.py). It assumes each request is mamba_state_slots_reference_tokens long and needs about 2 states. The number of requests that can be restored at the same time is not part of the sizing.

Reproduction

We reproduced this on CPU with a grouped tiering manager: one token group with enough slots, one recurrent group with 3 slots, and 4 documents written to the fs tier. A new process looks up all 4 documents in the same step. The first 3 state lookups return None (promotion started). The 4th returns False, although its file is on disk. With 4 state slots, all 4 return None. We also saw it on GPU (4 × V100, TP4 + DCP2) in a restart read-back test of a development branch that enables the same grouped layout for DFlash2 (#46, not merged): of 4 documents restored at the same time, one got 0 cached tokens, while all of its 71 page files were on disk. We have not yet run the same GPU test with MTP4 on a released version.

Workarounds

• Use a larger cpu_bytes_to_use. As a rule of thumb, keep 0.73 × pool GiB above the number of requests that may restore from disk at the same time.
• Or lower mamba_state_slots_reference_tokens (for example 16384). By the sizing formula this should give each recurrent group about 4 times more state slots from the same budget (not measured), at the cost of fewer token slots.
• Fewer concurrent restores (for example a lower --max-num-seqs) also avoids it.

Possible fixes

1. When the primary tier is full only because of in-flight promotions, return None (wait) from lookup instead of False.
2. Give each recurrent-group pool at least max_num_seqs × states per request slots.
3. When the states for the full prefix cannot be promoted, fall back to the longest prefix whose states fit.