diff --git a/README.md b/README.md index 4ee20d4..8318527 100644 --- a/README.md +++ b/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 | diff --git a/docs/FINDINGS.md b/docs/FINDINGS.md index c552346..db48be2 100644 --- a/docs/FINDINGS.md +++ b/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 diff --git a/runtime/mtp_serve.sh b/runtime/mtp_serve.sh index 4c01db6..5223378 100755 --- a/runtime/mtp_serve.sh +++ b/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 \ diff --git a/runtime/patch_int8_lmhead_v3.py b/runtime/patch_int8_lmhead_v3.py index 4554af9..c789bd5 100644 --- a/runtime/patch_int8_lmhead_v3.py +++ b/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 =================== ''' diff --git a/runtime/serve.sh b/runtime/serve.sh index 5451514..05d6c57 100755 --- a/runtime/serve.sh +++ b/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 diff --git a/scripts/stress_mem.sh b/scripts/stress_mem.sh new file mode 100644 index 0000000..8f5b796 --- /dev/null +++ b/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==""||$2maxs)maxs=$3} END{print " min_avail_MiB", min, " max_swap_MiB", maxs+0}' "$LOG" +rm -f "$LOG" +echo "[stress] conc_workloads exit=$RC"