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:
+11
-1
@@ -86,7 +86,17 @@ transfer to vLLM 0.23 + DFlash:
|
||||
~6.5–8.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 4–12.
|
||||
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user