diff --git a/README.md b/README.md index c358d18..0c5d8ed 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,36 @@ Preview without running: `... | bash -s -- --help`. GB10 is detected via `nvidia-smi --query-gpu=compute_cap` returning `12.1`; anything else needs `--force`. +### Memory & context (defaults are tuned for max context + KV depth) + +This model is **36/48 linear-attention (GDN) + 12 full-attention** layers, and the +attention layers use GQA `num_key_value_heads=2`, `head_dim=256` → **only +~24 KiB/token of KV**. A full **262 144** context (the model's native max) is just +~6 GiB of KV; the GDN layers hold a *fixed* per-sequence state (~0.2 GiB) that +does **not** grow with context. So on the 128 GB (119 GiB) GB10, reserving 14 GB +for the OS leaves ~106 GiB for vLLM: + +| | | +|---|---| +| weights (INT4) + DFlash drafter | ~64 GiB | +| CUDA graphs + activations | ~10 GiB | +| **KV pool** | **~32 GiB ≈ 1.38 M tokens** | + +The KV pool (~1.38 M tokens) dwarfs a single 262 144 context, so single context is +capped by the *model*, not memory. Defaults (override via flags/env): + +| Flag / env | Default | Note | +|---|---|---| +| `--gpu-mem` / `GPU_MEM` | **0.89** | reserves ~14 GB; drop to 0.87 if the OOM-guard fires on first load | +| `--ctx` / `CTX` (`MAX_MODEL_LEN`) | **262144** | model native max | +| `--max-num-seqs` / `MAX_NUM_SEQS` | **1** | single-stream; raising it is nearly free (pool ≫ one context) | +| `--max-batched-tokens` / `MAX_BATCHED_TOKENS` | **8192** | chunked-prefill chunk — kept **below** ctx so a long prefill doesn't batch all at once | + +> Unified-memory OOM **hard-freezes** the box, and vLLM's profiler can undershoot +> peak by a couple GB — always bring the server up under +> [`scripts/monitor.sh`](scripts/monitor.sh) (OOM auto-kill guard) the first time +> at a new `gpu-mem`/`ctx`. + ## What you get — profiles Pick with `--profile`: diff --git a/install.sh b/install.sh index 1b6a87f..959f494 100755 --- a/install.sh +++ b/install.sh @@ -55,8 +55,10 @@ NAME="${NAME:-qwen-spark}" PROFILE="dflash" # dflash | dense | base | mtp NSPEC="" # override num_speculative_tokens (default per profile) PORT="${PORT:-8000}" -CTX="${CTX:-16384}" -GPU_MEM="${GPU_MEM:-0.8}" +CTX="${CTX:-262144}" # max-model-len: model native max (KV is ~24 KiB/token) +GPU_MEM="${GPU_MEM:-0.89}" # reserves ~14 GB on a 128 GB (119 GiB) GB10 +MAX_NUM_SEQS="${MAX_NUM_SEQS:-1}" # single-stream; raise for concurrency (cheap here) +MAX_BATCHED_TOKENS="${MAX_BATCHED_TOKENS:-8192}" # chunked-prefill chunk (decoupled from ctx) BACKEND="${BACKEND:-flash_attn}" FORCE_HW=0 @@ -91,14 +93,16 @@ Flags: --hf-home DIR Use/populate this HF cache dir (default: $HF_HOME). --nspec N num_speculative_tokens (default 12 dflash/dense, 2 mtp, 0 base). --port N Server port (default: $PORT). - --ctx N max-model-len (default: $CTX). - --gpu-mem F gpu-memory-utilization, keep <=0.84 on 128GB (default: $GPU_MEM). + --ctx N max-model-len (default: $CTX = model native max). + --gpu-mem F gpu-memory-utilization (default: $GPU_MEM reserves ~14 GB on 119 GiB). + --max-num-seqs N concurrent sequences (default: $MAX_NUM_SEQS, single-stream). + --max-batched-tokens N chunked-prefill chunk (default: $MAX_BATCHED_TOKENS; keep < ctx). --force Skip the GB10/SM121 host check. --no-smoke Start the server but skip the Paris smoke test. Environment equivalents: QWEN_IMAGE TARGET_REPO DRAFT_REPO FP8_REPO HF_HOME HYBRID_DIR - NAME PORT CTX GPU_MEM BACKEND REPO_DIR REPO_URL + NAME PORT CTX GPU_MEM MAX_NUM_SEQS MAX_BATCHED_TOKENS BACKEND REPO_DIR REPO_URL EOF } @@ -116,6 +120,8 @@ while [[ $# -gt 0 ]]; do --port) PORT="$2"; shift 2 ;; --ctx) CTX="$2"; shift 2 ;; --gpu-mem) GPU_MEM="$2"; shift 2 ;; + --max-num-seqs) MAX_NUM_SEQS="$2"; shift 2 ;; + --max-batched-tokens) MAX_BATCHED_TOKENS="$2"; shift 2 ;; --backend) BACKEND="$2"; shift 2 ;; --force) FORCE_HW=1; shift ;; --no-smoke) SKIP_SMOKE=1; shift ;; @@ -270,7 +276,8 @@ start_server() { docker rm -f "$NAME" >/dev/null 2>&1 || true # shellcheck disable=SC2086 docker run -d --name "$NAME" --gpus all --net=host --ipc=host --ulimit memlock=-1:-1 \ - -e HF_HOME=/hf -e MAX_MODEL_LEN="$CTX" -e GPU_MEM="$GPU_MEM" ${HF_TOKEN:+-e HF_TOKEN="$HF_TOKEN"} \ + -e HF_HOME=/hf -e MAX_MODEL_LEN="$CTX" -e GPU_MEM="$GPU_MEM" \ + -e MAX_NUM_SEQS="$MAX_NUM_SEQS" -e MAX_BATCHED_TOKENS="$MAX_BATCHED_TOKENS" ${HF_TOKEN:+-e HF_TOKEN="$HF_TOKEN"} \ "${model_env[@]}" \ -v "$HF_HOME:/hf" -v "$REPO_DIR/runtime:/host:ro" "${mounts[@]}" \ --entrypoint bash "$IMAGE" "$wrapper" $serve_args >/dev/null diff --git a/runtime/mtp_serve.sh b/runtime/mtp_serve.sh index d971725..140e6aa 100755 --- a/runtime/mtp_serve.sh +++ b/runtime/mtp_serve.sh @@ -7,16 +7,18 @@ set -euo pipefail NSPEC="${1:-2}" BACKEND="${2:-flash_attn}" MODEL="${MODEL:-Intel/Qwen3.5-122B-A10B-int4-AutoRound}" -MAX_MODEL_LEN="${MAX_MODEL_LEN:-16384}" -GPU_MEM="${GPU_MEM:-0.8}" +MAX_MODEL_LEN="${MAX_MODEL_LEN:-262144}" +GPU_MEM="${GPU_MEM:-0.89}" +MAX_NUM_SEQS="${MAX_NUM_SEQS:-1}" +MAX_BATCHED_TOKENS="${MAX_BATCHED_TOKENS:-8192}" PORT="${PORT:-8000}" echo "[mtp] qwen3_5_mtp — backend=$BACKEND, num_speculative_tokens=$NSPEC, model=$MODEL" exec vllm serve "$MODEL" \ --served-model-name qwen \ --host 0.0.0.0 --port "$PORT" \ --max-model-len "$MAX_MODEL_LEN" \ - --max-num-seqs 16 \ - --max-num-batched-tokens "$MAX_MODEL_LEN" \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --max-num-batched-tokens "$MAX_BATCHED_TOKENS" \ --gpu-memory-utilization "$GPU_MEM" \ --no-enable-prefix-caching \ --enable-chunked-prefill \ diff --git a/runtime/serve.sh b/runtime/serve.sh index 3dad0d7..89ef9f7 100755 --- a/runtime/serve.sh +++ b/runtime/serve.sh @@ -19,8 +19,10 @@ NSPEC="${1:-12}" BACKEND="${2:-flash_attn}" MODEL="${MODEL:-Intel/Qwen3.5-122B-A10B-int4-AutoRound}" DRAFT="${DRAFT:-z-lab/Qwen3.5-122B-A10B-DFlash}" -MAX_MODEL_LEN="${MAX_MODEL_LEN:-16384}" -GPU_MEM="${GPU_MEM:-0.8}" +MAX_MODEL_LEN="${MAX_MODEL_LEN:-262144}" # model native max; KV is ~24 KiB/token so it fits +GPU_MEM="${GPU_MEM:-0.89}" # ~14 GB reserved on a 128 GB (119 GiB) GB10 +MAX_NUM_SEQS="${MAX_NUM_SEQS:-1}" # single-stream; each seq also reserves GDN state +MAX_BATCHED_TOKENS="${MAX_BATCHED_TOKENS:-8192}" # chunked-prefill chunk (NOT = max-model-len) PORT="${PORT:-8000}" # FLA sm121 big-tile shmem fix (prefill/TTFT only on sm121; harmless, free). @@ -49,8 +51,8 @@ exec vllm serve "$MODEL" \ --served-model-name qwen \ --host 0.0.0.0 --port "$PORT" \ --max-model-len "$MAX_MODEL_LEN" \ - --max-num-seqs 16 \ - --max-num-batched-tokens "$MAX_MODEL_LEN" \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --max-num-batched-tokens "$MAX_BATCHED_TOKENS" \ --gpu-memory-utilization "$GPU_MEM" \ --no-enable-prefix-caching \ --enable-chunked-prefill \