diff --git a/README.md b/README.md index 462fd6c..76e7656 100644 --- a/README.md +++ b/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) | diff --git a/docs/FINDINGS.md b/docs/FINDINGS.md index db48be2..88e0489 100644 --- a/docs/FINDINGS.md +++ b/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 diff --git a/runtime/patch_prefix_align.py b/runtime/patch_prefix_align.py new file mode 100644 index 0000000..d1f80ad --- /dev/null +++ b/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) diff --git a/runtime/serve.sh b/runtime/serve.sh index a25ab35..b41613c 100755 --- a/runtime/serve.sh +++ b/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_size0.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" \