#!/usr/bin/env bash
# =============================================================================
# setup_deepseek_v4_sm89.sh   (stage 2: MTP speculative decoding)
#
# Model repos (byte-identical; either works):
#   https://huggingface.co/SinclairSchneider/DeepSeek-V4-Flash-W4A16-FP8-MTP-Ada
#   https://huggingface.co/canada-quant/DeepSeek-V4-Flash-W4A16-FP8-MTP
#     (the original -- all credit for the quantization to canada-quant)
#
# Recreates, from scratch, the verified-working environment for serving
# DeepSeek-V4-Flash (W4A16 INT4 experts + FP8 attention + BF16 MTP head,
# 159 GB) on 4x RTX 6000 Ada (SM89) with vLLM, including MTP speculative
# decoding.
#
# Measured on 2026-07-21 (4x RTX 6000 Ada, TP=4, fp8_ds_mla KV cache):
#   - plain decoding:            ~80 tok/s single-stream
#   - MTP (num_spec_tokens=1):   115-133 tok/s sustained, draft acceptance
#                                96.6-100%, mean acceptance length ~1.97/2.00
#
# Verified-working state this script reproduces:
#   - vLLM  0.23.1rc1.dev145+g8c631d45e.cu130   <- the ONLY release of the
#     yhfgyyf/vllm-deepseek-v4-sm89 fork whose SM89 path actually works.
#     (The newer dev1018 release regressed the SM89 indexer-logits fallback.)
#   - transformers 5.8.1, PINNED AND INSTALLED LAST. Newer releases fail
#     config validation on this checkpoint: their ALLOWED_LAYER_TYPES no
#     longer accepts the 'hash_moe' entries in mlp_layer_types. Installing
#     the vLLM wheel afterwards would let pip drag a newer transformers back
#     in, so the pin must be the final install step (and is asserted below).
#   - flashinfer_python 0.6.14+sm89 (installed, but idle on the dev145
#     Triton path; kept to match the reference environment)
#   - flashinfer-cubin 0.6.13
#   - torch 2.11.0+cu130
#   - NO deep_gemm  (its mere presence hijacks a presence-gated code path
#     and crashes SM89 -- it must NOT be installed)
#   - ONE source patch on top of the wheel (applied automatically below):
#     a BF16/unquantized-wo_a fallback in nvidia/ops/o_proj.py. The fused
#     o-projection assumes FP8 block scales; the checkpoint's BF16 MTP
#     drafter head has none and crashes without it. The FP8 target model's
#     path is untouched.
#   - Sparse MLA runs via the portable Triton path: VLLM_TRITON_MLA_SPARSE=1
#     and NO --attention-backend flag.
#   - serve.sh resolves the model to a concrete snapshot DIRECTORY at launch
#     instead of trusting hub ref resolution: an online metadata check (e.g.
#     an interrupted huggingface-cli download after the upstream repo gained
#     a new commit) can leave refs/main pointing at a snapshot that was
#     never downloaded, which breaks HF_HUB_OFFLINE resolution.
#
# What this script does NOT do (on purpose):
#   - Download the 159 GB model (see the note printed at the end)
#   - Any system-level changes (apt, /usr/local/cuda symlinks). It only
#     VERIFIES that /usr/local/cuda-13.0 exists, because everything is
#     pinned to it at launch time.
# =============================================================================

# ----------------------------- PARAMETERS ------------------------------------
PROJECT_DIR="${1:-$HOME/deepseek-v4-serve}"      # folder to create (arg 1 overrides)
MODEL="${2:-SinclairSchneider/DeepSeek-V4-Flash-W4A16-FP8-MTP-Ada}"  # HF repo id OR local path (arg 2 overrides)
# The default is the Ada mirror (ships these scripts alongside the weights).
# The original repo works identically -- the checkpoints are byte-identical:
#   ./setup_deepseek_v4_sm89.sh <dir> canada-quant/DeepSeek-V4-Flash-W4A16-FP8-MTP

