Update setup_deepseek_v4_sm89.sh

setup_deepseek_v4_sm89.sh

@@ -63,12 +63,38 @@ SERVED_NAME="DeepSeek-V4-Flash"   # name exposed on the OpenAI-compatible API
 PORT=8002
 GPUS="0,1,2,3"
 TP_SIZE=4
-MAX_MODEL_LEN=262144
-GPU_MEM_UTIL=0.96
-MAX_NUM_SEQS=4
-MAX_NUM_BATCHED_TOKENS=2048       # author's spec-decode recommendation
+# MEMORY/THROUGHPUT: validated corner from a day of SWE-bench Verified pilots
+# (2026-07-23). KV is the thin slice on this model -- weights + activations +
+# cudagraph pools eat almost the whole 48 GB card. The values below trade
+# against each other; see the generated serve.sh header for the full story.
+# Short version:
+#   * util 0.97 leaves ~1.68 GiB/GPU free at idle, ~0.5-0.8 GiB under a long
+#     trajectory -- THIN. It is above the Triton JIT-OOM line only because
+#     warmup.sh pre-compiles the kernel shapes at startup and the systemd unit
+#     auto-restarts+re-warms. For an UNATTENDED or shared deployment, prefer
+#     GPU_MEM_UTIL=0.95 (restores ~2.4 GiB headroom; identical throughput at
+#     seq=1). See serve.sh header.
+#   * max-model-len 98304 fits KV and covers agentic trajectories (~30-100k);
+#     262144 will NOT boot (KV-fit ValueError). Prompt + max_tokens must be
+#     <= this, so keep the harness's requested output tokens modest.
+#   * max-num-seqs 1: at seq>1, multiple 30-100k trajectories don't fit KV at
+#     once and vLLM preempts + RE-PREFILLS from scratch (prefix-cache hit
+#     collapses <40%, ~10x wall-clock blowup). One resident stream (prefix hit
+#     ~85-91%) wins decisively here.
+#   * batched-tokens 4096: raising it toggles chunked prefill, whose buffers
+#     come out of the KV pool -- not free on this card.
+# Pilot results (50-instance SWE-bench Verified, error_instances=0):
+#   64k/8192-out=56%  |  96k/8192-out=60% (best)  |  96k/high-effort=60% (no gain)
+MAX_MODEL_LEN=98304
+GPU_MEM_UTIL=0.97
+MAX_NUM_SEQS=1
+MAX_NUM_BATCHED_TOKENS=4096       # 4096 tested; raising toggles chunked prefill (eats KV)
 NUM_SPEC_TOKENS=1                 # MTP draft depth (1 = validated config)
-REASONING_EFFORT="high"           # default chat-template reasoning effort
+REASONING_EFFORT="medium"         # default chat-template reasoning effort.
+                                  # high effort showed NO resolved-rate gain over
+                                  # medium on the pilots (60% both) at ~40% more
+                                  # wall-clock and more step-limit failures.
+WARMUP_WAIT_MAX=300               # warmup.sh: seconds to wait for the API before giving up
 SERVICE_USER="$(id -un)"          # user baked into the generated systemd unit
 
 PYTHON_BIN="python3.12"                    # venv interpreter (wheel is cp312)
@@ -344,11 +370,35 @@ cat > "$PROJECT_DIR/serve.sh" <<SERVE_EOF
 #    letting hub ref resolution run under HF_HUB_OFFLINE: an interrupted
 #    online check can leave refs/main pointing at a snapshot that was never
 #    downloaded (LocalEntryNotFoundError at boot).
+#  * MEMORY IS THE KNIFE-EDGE on this model (incident 2026-07-23). KV is the
+#    thin slice after weights/activations/cudagraphs; two failure modes bound
+#    the tuning window:
+#      - JIT-OOM (too little headroom): Triton compiles a kernel per prompt
+#        shape and loading the cubin (cuModuleLoadData) needs driver-level free
+#        VRAM outside the torch pool. At high util a never-seen shape mid-run
+#        OOM'd the load in _tf32_hc_prenorm_gemm and killed a worker, taking a
+#        2h benchmark down. Mitigated by warmup.sh (pre-compiles shapes at
+#        startup) + systemd auto-restart+re-warm.
+#      - KV-fit refusal (too little KV): boot fails if one full max-model-len
+#        sequence won't fit the pool (262144 needs ~1.84 GiB/GPU -> NO BOOT).
+#    Current values: util 0.97, max-model-len 98304, max-num-seqs 1. MEASURED
+#    ~1.68 GiB/GPU free at idle, ~0.5-0.8 GiB under a long trajectory -- thin.
+#    For UNATTENDED/shared serving, drop util to 0.95 (restores ~2.4 GiB
+#    headroom, same throughput at seq=1). Do NOT raise max-model-len toward
+#    262k without lowering util -- that reopens the JIT-OOM hole.
+#  * max-num-seqs 1: at seq>1 the KV pool can't hold multiple 30-100k agentic
+#    trajectories, so vLLM preempts and RE-PREFILLS evicted sequences (prefix
+#    hit <40%, ~10x wall-clock blowup). seq=1 keeps one context resident
+#    (prefix hit ~85-91%) decoding continuously; it wins decisively here.
+#  * PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True keeps the torch pool from
+#    overgrowing past its profiled peak on long prefills and eating the JIT
+#    headroom. (If graph capture ever rejects it on a future wheel, drop it.)
 #  * MTP measured at 115-133 tok/s (~1.5-1.65x over the ~80 tok/s baseline),
 #    96.6-100% draft acceptance. To run WITHOUT speculative decoding, delete
 #    the --speculative-config line (and optionally --max-num-batched-tokens).
