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.

README.md

@@ -290,14 +290,15 @@ which no dense-weight quantization touches. For base / low-acceptance serving
 (`dense`), the stack is a real +28 %. Full derivation in
 [`docs/FINDINGS.md`](docs/FINDINGS.md).
 
-## Under the hood: the four runtime patches
+## Under the hood: the five runtime patches
 
 vLLM is unmodified on disk; [`runtime/serve.sh`](runtime/serve.sh) edits the
 installed package in-place before `vllm serve` (idempotent, sentinel-guarded):
 
 | Patch | Effect | Required by |
 |---|---|---|
-| [`patch_unify2.py`](runtime/patch_unify2.py) | scale-block KV-cache **unify** so the hybrid GDN+mamba target absorbs the drafter's attention spec (the upstream assert cannot); paired with `--no-enable-prefix-caching` to route to the no-hash-assert coordinator | DFlash (any spec profile) |
+| [`patch_unify2.py`](runtime/patch_unify2.py) | scale-block KV-cache **unify** so the hybrid GDN+mamba target absorbs the drafter's attention spec (the upstream assert cannot) | DFlash (any spec profile) |
+| [`patch_prefix_align.py`](runtime/patch_prefix_align.py) | makes `resolve_kv_cache_block_sizes`' mamba back-off **align-aware** so prefix caching coexists with DFlash (uses the GCD hash block size, 2240, instead of the LCM that fails the coordinator assert) | prefix caching ON (default) |
 | [`patch_inc_hybrid.py`](runtime/patch_inc_hybrid.py) | adds an `INCConfig.maybe_update_config` override that detects FP8 dense layers in the hybrid checkpoint and dispatches `Fp8LinearMethod` for `shared_expert` | `dense` |
 | [`patch_int8_lmhead_v3.py`](runtime/patch_int8_lmhead_v3.py) | replaces the lm-head matmul in `_get_logits` with a batched int8 w8a16 Triton GEMV (keeps the bf16 weight for the shared drafter) | `dense` |
 | [`patch_fla_shmem.py`](runtime/patch_fla_shmem.py) | allows the FLA GDN chunk kernels to use large tiles on sm121's 99 KiB shmem (prefill / TTFT only; harmless) | always (free) |

docs/FINDINGS.md

@@ -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

runtime/patch_prefix_align.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python3
+"""spark-prefix-align: let vLLM prefix caching coexist with the DFlash drafter on the
+hybrid GDN+mamba+MoE target (DGX Spark, AEON 0.23.0+aeon.sm121a.dflash).
+
+Root cause (validated on GB10, 2026-06-28): the DFlash drafter's attention layers carry
+a ~2x larger KV page than the target, so page-size unification scales the TARGET's mamba
++ attention blocks UP (2240 -> 4480) to match the drafter's page, while the drafter group
+stays at 2240. `resolve_kv_cache_block_sizes` then sees a MambaSpec whose block_size
+(4480) != cache_config.block_size (2240) and takes its back-off branch, forcing
+hash_block_size = LCM = 4480. The drafter group (block 2240) is not divisible by 4480, so
+HybridKVCacheCoordinator.__init__ aborts with "block_size must be divisible by
+hash_block_size".
+
+But the GCD (2240) divides EVERY group (4480 % 2240 == 0, 2240 % 2240 == 0) and is the
+correct finest-common hash granularity. The back-off only exists to disable fine hashing
+for *non-align* mamba (its own comment says "mamba_cache_mode != align"); the
+`block_size != cache_config.block_size` test is a buggy proxy that also fires when an
+ALIGN-mode mamba block was merely scaled up by unification. This patch fixes the proxy:
+back off only when mamba is genuinely non-align. In align mode we fall through to the GCD
+path, which is exactly what hash_block_size (finer than block_size, merged up per group)
+was designed for (vLLM #29143).
+"""
+import pathlib, sys
+
+P = pathlib.Path("/usr/local/lib/python3.12/site-packages/vllm/v1/core/kv_cache_utils.py")
+src = P.read_text()
+
+if "spark-prefix-align" in src:
+    print("[patch_prefix_align] already patched", flush=True)
+    sys.exit(0)
+
+OLD = '''    # Mamba groups with block_size != cache_config.block_size
+    # (mamba_cache_mode != "align") break divisibility; back off to the
+    # scheduler block size.
+    if any(
+        isinstance(g.kv_cache_spec, MambaSpec)
+        and g.kv_cache_spec.block_size != cache_config.block_size
+        for g in groups
+    ):
+        return scheduler_block_size, scheduler_block_size'''
+
+NEW = '''    # spark-prefix-align: back off to the scheduler block size ONLY when mamba is
+    # genuinely non-align. A mamba group whose block_size != cache_config.block_size
+    # while mamba_cache_mode == "align" just means page-size unification scaled the
+    # block up (e.g. to match a larger DFlash drafter page); the GCD below still
+    # divides every group, so fine hashing remains valid. (Upstream's proxy test
+    # `block_size != cache_config.block_size` wrongly tripped this branch for the
+    # hybrid+DFlash+prefix-caching case, forcing hash_block_size=LCM and breaking
+    # HybridKVCacheCoordinator's divisibility assert.)
+    if getattr(cache_config, "mamba_cache_mode", "none") != "align" and any(
+        isinstance(g.kv_cache_spec, MambaSpec)
+        and g.kv_cache_spec.block_size != cache_config.block_size
+        for g in groups
+    ):
+        return scheduler_block_size, scheduler_block_size'''
+
+if OLD not in src:
+    print("[patch_prefix_align] ERROR: back-off anchor not found — vLLM source differs", flush=True)
+    sys.exit(1)
+src = src.replace(OLD, NEW, 1)
+
+# Confirmation log right before the final return so the serve log proves the fix engaged.
+RET_OLD = '''    if any(bs % hash_block_size != 0 for bs in group_block_sizes):
+        raise ValueError(
+            f"Invalid hash_block_size={hash_block_size}; all KV cache group "
+            f"block sizes must be divisible by hash_block_size. "
+            f"Got group block sizes={group_block_sizes}."
+        )
+    return scheduler_block_size, hash_block_size'''
+RET_NEW = '''    if any(bs % hash_block_size != 0 for bs in group_block_sizes):
+        raise ValueError(
+            f"Invalid hash_block_size={hash_block_size}; all KV cache group "
+            f"block sizes must be divisible by hash_block_size. "
+            f"Got group block sizes={group_block_sizes}."
+        )
+    print("SPARK_PREFIX_ALIGN resolved scheduler=%r hash=%r (mode=%r blocks=%r)" % (
+        scheduler_block_size, hash_block_size,
+        getattr(cache_config, "mamba_cache_mode", "none"), group_block_sizes),
+        flush=True)
+    return scheduler_block_size, hash_block_size'''
+if RET_OLD in src:
+    src = src.replace(RET_OLD, RET_NEW, 1)
+else:
+    print("[patch_prefix_align] WARN: confirmation-log anchor not found (non-fatal)", flush=True)
+
+P.write_text(src)
+print("[patch_prefix_align] patched OK (align-aware back-off)", flush=True)

