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
+9 -1
View File
@@ -21,7 +21,7 @@ MODEL="${MODEL:-Intel/Qwen3.5-122B-A10B-int4-AutoRound}"
DRAFT="${DRAFT:-z-lab/Qwen3.5-122B-A10B-DFlash}"
MAX_MODEL_LEN="${MAX_MODEL_LEN:-262144}" # model native max; KV is ~24 KiB/token so it fits
GPU_MEM="${GPU_MEM:-0.82}" # VALIDATED: ~14 GiB free on 128 GB (119 GiB) GB10 (0.88+ over-subscribes -> swap)
MAX_NUM_SEQS="${MAX_NUM_SEQS:-3}" # 3 concurrent streams; KV pool ~457k tokens, 1.74x at full 262144
MAX_NUM_SEQS="${MAX_NUM_SEQS:-3}" # 3 concurrent streams; KV pool ~427k (dense) / ~457k (dflash) tokens at full 262144
MAX_BATCHED_TOKENS="${MAX_BATCHED_TOKENS:-8192}" # chunked-prefill chunk (NOT = max-model-len)
PORT="${PORT:-8000}"
# Read straight to the device (no mmap, no host staging) — the slow default safetensors
@@ -30,6 +30,14 @@ PORT="${PORT:-8000}"
# --safetensors-load-strategy eager via SAFETENSORS_STRATEGY).
LOAD_FORMAT="${LOAD_FORMAT:-fastsafetensors}"
# Reclaim vLLM's CUDA-graph memory OVER-estimate back to the KV pool. The profiler
# reserves ~0.7 GiB for the graph pool but capture actually uses ~0.14 GiB; disabling
# the estimate gives the difference (~0.6 GiB / ~9.5k tokens) to KV. The real capture
# then comes out of the (1 - gpu-mem) headroom, which is ~21 GiB at 0.82 -> no OOM
# risk at the shipped util. Set ESTIMATE_CUDAGRAPHS=1 to restore vLLM's default if you
# push gpu-mem very high (small headroom).
export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS="${VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS:-0}"
# FLA sm121 big-tile shmem fix (prefill/TTFT only on sm121; harmless, free).
echo "[serve] FLA sm121 big-tile shmem patch"
python3 /host/patch_fla_shmem.py || true