kv: recover dense pool 376k->427k (+13%) at the same 0.82 headroom

Reclaims over-reserved memory rather than spending headroom, so it honors
the conservative 0.82 / 14 GB-OS-reserve choice (still ~15 GiB free under
peak load) while closing most of the gap to the dflash pool (456k):

- int8 lm-head: build int8 on first warmup forward, then free the dead bf16
  copy (~1.4 GiB) + empty_cache() so the block lands in the KV pool. Safe:
  the DFlash drafter shares the one int8 lm_head module (verified one int8
  build; drafter ckpt has no lm_head/embed) and tie_word_embeddings=False,
  so bf16 has no reader (bias fallback = int8-GEMV + add-bias).
  SPARK_KEEP_BF16_LMHEAD=1 restores keep-bf16.
- VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 in serve.sh/mtp_serve.sh:
  returns the CUDA-graph over-estimate (~0.6 GiB; actual capture ~0.14 GiB
  from the wide 0.82 headroom). Overridable with =1.
- scripts/stress_mem.sh: drive the concurrency bank while sampling host RAM
  to prove a gpu-mem setting survives peak load without swapping.

Validated: pool 426,610 tokens, coherent output, drafter acceptance 4-12,
no new swap, throughput unchanged. max-batched-tokens 8192->4096 tested as
a KV lever and rejected (only +2.5k tokens, costs prefill speed).
This commit is contained in:
Entrpi
2026-06-24 17:41:07 +10:00
parent 1d0f2e1af3
commit 2cb5033fb8
6 changed files with 105 additions and 17 deletions
+11 -1
View File
@@ -86,7 +86,17 @@ transfer to vLLM 0.23 + DFlash:
~6.58.8 ms) — **~2× faster, argmax-exact**. Prior ports failed not on the
kernel but on **integration**: zeroing the lm-head weight corrupted the
*drafter-shared* head (garbage), and a per-row loop for B>4 was slower under
spec. v3 uses one batched kernel and **keeps** the bf16 weight.
spec. v3 uses one batched kernel; on the first warmup forward it builds the int8
copy, then **frees the now-dead bf16 weight** (~1.4 GiB) and `empty_cache()`s so
the block returns before vLLM sizes the KV pool. This is safe here because the
DFlash drafter *shares this same int8 lm_head module* (verified: exactly one int8
build; the drafter checkpoint carries no lm_head/embed tensors) and
`tie_word_embeddings=False` (so `.weight` is not aliased to `embed_tokens`) — the
bf16 copy has no remaining reader (the bias fallback does int8-GEMV + add-bias
instead). Set `SPARK_KEEP_BF16_LMHEAD=1` to restore the keep-bf16 behavior. With
`VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0` (reclaims the CUDA-graph
over-estimate), this lifts the `dense` KV pool 376,518 → **426,610 tokens**
(+13 %) at the same `0.82` headroom, validated coherent with acceptance 412.
Why the denominator matters: the 0.48 GB shared-expert saving is **0.7 % of the
71 GB on disk** but **~8 % of the ~6 GB *active per-token* footprint** (the disk