SERVED_NAME="DeepSeek-V4-Flash"   # name exposed on the OpenAI-compatible API
PORT=8002
GPUS="0,1,2,3"
TP_SIZE=4
# 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="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)
CUDA_HOME_DIR="/usr/local/cuda-13.0"       # MUST be 13.0 to match cu130 wheels
LOCAL_WHEEL_DIR="$HOME/wheels"             # checked first, before downloading

# Pinned artifacts (do not change casually -- this exact combination works)
GH_REPO="yhfgyyf/vllm-deepseek-v4-sm89"
VLLM_TAG="v0.23.1rc1.dev145-g8c631d45e-cu130-sm89"
VLLM_WHEEL="vllm-0.23.1rc1.dev145+g8c631d45e.cu130-cp312-cp312-linux_x86_64.whl"
FI_TAG="v0.23.1rc1.dev1018-g8aba6ae7e-cu130-sm89"   # flashinfer wheel only exists in this release
FI_WHEEL="flashinfer_python-0.6.14+sm89-py3-none-any.whl"
TORCH_SPEC="torch==2.11.0"
TORCH_INDEX="https://download.pytorch.org/whl/cu130"
FI_CUBIN_SPEC="flashinfer-cubin==0.6.13"
TRANSFORMERS_SPEC="transformers==5.8.1"    # MUST be installed LAST (see header)
# -----------------------------------------------------------------------------

set -euo pipefail

log()  { printf '\n\033[1;32m==> %s\033[0m\n' "$*"; }
warn() { printf '\n\033[1;33m!!  %s\033[0m\n' "$*"; }
die()  { printf '\n\033[1;31mXX  %s\033[0m\n' "$*" >&2; exit 1; }

# ----------------------------- SANITY CHECKS ---------------------------------
log "Sanity checks"

command -v "$PYTHON_BIN" >/dev/null 2>&1 \
    || die "$PYTHON_BIN not found (the pinned wheels are cp312)."

command -v nvidia-smi >/dev/null 2>&1 \
    || die "nvidia-smi not found -- NVIDIA driver missing?"

[ -x "$CUDA_HOME_DIR/bin/nvcc" ] \
    || die "$CUDA_HOME_DIR/bin/nvcc not found. cu130 wheels need CUDA 13.0 here."

"$CUDA_HOME_DIR/bin/nvcc" --version | grep -q "cuda_13\.0" \
    || die "$CUDA_HOME_DIR/bin/nvcc is not CUDA 13.0. Mixed toolchains break JIT builds (one-arg __cudaLaunch vs. 13.x headers)."

[ -e "$PROJECT_DIR" ] && die "$PROJECT_DIR already exists -- refusing to clobber it."

# ----------------------------- FOLDER + VENV ---------------------------------
log "Creating $PROJECT_DIR and virtual environment"
mkdir -p "$PROJECT_DIR/wheels"
"$PYTHON_BIN" -m venv "$PROJECT_DIR/.venv"
VPY="$PROJECT_DIR/.venv/bin/python"
"$VPY" -m pip install --upgrade pip >/dev/null

# ----------------------------- WHEEL RETRIEVAL -------------------------------
# Lessons encoded here:
#   * pip REFUSES wheels whose filename was changed -- keep canonical names.
#   * GitHub release asset URLs contain '+', which must be sent as %2B.
#   * gh CLI needs auth; the plain REST API does not.
fetch_asset() {  # fetch_asset <tag> <asset-filename>
    local tag="$1" name="$2" dest="$PROJECT_DIR/wheels/$2" url
    if [ -f "$LOCAL_WHEEL_DIR/$name" ]; then
        log "Using local copy of $name from $LOCAL_WHEEL_DIR"
        cp "$LOCAL_WHEEL_DIR/$name" "$dest"
        return
    fi
    log "Downloading $name from release $tag"
    url=$(curl -fsS "https://api.github.com/repos/$GH_REPO/releases/tags/$tag" \
        | "$VPY" -c "import json,sys;print([a['browser_download_url'] for a in json.load(sys.stdin)['assets'] if a['name']=='$name'][0])") \
        || die "Could not resolve asset $name in release $tag"
    curl -fL -o "$dest" "${url//+/%2B}" || die "Download failed: $name"
}