-#  * First requests JIT-compile ~a dozen Triton kernels (one-time latency
-#    spikes); throughput settles after the warm-up traffic.
+#  * First requests JIT-compile Triton kernels per shape (one-time latency
+#    spikes). Run ./warmup.sh once after every server start so the shapes
+#    compile at minute 0, not minute 100 of an unattended run.
 set -euo pipefail
 cd "\$(dirname "\$0")"
 
@@ -396,6 +446,7 @@ exec env -u VLLM_TRITON_MLA_SPARSE_HEAD_BLOCK_SIZE \\
     PATH="$CUDA_HOME_DIR/bin:\$PATH" \\
     FLASHINFER_DISABLE_VERSION_CHECK=1 \\
     VLLM_USE_FLASHINFER_SAMPLER=0 \\
+    PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \\
     HF_HUB_OFFLINE=1 \\
     CUDA_VISIBLE_DEVICES=$GPUS \\
     "\$(dirname "\$0")/.venv/bin/vllm" serve "\$MODEL_PATH" \\
@@ -429,6 +480,69 @@ echo
 CANARY_EOF
 chmod +x "$PROJECT_DIR/canary.sh"
 
+# ----------------------------- GENERATE warmup.sh ----------------------------
+# Forces the Triton per-shape specializations to JIT-compile at startup, in
+# fresh VRAM headroom and under supervision, instead of lazily mid-run into a
+# thin pool (incident 2026-07-23). Wired into the systemd unit's ExecStartPost.
+log "Writing $PROJECT_DIR/warmup.sh"
+cat > "$PROJECT_DIR/warmup.sh" <<WARMUP_EOF
+#!/usr/bin/env bash
+# Triton-specialization warmup for the DeepSeek-V4-Flash server.
+# -- generated by setup_deepseek_v4_sm89.sh
+#
+# The wait-for-server loop is BOUNDED (WAIT_MAX, default ${WARMUP_WAIT_MAX}s):
+# as ExecStartPost, an unbounded loop would pin the unit in
+# 'activating (start-post)' with a dead main PID if the engine dies during
+# init. Bail fast so systemd can fail and Restart= act.
+#
+# Usage: ./warmup.sh [host] [port]      (defaults: 127.0.0.1 $PORT)
+#        WAIT_MAX=<seconds> overrides the server-wait bound.
+set -euo pipefail
+HOST="\${1:-127.0.0.1}"
+PORT="\${2:-$PORT}"
+URL="http://\$HOST:\$PORT/v1/chat/completions"
+MODEL="$SERVED_NAME"
+WAIT_MAX="\${WAIT_MAX:-$WARMUP_WAIT_MAX}"
+MAXLEN="$MAX_MODEL_LEN"
+
+echo "==> waiting for server at \$HOST:\$PORT (up to \${WAIT_MAX}s) ..."
+deadline=\$(( \$(date +%s) + WAIT_MAX ))
+until curl -sf "http://\$HOST:\$PORT/v1/models" >/dev/null 2>&1; do
+    if [ "\$(date +%s)" -ge "\$deadline" ]; then
+        echo "XX server not reachable after \${WAIT_MAX}s -- engine likely died during init; giving up" >&2
+        exit 1
+    fi
+    sleep 5
+done
+echo "==> server is up"
+
+req() { # req <approx_token_count>
+    local n="\$1"
+    local prompt payload
+    prompt="\$(yes 'the' | head -n "\$n" | tr '\n' ' ')"
+    # Build JSON via printf to avoid nested-quote escaping hazards.
+    payload="\$(printf '{"model":"%s","messages":[{"role":"user","content":"%s"}],"max_tokens":16,"temperature":0}' "\$MODEL" "\$prompt")"
+    curl -s -o /dev/null -w "    HTTP %{http_code} in %{time_total}s  (n=\$n)\n" \\
+        "\$URL" -H 'Content-Type: application/json' --data "\$payload"
+}
+
+# Sweep shapes up to (but not over) --max-model-len; an over-cap request just
+# 400s and compiles nothing useful. Top entry sits just under MAXLEN.
+echo "==> sequential shape sweep (long ones take a while; that is the point)"
+top=\$(( MAXLEN > 4096 ? MAXLEN - 4096 : MAXLEN / 2 ))
+for n in 16 100 511 1024 2000 4096 8192 16384 32768 65536 "\$top"; do
+    [ "\$n" -le "\$MAXLEN" ] && req "\$n"
+done
+
+echo "==> warmup done. Driver-level headroom per GPU (want >= ~1 GiB free):"
+if command -v nvidia-smi >/dev/null 2>&1; then
+    nvidia-smi --query-gpu=index,memory.used,memory.total --format=csv
+else
+    echo "    (nvidia-smi not on PATH here; check on the host)"
+fi
+WARMUP_EOF
+chmod +x "$PROJECT_DIR/warmup.sh"
+
 # ----------------------------- GENERATE systemd unit -------------------------
 log "Writing $PROJECT_DIR/deepseek-v4-flash.service"
 cat > "$PROJECT_DIR/deepseek-v4-flash.service" <<UNIT_EOF