runtime/serve.sh

@@ -76,6 +76,28 @@ TOOL_ARG=()
 [ -n "$TOOL_PARSER" ] && TOOL_ARG+=(--enable-auto-tool-choice --tool-call-parser "$TOOL_PARSER")
 [ -n "$REASONING_PARSER" ] && TOOL_ARG+=(--reasoning-parser "$REASONING_PARSER")
 
+# Prefix caching: ON by default; set PREFIX_CACHE=0 to disable. It is a clear win for agentic
+# multi-turn / long-context re-reads and neutral for single-turn c=1, so it is on by default.
+# Enabling it selects vLLM's HybridKVCacheCoordinator. With the DFlash drafter 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 it. That makes the (align-mode) mamba block
+# != cache_config.block_size, which trips resolve_kv_cache_block_sizes' back-off and forces
+# hash_block_size = LCM (4480); the drafter group stays at 2240, so the coordinator's
+# `block_size % hash_block_size` assert dies. patch_prefix_align.py makes that back-off
+# align-aware, so resolve uses the GCD (2240) — which divides every group (4480 and 2240) and
+# is the correct finer hash granularity (vLLM's intended hash_block_size<block_size design).
+# Validated 2026-06-28 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.
+PREFIX_CACHE="${PREFIX_CACHE:-1}"
+if [ "$PREFIX_CACHE" != "0" ]; then
+  echo "[serve] prefix caching ON (default) — applying align-aware hash_block_size fix"
+  python3 /host/patch_prefix_align.py
+  PREFIX_ARG=(--enable-prefix-caching)
+else
+  echo "[serve] prefix caching OFF (PREFIX_CACHE=0)"
+  PREFIX_ARG=(--no-enable-prefix-caching)
+fi
+
 exec vllm serve "$MODEL" \
   --served-model-name qwen \
   --host 0.0.0.0 --port "$PORT" \
@@ -83,7 +105,7 @@ exec vllm serve "$MODEL" \
   --max-num-seqs "$MAX_NUM_SEQS" \
   --max-num-batched-tokens "$MAX_BATCHED_TOKENS" \
   --gpu-memory-utilization "$GPU_MEM" \
-  --no-enable-prefix-caching \
+  "${PREFIX_ARG[@]}" \
   --enable-chunked-prefill \
   --trust-remote-code \
   --load-format "$LOAD_FORMAT" \