Skip to content

[Perf][KV Offload] DFlash2: per-group CPU pool and disk pages when block spans differ - #46

Draft
Leonccaa wants to merge 2 commits into
mainfrom
f/0.1.5-a3-dflash2-cpu-partition
Draft

Leonccaa wants to merge 2 commits into
mainfrom
f/0.1.5-a3-dflash2-cpu-partition

Conversation

@Leonccaa

@Leonccaa Leonccaa commented Oct 27, 2026 •

Copy link
Copy Markdown
Owner

Draft: CPU work is complete; GPU validation (readback bitwise and hit counts with the CPU pool and the NVMe tier, real per-group page sizes and slot counts, and the MTP4 no-change check) is planned for the next GPU window. Not for merge before that.

Change

With DFlash2, the drafter cache group uses 832-token blocks while the target groups use 1,664. The grouped CPU offload layout, where each cache group has a private pool that holds only its own pages, required equal block sizes, so DFlash2 fell back to the shared-row layout: every CPU slot and every disk page carried one full row of all KV tensors (25,559,040 bytes across four workers), and the recurrent and drafter groups each paid that full row. This is why DFlash2 needs about 2.4 times MTP4's bytes per token in the CPU pool and on disk (docs/kv-offload.md).

The grouped layout now applies whenever there is a recurrent group and more than one prefix-cacheable group, also when block spans differ. Slot counts are solved per group: a group whose offloaded block covers fewer tokens gets proportionally more slots (the drafter gets twice the target's), so every group covers the same number of tokens within the same byte budget. Recurrent groups keep the retention-based state slot count, now computed from each group's own spec; recurrent groups with different geometry fail at startup with a clear error instead of being sized from the last spec seen.

With equal spans the sizing is exactly the previous grouped sizing, so MTP4 and other equal-span models are unchanged, including their disk-tier namespace. Unequal spans use a new persistent layout version, so DFlash2 disk files written by 0.1.4 (shared rows) are not reused after the upgrade; the first requests after the upgrade compute cold and store again. "cpu_group_partition": "off" in kv_connector_extra_config restores the shared-row layout as a rollback switch ("auto" is the default). Startup logs one line per group with page size, slot count and token span.

Expected effect (estimate from the group layout, to be measured): DFlash2 CPU and disk bytes per token fall by about 1.8–2.3 times, since the target groups already fill most of the row. docs/kv-offload.md describes the per-group pools and says the 0.1.4 DFlash2 figures are being re-measured; the numbers will be updated from the GPU run before this is merged.

Validation

CPU, in the ReVolta build image: pytest tests/v1/kv_offload/ tests/v1/core/test_qwen4_exp_kv_cache.py tests/v1/core/test_qwen4_exp_dflash_kv_cache.py tests/v1/attention/test_kv_storage_contract.py tests/v1/kv_connector/unit/offloading_connector/ → 395 passed, 24 failed, 6 skipped; the 24 failures are CUDA-only transfer tests and are exactly the same set (by node id) as on the base. New tests cover the gate cases, DFlash2-like group pages and balanced token capacity (the drafter gets twice the target's slots and every group covers the same tokens; the solve is tight against the budget), equal-span sizing identical to a copy of the previous code including the measured MTP4 layout, per-group region geometry, disk page sizes per group, the MTP4 persistent layout unchanged, a non-grouped tier unchanged, and the file mapper namespace. The existing DFlash2 offload-budget test now compares real per-group maps instead of two empty dicts. pre-commit and mypy 3.12 pass on the changed files.

Not yet run on GPU: real DFlash2 page sizes and slot counts, readback bitwise and hit counts from the CPU pool and the NVMe tier, restart reuse, and the MTP4 no-change check.

Attribution

The design, the implementation and its tests are the work of a local model: Qwen3.8-Flash-Next (Uncensored AWQ g32) served by ReVolta 0.1.4 on our 4 × V100 server, working through the pi coding agent. It worked out the group page sizes, the per-group slot solve and the disk namespace in a read-only design pass, then implemented phase 1 and wrote and ran the tests. Claude Code reviewed it, fixed one mypy error and opened this PR. No upstream code. Commits are signed off under the DCO.

🤖 Generated with Claude Code

Leonccaa and others added 2 commits October 27, 2026 16:49
…spans differ

DFlash2's drafter cache group uses 832-token blocks against the target's 1,664, so the grouped CPU layout (one private pool per cache group, holding only that group's pages) was skipped and every CPU slot and disk page carried a full shared row of all KV tensors. The gate now partitions whenever there is a recurrent group and more than one prefix-cacheable group, also with different block spans.

Slot counts are solved per group: a group whose offloaded block covers fewer tokens gets proportionally more slots, so every group covers the same number of tokens within the same byte budget. Recurrent groups keep the retention-based state slot count, now computed from each group's own spec; recurrent groups with different geometry are rejected at startup instead of being sized from the last spec seen. With equal spans the sizing is exactly the previous grouped sizing, so MTP4 and other equal-span models are unchanged, including the disk namespace; unequal spans use a new persistent layout version so their disk files never mix with shared-row files. `cpu_group_partition` in the connector's extra config ("auto" default, "off" restores the shared-row layout) is a rollback switch. Startup logs one line per group with page size, slot count and token span.

Tests (CPU): gate cases, DFlash2-like group pages and balanced token capacity, equal-span sizing identical to a copy of the previous code (including the measured MTP4 layout), per-group region geometry, disk page sizes per group, MTP4 persistent layout unchanged, non-grouped tier unchanged, and the file mapper namespace. Not yet validated on GPU.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com>
… loop

mypy: the loop reused the name of the KVCacheGroupSpec variable above.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant