tune defaults for max context + KV depth: gpu-mem 0.89 (reserve ~14GB), ctx 262144, max-num-seqs 1, decouple max-batched-tokens (8192)

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`:

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}"
+CTX="${CTX:-262144}"               # max-model-len: model native max (KV is ~24 KiB/token)
-GPU_MEM="${GPU_MEM:-0.8}"
+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).
+  --ctx N                 max-model-len (default: $CTX = model native max).
-  --gpu-mem F             gpu-memory-utilization, keep <=0.84 on 128GB (default: $GPU_MEM).
+  --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

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}"
+MAX_MODEL_LEN="${MAX_MODEL_LEN:-262144}"
-GPU_MEM="${GPU_MEM:-0.8}"
+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-seqs "$MAX_NUM_SEQS" \
-  --max-num-batched-tokens "$MAX_MODEL_LEN" \
+  --max-num-batched-tokens "$MAX_BATCHED_TOKENS" \
   --gpu-memory-utilization "$GPU_MEM" \
   --no-enable-prefix-caching \
   --enable-chunked-prefill \

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}"
+MAX_MODEL_LEN="${MAX_MODEL_LEN:-262144}"      # model native max; KV is ~24 KiB/token so it fits
-GPU_MEM="${GPU_MEM:-0.8}"
+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-seqs "$MAX_NUM_SEQS" \
-  --max-num-batched-tokens "$MAX_MODEL_LEN" \
+  --max-num-batched-tokens "$MAX_BATCHED_TOKENS" \
   --gpu-memory-utilization "$GPU_MEM" \
   --no-enable-prefix-caching \
   --enable-chunked-prefill \