@@ -438,6 +552,13 @@ cat > "$PROJECT_DIR/deepseek-v4-flash.service" <<UNIT_EOF
 # supervises it. systemd's clean service environment additionally guarantees
 # no stale shell-profile exports can leak in.
 #
+# Incident 2026-07-23: a Triton JIT module load OOM'd and killed a worker
+# mid-run; the API server then exited *cleanly* (main PID rc 0), so the old
+# Restart=on-failure never fired and the port stayed dead. Hence Restart=always
+# below -- for an inference server a crash that looks like a clean exit is
+# still a crash; only 'systemctl stop' should keep it down. ExecStartPost runs
+# warmup.sh after every (re)start so Triton specializations compile up front.
+#
 # Install:
 #   sudo cp $PROJECT_DIR/deepseek-v4-flash.service /etc/systemd/system/
 #   sudo systemctl daemon-reload
@@ -466,7 +587,16 @@ ExecStartPre=/usr/bin/nvidia-smi
 
 ExecStart=$PROJECT_DIR/serve.sh
 
-Restart=on-failure
+# Warm the Triton shape specializations after EVERY start (incl. auto-restart).
+# warmup.sh waits for the API itself and is bounded, so it cannot hang the
+# unit; '-' prefix + timeout mean a wedged warmup never fails the start.
+ExecStartPost=-/usr/bin/timeout 900 $PROJECT_DIR/warmup.sh
+# Budget must cover ExecStartPre + weight load + bounded warmup.
+TimeoutStartSec=1200
+
+# Restart=always, NOT on-failure: a JIT-OOM crash can present as a clean exit
+# (see incident note above); only an explicit 'systemctl stop' should end it.
+Restart=always
 # Weight load alone is ~80 s (+ ~5 s MTP drafter); don't hammer restarts.
 RestartSec=15
 # Give the engine time to tear down 4 TP workers and NCCL cleanly on stop.
@@ -491,13 +621,15 @@ Next steps:
      run once:
          $PROJECT_DIR/.venv/bin/huggingface-cli download "$MODEL"
   2. Start the server:      $PROJECT_DIR/serve.sh
-     Expect ~80 s of weight loading (+ ~5 s for the MTP drafter), brief
-     Triton JIT compiles during warmup and the first requests, then
+     Expect ~80 s of weight loading (+ ~5 s for the MTP drafter), then
      'Application startup complete'.
-  3. Smoke test:            $PROJECT_DIR/canary.sh
+  3. Warm up (once/start): $PROJECT_DIR/warmup.sh
+     Compiles Triton shapes up front and prints free VRAM/GPU (want >= ~1 GiB).
+     Runs automatically via ExecStartPost under the systemd unit.
+  4. Smoke test:            $PROJECT_DIR/canary.sh
      A correct one-sentence answer naming Muenchen = numerics are sound.
      Watch the SpecDecoding metrics lines: draft acceptance should sit
      around 95-100% with mean acceptance length near 2.00.
-  4. Run as a service:      see header of
+  5. Run as a service:      see header of
          $PROJECT_DIR/deepseek-v4-flash.service
 NEXT_EOF