fetch_asset "$FI_TAG"   "$FI_WHEEL"
fetch_asset "$VLLM_TAG" "$VLLM_WHEEL"

# ----------------------------- INSTALL ORDER ---------------------------------
# Order matters: torch cu130 first so nothing drags in a CPU/other-CUDA torch,
# then the SM89 flashinfer wheel so pip never fetches the official one, then
# the vLLM wheel (with deps), and transformers==5.8.1 LAST so that nothing
# can override the pin afterwards.
log "Installing torch 2.11.0 (cu130)"
"$VPY" -m pip install "$TORCH_SPEC" --index-url "$TORCH_INDEX"

log "Installing flashinfer-cubin"
"$VPY" -m pip install "$FI_CUBIN_SPEC"

log "Installing SM89 FlashInfer wheel"
"$VPY" -m pip install "$PROJECT_DIR/wheels/$FI_WHEEL"

log "Installing vLLM dev145 (SM89 golden release) + dependencies"
"$VPY" -m pip install "$PROJECT_DIR/wheels/$VLLM_WHEEL"

log "Ensuring deep_gemm is ABSENT (presence breaks SM89)"
"$VPY" -m pip uninstall -y deep-gemm deep_gemm deepgemm >/dev/null 2>&1 || true

log "Pinning transformers 5.8.1 (LAST install -- newer releases reject the checkpoint's hash_moe layer types)"
"$VPY" -m pip install "$TRANSFORMERS_SPEC"

# ----------------------------- SOURCE PATCH ----------------------------------
# BF16/unquantized-wo_a fallback for the fused o-projection. Required for the
# BF16 MTP drafter head; harmless otherwise (guarded, target path unchanged).
log "Writing and applying the BF16 o_proj patch"
cat > "$PROJECT_DIR/patch_dev145_bf16_oproj.py" <<'PATCH_EOF'
#!/usr/bin/env python3
"""
patch_dev145_bf16_oproj.py

Fixes: AttributeError: 'ColumnParallelLinear' object has no attribute
'weight_scale' in vllm/models/deepseek_v4/nvidia/ops/o_proj.py:68 when
running --speculative-config '{"method":"mtp",...}' on
canada-quant/DeepSeek-V4-Flash-W4A16-FP8-MTP with the dev145 SM89 wheel.

Root cause: deep_gemm_fp8_o_proj() unconditionally assumes an FP8-quantized
wo_a (reads weight_scale_inv / weight_scale). The target model's attention is
FP8 (compressed-tensors) and sails through; the checkpoint's MTP drafter head
is BF16/unquantized, so its wo_a is a plain ColumnParallelLinear with no
scale attribute.

Fix: insert an early branch for scale-less wo_a that replicates the fused
op's math in plain torch: inverse RoPE (interleaved even/odd pairs on the
LAST rope_dim dims of each head, rotation by -theta, mirroring
_fused_inv_rope_fp8_quant_per_head), grouped bf16 einsum against
wo_a.weight viewed [n_groups, o_lora_rank, heads_per_group*head_dim], then
wo_b.

Idempotent (marker-guarded); writes <file>.bak_bf16 once before modifying.
"""
import pathlib
import py_compile
import sys

MARKER = "SM89 patch: BF16/unquantized wo_a"

OLD = '''    Shared by the FlashMLA and FlashInfer CUDA backends. ``einsum_recipe`` /
    ``tma_aligned_scales`` come from ``compute_fp8_einsum_recipe``.
    """
    o_fp8, o_scale = fused_inv_rope_fp8_quant(
'''

NEW = '''    Shared by the FlashMLA and FlashInfer CUDA backends. ``einsum_recipe`` /
    ``tma_aligned_scales`` come from ``compute_fp8_einsum_recipe``.
    """
    if (
        getattr(wo_a, "weight_scale_inv", None) is None
        and getattr(wo_a, "weight_scale", None) is None
    ):
        # ---- SM89 patch: BF16/unquantized wo_a (e.g. the BF16 MTP drafter
        # head in canada-quant/DeepSeek-V4-Flash-W4A16-FP8-MTP). The fused
        # FP8 path below requires block scales; replicate its math in plain
        # torch: inverse RoPE (interleaved even/odd pairs on the LAST
        # ``rope_dim`` dims, rotation by -theta, mirroring
        # _fused_inv_rope_fp8_quant_per_head), grouped bf16 einsum, wo_b.
        num_tokens, num_heads, head_dim = o.shape
        of = o.to(torch.float32)
        cs = cos_sin_cache[positions.to(torch.long)]
        half = rope_dim // 2
        cos = cs[:, :half].unsqueeze(1)
        sin = cs[:, half:].unsqueeze(1)
        rope = of[..., nope_dim:]
        x1 = rope[..., 0::2]
        x2 = rope[..., 1::2]
        rope_inv = torch.stack(
            (x1 * cos + x2 * sin, x2 * cos - x1 * sin), dim=-1
        ).flatten(-2)
        o_inv = torch.cat((of[..., :nope_dim], rope_inv), dim=-1)
        o_grouped = o_inv.to(torch.bfloat16).reshape(
            num_tokens, n_groups, heads_per_group * head_dim
        )
        w = wo_a.weight.to(torch.bfloat16).view(
            n_groups, o_lora_rank, heads_per_group * head_dim
        )
        z_bf16 = torch.einsum("bhr,hdr->bhd", o_grouped, w)
        return wo_b(z_bf16.flatten(1))
    o_fp8, o_scale = fused_inv_rope_fp8_quant(
'''


def resolve_installed_target() -> pathlib.Path:
    import vllm  # noqa: PLC0415

    return (
        pathlib.Path(vllm.__file__).parent
        / "models" / "deepseek_v4" / "nvidia" / "ops" / "o_proj.py"
    )


def main() -> int:
    if len(sys.argv) > 1:
        target = pathlib.Path(sys.argv[1])
    else:
        target = resolve_installed_target()

    if not target.is_file():
        print(f"XX target not found: {target}", file=sys.stderr)
        return 2

    src = target.read_text()

    if MARKER in src:
        print(f"OK already patched, nothing to do: {target}")
        return 0

    if OLD not in src:
        print(
            "XX expected code block not found -- file differs from the "
            "dev145 (g8c631d45e) layout this patch targets. Refusing to "
            f"guess. File: {target}",
            file=sys.stderr,
        )
        return 2

    backup = target.with_suffix(target.suffix + ".bak_bf16")
    if not backup.exists():
        backup.write_text(src)
        print(f"OK backup written: {backup}")

    target.write_text(src.replace(OLD, NEW, 1))
    py_compile.compile(str(target), doraise=True)
    print(f"OK patched + compiled: {target}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
PATCH_EOF
"$VPY" "$PROJECT_DIR/patch_dev145_bf16_oproj.py"

# ----------------------------- VERIFICATION ----------------------------------
log "Verifying installed environment"
"$VPY" - <<'PYEOF'
import pathlib
import sys

import torch
assert torch.__version__.startswith("2.11.0"), f"torch is {torch.__version__}, expected 2.11.0 (a vLLM dependency probably replaced it -- reinstall torch with the cu130 index)"
assert torch.version.cuda == "13.0", f"torch CUDA is {torch.version.cuda}, expected 13.0"

import transformers
assert transformers.__version__ == "5.8.1", f"transformers is {transformers.__version__}, expected 5.8.1 (newer releases reject this checkpoint's hash_moe layer types) -- something installed after the pin overrode it"

import vllm
assert vllm.__version__.startswith("0.23.1rc1.dev145"), f"vLLM is {vllm.__version__}, expected dev145 (dev1018 is the broken release)"

from vllm.models.deepseek_v4.nvidia.ops import sm12x_deep_gemm_fallbacks as m
assert hasattr(m, "_fp8_paged_mqa_logits_sm12x"), "SM89 indexer fallbacks missing from this vLLM build"

o_proj = pathlib.Path(vllm.__file__).parent / "models/deepseek_v4/nvidia/ops/o_proj.py"
assert "SM89 patch: BF16/unquantized wo_a" in o_proj.read_text(), "BF16 o_proj patch not applied"

import flashinfer
try:
    import deep_gemm  # noqa: F401
    sys.exit("deep_gemm is importable -- it MUST NOT be installed on SM89")
except ImportError:
    pass

print(f"OK  torch {torch.__version__} | transformers {transformers.__version__} | vLLM {vllm.__version__} | flashinfer {getattr(flashinfer, '__version__', '?')} | deep_gemm absent | sm12x fallbacks present | bf16 o_proj patch applied")
if torch.cuda.is_available():
    for i in range(torch.cuda.device_count()):
        print(f"OK  GPU {i}: {torch.cuda.get_device_name(i)} capability {torch.cuda.get_device_capability(i)}")
PYEOF

# ----------------------------- GENERATE serve.sh -----------------------------
log "Writing $PROJECT_DIR/serve.sh"
cat > "$PROJECT_DIR/serve.sh" <<SERVE_EOF
#!/usr/bin/env bash
# Launch DeepSeek-V4-Flash with MTP speculative decoding on 4x SM89
# -- generated by setup_deepseek_v4_sm89.sh
#
#  * VLLM_TRITON_MLA_SPARSE=1 is LOAD-BEARING: it selects the portable Triton
#    sparse-MLA path, the only one validated on Ada. Do NOT add an
#    --attention-backend flag; the FlashInfer path crashes on SM89 (SWA cache
#    shape mismatch).
#  * 'env -u' scrubs stale profile exports that dev145 would otherwise read.
#  * CUDA is pinned to 13.0 because this machine has six toolkits and an apt
#    CUDA 12.0 nvcc in /usr/bin; unpinned JIT builds pick the wrong one.
#  * The model is resolved to a concrete snapshot DIRECTORY below instead of
#    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 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")"

MODEL="$MODEL"

# Resolve a repo id to its local snapshot directory. Prefer refs/main when it
# points at a snapshot that actually exists; otherwise fall back to the
# newest snapshot directory present.
resolve_model() {
    local m="\$1"
    if [ -d "\$m" ]; then printf '%s\n' "\$m"; return; fi
    local hub="\${HF_HOME:-\$HOME/.cache/huggingface}/hub"
    local repo_dir="\$hub/models--\${m//\//--}"
    local snaps="\$repo_dir/snapshots"
    if [ -d "\$snaps" ]; then
        local ref=""
        [ -f "\$repo_dir/refs/main" ] && ref="\$(cat "\$repo_dir/refs/main" 2>/dev/null || true)"
        if [ -n "\$ref" ] && [ -d "\$snaps/\$ref" ]; then
            printf '%s\n' "\$snaps/\$ref"; return
        fi
        local newest
        newest="\$(ls -1t "\$snaps" 2>/dev/null | head -1)"
        if [ -n "\$newest" ]; then printf '%s\n' "\$snaps/\$newest"; return; fi
    fi
    printf '\n'
}

MODEL_PATH="\$(resolve_model "\$MODEL")"
if [ -z "\$MODEL_PATH" ]; then
    echo "XX model '\$MODEL' not found locally. Download it once with:" >&2
    echo "   \$(dirname "\$0")/.venv/bin/huggingface-cli download \$MODEL" >&2
    exit 1
fi
echo "==> serving model from: \$MODEL_PATH"

exec env -u VLLM_TRITON_MLA_SPARSE_HEAD_BLOCK_SIZE \\
         -u VLLM_TRITON_MLA_SPARSE_TOPK_CHUNK_SIZE \\
         -u VLLM_TRITON_MLA_SPARSE_QUERY_CHUNK_SIZE \\
         -u VLLM_TRITON_MLA_SPARSE_MATMUL_DECODE \\
         -u VLLM_VERSION_OVERRIDE \\
         -u VLLM_PYTHON_EXECUTABLE \\
    VLLM_TRITON_MLA_SPARSE=1 \\
    CUDA_HOME=$CUDA_HOME_DIR \\
    CUDA_PATH=$CUDA_HOME_DIR \\
    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" \\
        --served-model-name "$SERVED_NAME" \\
        --host 0.0.0.0 --port $PORT \\
        --tensor-parallel-size $TP_SIZE \\
        --kv-cache-dtype fp8_ds_mla \\
        --block-size 256 \\
        --max-model-len $MAX_MODEL_LEN \\
        --gpu-memory-utilization $GPU_MEM_UTIL \\
        --max-num-seqs $MAX_NUM_SEQS \\
        --max-num-batched-tokens $MAX_NUM_BATCHED_TOKENS \\
        --tokenizer-mode deepseek_v4 \\
        --reasoning-parser deepseek_v4 \\
        --enable-auto-tool-choice --tool-call-parser deepseek_v4 \\
        --default-chat-template-kwargs '{"reasoning_effort": "$REASONING_EFFORT"}' \\
        --speculative-config '{"method":"mtp","num_speculative_tokens":$NUM_SPEC_TOKENS}' \\
        --trust-remote-code
SERVE_EOF
chmod +x "$PROJECT_DIR/serve.sh"

# ----------------------------- GENERATE canary.sh ----------------------------
log "Writing $PROJECT_DIR/canary.sh"
cat > "$PROJECT_DIR/canary.sh" <<CANARY_EOF
#!/usr/bin/env bash
# German coherence smoke test -- token salad here means broken numerics.
curl -s http://localhost:$PORT/v1/chat/completions \\
  -H 'Content-Type: application/json' \\
  -d '{"model":"$SERVED_NAME","messages":[{"role":"user","content":"Antworte in einem Satz: Was ist die Hauptstadt von Bayern?"}],"max_tokens":128}'
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
# systemd unit for the DeepSeek-V4-Flash vLLM server (MTP speculative
# decoding) -- generated by setup_deepseek_v4_sm89.sh. serve.sh is the single
# source of truth for environment variables and launch flags; this unit only
# 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
#   sudo systemctl enable --now deepseek-v4-flash
# Logs:
#   journalctl -fu deepseek-v4-flash

[Unit]
Description=vLLM OpenAI API server - DeepSeek-V4-Flash + MTP on 4x RTX 6000 Ada (SM89)
After=network-online.target
Wants=network-online.target
# A boot loop on a 159 GB model is expensive: allow 3 failed starts per 10 min,
# then stay down until 'systemctl reset-failed deepseek-v4-flash'.
StartLimitIntervalSec=600
StartLimitBurst=3

[Service]
Type=exec
User=$SERVICE_USER
Group=$SERVICE_USER
WorkingDirectory=$PROJECT_DIR

# Fail fast (and retry via Restart=) if the NVIDIA driver is not up yet,
# e.g. when racing device initialization right after boot.
ExecStartPre=/usr/bin/nvidia-smi

ExecStart=$PROJECT_DIR/serve.sh

# 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.
TimeoutStopSec=90
# SIGTERM to the API server first, then SIGKILL the whole cgroup
# (EngineCore + 4 worker processes) if anything lingers.
KillMode=mixed
LimitNOFILE=65535
LimitMEMLOCK=infinity

[Install]
WantedBy=multi-user.target
UNIT_EOF

# ----------------------------- DONE ------------------------------------------
log "Environment ready in $PROJECT_DIR"
cat <<NEXT_EOF

Next steps:
  1. Model weights: serve.sh resolves "$MODEL" from the local HF
     cache (or uses it directly if it is a path). To fetch it (~159 GB),
     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), then
     'Application startup complete'.
  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.
  5. Run as a service:      see header of
         $PROJECT_DIR/deepseek-v4-flash.service
NEXT_EOF