TBNilles/qwen3.5-122B-A10B-on-spark
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+3 -2
View File
@@ -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) |