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).

README.md

@@ -98,9 +98,17 @@ target hardware at the shipped defaults (`gpu-mem 0.82`, `ctx 262144`, `seqs 3`)
 
 | Measurement (default `dense` profile) | Value |
 |---|---|
-| Free memory at READY | **~18 GiB** (responsive, no swap) |
-| GPU KV cache pool | **376,518 tokens** (`dflash` profile: 456,664) |
-| Max concurrency at full 262 144 | **1.44×** (`dflash`: 1.74×) |
+| Free memory at READY | **~16 GiB** (responsive, no swap; ~15 GiB under peak 3-stream load) |
+| GPU KV cache pool | **426,610 tokens** (`dflash` profile: 456,664) |
+| Max concurrency at full 262 144 | **1.63×** (`dflash`: 1.74×) |
+
+The `dense` pool was lifted from 376,518 → **426,610 tokens** (+13 %) at the *same*
+`0.82` headroom by reclaiming over-reserved memory rather than spending headroom:
+the int8 lm-head frees its now-dead bf16 copy (~1.4 GiB — the DFlash drafter shares
+the same int8 lm_head, and `tie_word_embeddings=False`, so it is genuinely unused
+after quantization), and `VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0` returns the
+CUDA-graph over-estimate (~0.6 GiB; actual capture is ~0.14 GiB, drawn from the wide
+0.82 headroom). Both are validated coherent with full drafter acceptance (4–12).
 
 Decode-only throughput (streaming, excludes prefill), by workload and concurrency:
 
@@ -116,16 +124,16 @@ exceeds the ~81 headline, which is the median over 10 varied real turns includin
 longer, slower-prefilling contexts.)
 
 A typical load — three streams under ~100 k each (≈ <300 k tokens) — fits the
-376 k pool with margin, and a single stream can still reach the full 262 144
+426 k pool with margin, and a single stream can still reach the full 262 144
 context. At `gpu-mem` 0.88–0.89 the static footprint leaves only ~5 GiB free; the
-host then swaps and requests stall. `0.82` is the validated value (~18 GiB free
-on `dense`). Defaults (override via flags or environment variables):
+host then swaps and requests stall. `0.82` is the validated value (~16 GiB free
+on `dense`, ~15 GiB under peak load). Defaults (override via flags or environment variables):
 
 | Flag / env | Default | Note |
 |---|---|---|
 | `--gpu-mem` / `GPU_MEM` | **0.82** | ~14 GiB free (validated); 0.88+ over-subscribes and swaps |
 | `--ctx` / `CTX` (`MAX_MODEL_LEN`) | **262144** | model native max; a single stream can reach any length up to this. Costs only KV-pool sizing — the CUDA-graph compile range tracks `max-batched-tokens`, not `ctx` |
