feat(serve): enable prefix caching with DFlash (on by default)

Combining --enable-prefix-caching with the DFlash drafter crashed engine
init with "block_size must be divisible by hash_block_size" in
HybridKVCacheCoordinator. The drafter's attention KV page is ~2x the
target's, so vLLM's page-size unification scales the target's mamba+attn
block 2240->4480 to match; the (align-mode) mamba block then differs from
cache_config.block_size and resolve_kv_cache_block_sizes backs off to
hash_block_size = LCM (4480), which the drafter group (still 2240) is not
divisible by. But the GCD (2240) divides every group and is the correct
finer hash granularity; the back-off's `block_size != cache_block` test is
a buggy proxy for "non-align mamba".

patch_prefix_align.py makes the back-off align-aware (only back off when
mamba_cache_mode != "align"), so resolve uses the GCD. Prefix caching is on
by default (a win for multi-turn / long-context, neutral for single-turn
c=1); set PREFIX_CACHE=0 to disable.

Validated on GB10: READY, DFlash accept ~7.7 tok/step on code, ~13x
warm-prefix TTFT (2.30s->0.18s), KV pool ~422k tokens (no regression).
See docs/FINDINGS.md section 1a.
This commit is contained in:
Entrpi
2026-06-28 21:29:40 +10:00
parent 3f05a87cdf
commit 41d92cb6f9
4 changed files with 152 additions and 4 deletions
+39 -1
View File
@@ -32,7 +32,9 @@ the drafter's attention spec:
makes the mamba group block ≠ cache block, tripping the coordinator hash
assert. Omit it.
- `--no-enable-prefix-caching` routes to `KVCacheCoordinatorNoPrefixCache`, which
has no line-504 hash assert. Prefix caching is irrelevant at c=1 anyway.
has no hash assert. It is the default; prefix caching is irrelevant at c=1
single-turn but a large win for agentic multi-turn / long-context re-reads —
see §1a for enabling it *with* DFlash.
Working stack: `patch_unify2` + prefix-off + INT4 (bf16 KV) target + drafter
pinned to `FLASH_ATTN`. **No FA4 shim needed** — vLLM gates FA4 to cap families
@@ -40,6 +42,42 @@ pinned to `FLASH_ATTN`. **No FA4 shim needed** — vLLM gates FA4 to cap familie
SGLang-only; SGLang's DFlash works too but its sm121 base decode is ~2× slower
than vLLM's, so it loses on absolute throughput.)
## 1a. Prefix caching + DFlash together (`PREFIX_CACHE=1`)
Enabling `--enable-prefix-caching` *with* the DFlash drafter crashed engine init with
`AssertionError: block_size must be divisible by hash_block_size`
(`HybridKVCacheCoordinator.__init__`). Prefix caching alone (no DFlash) and DFlash alone
both work; only the combination broke. Diagnosed on the GB10 (2026-06-28):
- The DFlash drafter's attention layers carry a **2× larger KV page** (9 175 040 vs the
target's 4 587 520 B — the drafter has ~2× the KV heads). So the drafter is `max_page`.
- `unify_kv_cache_spec_page_size` therefore **scales the *target's* mamba + attention block
2240 → 4480** (ratio 2) to match the drafter page; the drafter group stays at **2240**.
- In `resolve_kv_cache_block_sizes` the (align-mode) mamba block is now 4480 ≠
`cache_config.block_size` (2240), so its **back-off branch** fires and forces
`hash_block_size = LCM = 4480`. The drafter group is 2240, and `2240 % 4480 ≠ 0` → assert.
- But `GCD = 2240` divides **every** group (4480 and 2240) and is the correct finer hash
granularity (vLLM's `hash_block_size < block_size` merge-up, #29143). The back-off only
exists to disable fine hashing for *non-align* mamba — its `block_size != cache_block`
test is a **buggy proxy** that also trips when an align-mode block was merely scaled up.
**Fix** ([`patch_prefix_align.py`](../runtime/patch_prefix_align.py)): make the back-off
align-aware — only back off when `mamba_cache_mode != "align"`. In align mode it falls
through to the GCD path. One-condition change; no drafter-geometry surgery, no extra memory.
(vLLM #45181's pad-don't-scale path does *not* apply here: our pages are an exact 2×, so
unify *scales* rather than pads — #45181 only adds a branch for the non-divisible case.)
The KV reshape already reads padded/strided pages correctly in this image
(`get_kv_cache_block_dim``physical_block_dim`), so there is **no acceptance risk** — the
old 1.47-accept "pad" regression was a *different*, since-fixed reshape bug, not this path.
Validated (prefix-ON + DFlash n=12, flash_attn): **READY**; **DFlash accept ~7.7 tok/step**
on code (240 accepted / 36 drafts; per-position 34→11, unchanged from prefix-off);
**warm-prefix TTFT 2.30 s → 0.18 s (~13×)** on a 4k shared prefix with a real
`prefix_cache_hits_total` bump; coherent output; **KV pool 421 888 tokens** (≈ the prefix-off
~427k, no regression). **On by default** (win for multi-turn / long-context, neutral for
single-turn c=1); set `PREFIX_CACHE=0` to disable.
## 2. DFlash vs MTP — acceptance is task-dependent
MTP-2 (the native head) drafts 2 tokens **sequentially**, so acceptance caps at