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.
2026-06-28 21:29:40 +10:00
parent 3f05a87cdf
commit 41d92cb6f9
4 changed files with 152 additions and 4 deletions
@@ -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)
@@ -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" \