-| `--max-num-seqs` / `MAX_NUM_SEQS` | **3** | concurrent-stream cap; the pool holds ~1.4× (`dense`) / ~1.7× (`dflash`) a full-262 k context, ample for <100 k streams |
+| `--max-num-seqs` / `MAX_NUM_SEQS` | **3** | concurrent-stream cap; the pool holds ~1.6× (`dense`) / ~1.7× (`dflash`) a full-262 k context, ample for <100 k streams |
 | `--max-batched-tokens` / `MAX_BATCHED_TOKENS` | **8192** | chunked-prefill chunk, kept **below** `ctx` so a long prefill does not batch all at once |
 
 The default operating point is a single stream (no contention; ~81 tok/s on agent
@@ -159,7 +167,7 @@ Selected with `--profile`:
 | Profile | Stack | Best for | Measured |
 |---|---|---|---|
 | **`dense`** *(default)* | hybrid INT4+FP8 + int8 lm-head + DFlash n=12 | general — downloads the prebuilt hybrid; ≈ dflash on agents, +28% on base | 36.0 base (+28%) · 59.0 albond-bench · ~81 Hermes |
-| `dflash` | INT4 + DFlash n=12 | agent path; largest KV pool (456k vs 376k) | **~81 tok/s** Hermes · 53.7 albond-bench |
+| `dflash` | INT4 + DFlash n=12 | agent path; largest KV pool (456k vs 426k) | **~81 tok/s** Hermes · 53.7 albond-bench |
 | `base` | plain INT4, no speculative decode | airtight baseline | 28.2 tok/s c=1 |
 | `mtp` | INT4 + native MTP-2 head | comparison | ~40 tok/s Hermes |
 

docs/FINDINGS.md

@@ -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

runtime/mtp_serve.sh

@@ -13,6 +13,8 @@ MAX_NUM_SEQS="${MAX_NUM_SEQS:-3}"
 MAX_BATCHED_TOKENS="${MAX_BATCHED_TOKENS:-8192}"
 LOAD_FORMAT="${LOAD_FORMAT:-fastsafetensors}"
 PORT="${PORT:-8000}"
+# Reclaim the CUDA-graph memory over-estimate to KV (see serve.sh). Set =1 to restore.
+export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS="${VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS:-0}"
 echo "[mtp] qwen3_5_mtp — backend=$BACKEND, num_speculative_tokens=$NSPEC, model=$MODEL"
 exec vllm serve "$MODEL" \
   --served-model-name qwen \

runtime/patch_int8_lmhead_v3.py

@@ -5,9 +5,13 @@ batched int8 GEMV (one kernel launch for any batch), leaving the existing TP
 gather + org_vocab_size trim untouched.
 
 Fixes vs the broken v2 port:
-  * KEEPS the bf16 lm_head weight (DFlash drafter shares it) — does NOT zero it.
-    Trades the memory saving for correctness; the speed win is the int8 read in
-    _get_logits, independent of keeping bf16 around.
+  * Builds the int8 copy lazily on the first warmup forward, then FREES the dead
+    bf16 weight (~1.4 GiB -> KV pool) + empty_cache(). Safe here: the DFlash
+    drafter SHARES this same int8 lm_head module (one int8 build; drafter ckpt has
+    no lm_head/embed) and tie_word_embeddings=False, so bf16 has no remaining
+    reader (bias fallback does int8-GEMV + add-bias). v2's "garbage" was zeroing
+    bf16 eagerly before the shared/int8 path was ready. SPARK_KEEP_BF16_LMHEAD=1
+    restores the keep-bf16 behavior.
   * Single BATCHED kernel (dot-based, pad B->16) for ALL B — no per-row Python
     loop (the v2 B>4 loop was what made spec decode SLOWER).
   * Fixed proven config (N128/K128/w4/s3, ~227 GB/s, argmax-exact vs bf16 on the
@@ -82,6 +86,7 @@ def _spark_int8_gemm(hidden, w_int8, w_scale):
 
 
 def _spark_int8_lmhead_apply(self, lm_head, hidden_states, embedding_bias):
+    import os
     import sys
     import torch
     if not getattr(lm_head, "_spark_int8_ready", None) is True and \\
@@ -95,12 +100,34 @@ def _spark_int8_lmhead_apply(self, lm_head, hidden_states, embedding_bias):
             lm_head._spark_w_int8 = w_int8.contiguous()
             lm_head._spark_w_scale = scales.to(torch.float16)
             lm_head._spark_int8_ready = True
-            print("DGX_SPARK_INT8_LMHEAD_V3: lm_head -> int8 (%s), bf16 kept for shared drafter"
-                  % (list(w_int8.shape),), file=sys.stderr, flush=True)
+            # Free the now-dead bf16 copy to give its ~1.4 GiB back to the KV pool.
+            # SAFE here because: (1) the DFlash drafter SHARES this same lm_head
+            # module (verified: one int8 build, drafter ckpt has no lm_head) so it
+            # also uses the int8 path; (2) tie_word_embeddings=False so .weight is
+            # NOT shared with embed_tokens; (3) this runs in the init warmup BEFORE
+            # KV-cache sizing, so the pool grows. The only bf16 reader left is the
+            # bias fallback, handled post-hoc below. Toggle off with
+            # SPARK_KEEP_BF16_LMHEAD=1 (reverts to the v3 keep-bf16 behavior).
+            if os.environ.get("SPARK_KEEP_BF16_LMHEAD", "0") != "1":
+                lm_head.weight.data = torch.empty(0, dtype=w.dtype, device=w.device)
+                lm_head._spark_bf16_freed = True
+                # Return the freed block to the driver NOW so vLLM's mem_get_info
+                # based KV sizing (which runs right after this warmup) actually
+                # counts it — otherwise the caching allocator holds most of it.
+                torch.cuda.empty_cache()
+                _spark_msg = "bf16 FREED (int8-only, ~1.4 GiB -> KV)"
+            else:
+                _spark_msg = "bf16 kept for shared drafter"
+            print("DGX_SPARK_INT8_LMHEAD_V3: lm_head -> int8 (%s), %s"
+                  % (list(w_int8.shape), _spark_msg), file=sys.stderr, flush=True)
         else:
             lm_head._spark_int8_disabled = True
-    if getattr(lm_head, "_spark_int8_ready", False) and embedding_bias is None:
-        return _spark_int8_gemm(hidden_states, lm_head._spark_w_int8, lm_head._spark_w_scale)
+    if getattr(lm_head, "_spark_int8_ready", False):
+        if embedding_bias is None:
+            return _spark_int8_gemm(hidden_states, lm_head._spark_w_int8, lm_head._spark_w_scale)
+        if getattr(lm_head, "_spark_bf16_freed", False):
+            # bias path can't read the freed bf16 weight; do int8 GEMV + add bias.
+            return _spark_int8_gemm(hidden_states, lm_head._spark_w_int8, lm_head._spark_w_scale) + embedding_bias
     return lm_head.quant_method.apply(lm_head, hidden_states, bias=embedding_bias)
 # =================== end DGX_SPARK_INT8_LMHEAD_V3 ===================
 '''

runtime/serve.sh

@@ -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

scripts/stress_mem.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+# stress_mem.sh — drive the server with the concurrency bank while sampling host
+# memory, to prove a gpu-memory-utilization setting survives peak load WITHOUT
+# swapping (the failure mode that hard-freezes the Spark). Reports min available
+# RAM and max swap observed across the run.
+#
+#   usage: stress_mem.sh [LEVELS]   (default "1,2,3")
+# Run on the box; the server (container qwen-spark) must be READY on :8000.
+set -uo pipefail
+LEVELS="${1:-1,2,3}"
+HERE="$(cd "$(dirname "$0")" && pwd)"
+LOG=/tmp/stress_memlog.$$
+
+# baseline
+echo "[stress] swap before:"; free -m | awk '/Swap:/{print "  swap_used_MiB",$3}'
+echo "[stress] avail before:"; free -m | awk '/Mem:/{print "  avail_MiB",$7}'
+
+# background sampler: timestamp, avail MiB, swap-used MiB, every 1s
+( for i in $(seq 1 900); do
+    free -m | awk -v t="$i" '/Mem:/{a=$7} /Swap:/{s=$3} END{print t, a, s}'
+    sleep 1
+  done ) > "$LOG" 2>/dev/null &
+SPID=$!
+
+echo "[stress] running conc_workloads --levels $LEVELS ..."
+python3 "$HERE/conc_workloads.py" --levels "$LEVELS"
+RC=$?
+
+kill "$SPID" 2>/dev/null
+echo "[stress] === memory envelope during load ==="
+awk 'NF>=3 {if(min==""||$2<min)min=$2; if($3>maxs)maxs=$3} END{print "  min_avail_MiB", min, " max_swap_MiB", maxs+0}' "$LOG"
+rm -f "$LOG"
+echo "[stress] conc_workloads exit=$RC"