#!/usr/bin/env bash
# Serve a Qwen3.6 / Qwen3.5 -35B-A3B 2-bit eschamoe export with escha-sglang, OpenAI-compatible,
# usable by any OpenAI-compatible client. eschamoe quant is auto-detected from config.json.
#
# Linux x86-64 + NVIDIA CUDA 12.x only (native or WSL2). See README.md.
#
# Required: set MODEL to your downloaded model dir (a flat folder of safetensors +
# config + tokenizer — the model repo has no nested subfolder). Everything else has
# safe defaults for a 24 GB card.
#   MODEL=./Qwen3.6-35B-A3B-Escha-W2 SERVED_NAME=escha-qwen36-35b-a3b-w2 \
#     bash serve.sh
set -euo pipefail

# ---- paths ----------------------------------------------------------------------------------
MODEL=${MODEL:?set MODEL=/path/to/your/eschamoe export dir}
SERVED_NAME=${SERVED_NAME:-escha-qwen36-35b-a3b-w2}   # opencode/LM Studio must use this id verbatim
VENV=${VENV:-}              # optional: path to your venv; if set, we activate it
ESCHA_SRC=${ESCHA_SRC:-}    # optional: path to escha *source* tree's src/ (only if not pip-installed)

# ---- server ---------------------------------------------------------------------------------
HOST=${HOST:-127.0.0.1}     # 0.0.0.0 to serve other machines — do this behind a VPN or tunnel and
                            #     set API_KEY; never open the raw port to the internet.
PORT=${PORT:-30000}
API_KEY=${API_KEY:-}        # REQUIRED if HOST=0.0.0.0 / exposed: clients must send this as the
                            #     Bearer key. Empty = no auth (fine only for localhost).
MEM=${MEM:-0.78}            # mem-fraction-static. 0.78 is the right default for a 24 GB card at the
                            #     CTXLEN below: mem-fraction-static sizes the pool the KV cache is
                            #     carved from, so too LOW is the common failure — 0.70 dies at startup
                            #     with "Not enough memory. Please try to increase --mem-fraction-static"
                            #     (measured, 4090 + CTXLEN 32768). Lower it only together with CTXLEN;
                            #     ~0.60-0.64 is fine for short-generation throughput work.
CTXLEN=${CTXLEN:-32768}     # KV is cheap on this hybrid model; raise once it works (see guide §2.5)
CHUNK=${CHUNK:-2048}        # chunked-prefill-size
MAXREQ=${MAXREQ:-0}         # max concurrent requests; 0 = auto. Lower (e.g. 4) frees Mamba pool for KV.
                            #     NOTE: this is CLAMPED by the mamba req-slot pool —
                            #     effective = min(MAXREQ, MAXMAMBA // ratio), ratio ~= 4 with the radix
                            #     cache ON (RADIX=1) and ~= 1 with RADIX=0. So MAXREQ=8 MAXMAMBA=8 RADIX=1
                            #     silently runs only ~2 concurrent. For bs 16: RADIX=0 MAXMAMBA=16, or keep
                            #     radix and set MAXMAMBA=64. A clamp is logged at startup ("max_num_reqs ...").
MAXMAMBA=${MAXMAMBA:-0}     # --max-mamba-cache-size; 0 = auto. Raise (e.g. 32) for deeper prefix cache /
                            #     to lift the MAXREQ clamp above.
MAMBA_RATIO=${MAMBA_RATIO:-} # --mamba-full-memory-ratio; empty = default (0.9). Lower (e.g. 0.6) frees the
                            #     pool for spec + CUDA-graph capture headroom. Set when SPEC=1 + GRAPHS=1.

# ---- behavior -------------------------------------------------------------------------------
GRAPHS=${GRAPHS:-1}        # 1 = CUDA graphs ON (DEFAULT, 2026-07-23). This hybrid MoE launches many
                           #     small kernels per token (30 GDN layers + 40 MoE layers + W8A16 GEMVs),
                           #     so eager decode is launch-bound: graphs are worth ~4.4x (measured
                           #     16 GB sm_120: 27.7 -> 122 tok/s bs1) and captured cleanly on every
                           #     card tested (sm_89 4090, sm_120 5060 Ti / 5090). Set GRAPHS=0 only
                           #     to debug a capture failure — expect ~4x slower decode there.
CUDA_GRAPH_BS=${CUDA_GRAPH_BS:-"1 2 4 8 16"}  # batch sizes to capture when GRAPHS=1. With SPEC=1 keep
                           #     max(bs) * DRAFT_TOKENS <= 64 (escham_moe verify-op limit) — e.g. "1 2 4 8" at
                           #     draft=4 for an 8-user pool.
THINK=${THINK:-1}          # 1 = thinking-ON + --reasoning-parser qwen3 (simplest, always works).
                           # 0 = thinking-OFF (faster; clean `content`): patches the chat
                           #     template default to non-thinking AND drops the reasoning parser.
                           # A per-request chat_template_kwargs:{"enable_thinking":false} is
                           #     honoured on a THINK=1 server — `content` is populated and tool
                           #     calls are parsed, so THINK=0 is not required for tool-calling
                           #     clients that disable thinking. With thinking ON, the reasoning
                           #     arrives in `reasoning_content` and the answer in `content`:
                           #     read both.
RADIX=${RADIX:-1}          # 1 = prefix caching ON (MambaRadixCache) — multi-turn agent reuse.
ATTN_BACKEND=${ATTN_BACKEND:-}  # full-attention backend override. Empty = sglang default (fine on
                           #     Ada/4090). REQUIRED on consumer Blackwell (RTX 50 / sm_120): the
                           #     default resolves to flashinfer but the fork asserts triton/trtllm_mha
                           #     /fa4 for hybrid-GDN on Blackwell -> set ATTN_BACKEND=triton.
TOOL_PARSER=${TOOL_PARSER:-qwen3_coder}   # matches <function=..>/<parameter=..> XML. Verify w/ curl.
REASONING_PARSER=${REASONING_PARSER:-qwen3}
ENABLE_CLP=${ENABLE_CLP:-0}   # 1 = --enable-custom-logit-processor (per-request thinking_budget cap etc.)
JSON_WS=${JSON_WS:-1}      # 1 = --constrained-json-disable-any-whitespace (DEFAULT).
                           #     With permissive whitespace, a json_schema `integer`/`number`
                           #     property that follows a string property can take a spurious
                           #     leading `-`: the grammar allows a "\n    " branch after the
                           #     colon, and from that off-distribution state the sign/digit
                           #     choice is near-tied. Measured on a 5080: 2-3 of 3 identical
                           #     greedy requests returned a NEGATIVE population for Tokyo --
                           #     schema-valid, so no client validation catches it. Disabling
                           #     any-whitespace removed it 0/6 and 0/3. Compact JSON is what
                           #     structured-output consumers want anyway. `JSON_WS=0` opts out.
DETERMINISTIC=${DETERMINISTIC:-0}  # 1 = --enable-deterministic-inference. Greedy is NOT
                           #     bit-reproducible across requests by default (fp16 accumulation
                           #     order): long reasoning chains diverge once a near-tie flips.
                           #     Set 1 for eval reproducibility; expect a throughput cost.
                           #     NOT AVAILABLE on consumer Blackwell (RTX 50 / sm_120): the
                           #     deterministic attention kernel asks for 104 KB of shared memory
                           #     per block, above the sm_120 limit, and the server exits at
                           #     startup. Works on Ampere / Ada / Hopper.

# ---- speculative decoding: native MTP / NEXTN (self-draft) ----------------------------------
# Large single-stream + low-concurrency speedup (bs1 ~1.55-1.77x; still +1.1x at 8 concurrent; turns
# net-NEGATIVE only above ~12 concurrent -> the right default for a <=8-person team). The draft is the
# model's OWN pretrained MTP head dequantized to fp16 (NOT a 2-bit quant — quantizing the draft collapses
# acceptance); accept-len ~3.5 at draft=4. Evidence: docs/SOTA/serving_perf_notes.md (2026-06-29).
SPEC=${SPEC:-0}            # 1 = enable MTP/NEXTN speculative decode. Needs DRAFT; pair with GRAPHS=1.
                           # NOTE: the published W2 model repo does NOT include an MTP draft export
                           #     (the mtp.* head in the checkpoint is not servable as a draft), so
                           #     SPEC=1 is unusable with that release as shipped — all published
                           #     serving numbers are non-speculative.
DRAFT=${DRAFT:-}          # MTP draft dir (fp16, the *-MTP-draft export). REQUIRED when SPEC=1.
STEPS=${STEPS:-3}        # --speculative-num-steps  (draft=4 / steps=3 is the measured sweet spot)
TOPK=${TOPK:-1}          # --speculative-eagle-topk
DRAFT_TOKENS=${DRAFT_TOKENS:-4}   # --speculative-num-draft-tokens (M = MAXREQ * this must stay <= 64)
MAMBA_SCHED=${MAMBA_SCHED:-}      # --mamba-scheduler-strategy. AUTO -> extra_buffer when SPEC=1 + radix ON
                                  #   (this hybrid-SSM arch REQUIRES extra_buffer + Spec-V2 to run radix +
                                  #   spec together; otherwise the default no_buffer is used).

# ---- environment the escha stack needs ------------------------------------------------------
[[ -n "$VENV" && -f "$VENV/bin/activate" ]] && source "$VENV/bin/activate"
PYBIN=$(command -v python)
TORCH_LIB=$("$PYBIN" -c 'import torch,os;print(os.path.join(os.path.dirname(torch.__file__),"lib"))')
export LD_LIBRARY_PATH="$TORCH_LIB${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"   # the reference codec_ext needs libc10
[[ -n "$ESCHA_SRC" ]] && export PYTHONPATH="$ESCHA_SRC${PYTHONPATH:+:$PYTHONPATH}"
export SGLANG_MAMBA_CONV_DTYPE=float16          # must match --dtype float16 or causal_conv1d crashes
export SGLANG_DISABLE_CUDNN_CHECK=1
# These qwen3_5_moe Escha checkpoints are TEXT-ONLY (the vision tower is in the quant `ignore`
# list — the checkpoint has zero visual.* tensors). Default to NOT instantiating the vision tower:
# it saves ~0.8 GB VRAM on every card, and prevents an image request from silently decoding through
# a random-init tower (garbage instead of an error). Override SGLANG_VLM_TEXT_ONLY=0 only if you
# ever serve a checkpoint that actually ships vision weights.
export SGLANG_VLM_TEXT_ONLY=${SGLANG_VLM_TEXT_ONLY:-1}
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True   # correct var name — cuts reserved-vs-allocated ~1.3GB
export PYTORCH_ALLOC_CONF=expandable_segments:True         # (alias on newer torch; keep both)
export HF_HUB_OFFLINE=1

# ---- int8-as-stored profile -----------------------------------------------------------------
# The runtime can keep lm_head/embed/attn/GDN/shared-expert INT8 in VRAM ("int8-as-stored",
# served via the w8a16 kernel) instead of dequantizing to fp16 at load. The right choice depends
# on CONCURRENCY, not just VRAM. At bs1 decode is memory-bandwidth-bound and int8 wins, by a
# card-dependent margin (measured 2026-07-27): +21–24% tok/s on an RTX 5090, +35–42% on an
# RTX 3090 — and −2.2 GB VRAM either way. Past the crossover (between concurrency 4 and 8)
# decode turns compute-bound and it inverts: −26% aggregate at bs16 on the 5090, to −52% on
# other cards. Pick with one knob:
#   INT8=auto (default) — VRAM-threshold rule: ON at ≤ SGLANG_INT8_AUTO_VRAM_GB (24), OFF above.
#   INT8=on   — force ON:  single-user / latency-sensitive serving on ANY card size
#               (e.g. a 32 GB RTX 5090 serving one user leaves ~20% single-stream on the
#               table under auto).
#   INT8=off  — force OFF: batched / high-concurrency throughput serving on >24 GB cards.
INT8=${INT8:-auto}
case "$INT8" in
  on)  for _v in LM_HEAD EMBED ATTN GDN SHEXP; do export "SGLANG_INT8_${_v}=1"; done ;;
  off) for _v in LM_HEAD EMBED ATTN GDN SHEXP; do export "SGLANG_INT8_${_v}=0"; done ;;
  auto) : ;;   # leave the five SGLANG_INT8_* vars as the caller set them (or unset -> AUTO)
  *) echo "[serve] ERROR: INT8 must be auto|on|off (got '$INT8')" >&2; exit 1 ;;
esac

# The fused MoE decode kernel is the default, always-preferred path (auto-detected; falls back
# to the slower per-expert sorted dispatch only if the fused op/kernel build is unavailable).
# DEFAULT-ON fail-loud (2026-07-23): the silent fallback is ~40x slower — exactly what a
# benchmarking or production user must not hit unknowingly. Set ESCHA_MOE_REQUIRE_FUSED=0
# only while debugging a box where the fused path genuinely cannot run.
# Renamed 2026-07-25: DI_ESCHA_MOE_REQUIRE_FUSED -> ESCHA_MOE_REQUIRE_FUSED (every public
# knob is ESCHA_*). The old name is still honored, with a notice, so existing launch
# scripts keep working.
if [[ -n "${DI_ESCHA_MOE_REQUIRE_FUSED:-}" && -z "${ESCHA_MOE_REQUIRE_FUSED:-}" ]]; then
  echo "[serve] DI_ESCHA_MOE_REQUIRE_FUSED is deprecated; use ESCHA_MOE_REQUIRE_FUSED"
  ESCHA_MOE_REQUIRE_FUSED="$DI_ESCHA_MOE_REQUIRE_FUSED"
fi
export ESCHA_MOE_REQUIRE_FUSED=${ESCHA_MOE_REQUIRE_FUSED:-1}
# Mixed-bit (K2 gate_up / K3 down) fused epilogue. Worth +10% single-user decode
# on this model — measured 2026-07-25 on an RTX 4090 (sm_89) with the shipped
# wheel: 207.4-208.6 -> 228.5-228.8 tok/s at bs1, identical answers (arithmetic,
# recall, long-context needle). It was opt-in only so that older kernel builds
# without the capability probe would not trip; the probe
# (`escham_moe_prehad_k3_available`) is present in this build and auto-falls-back
# when absent, so ON is the right default. `ESCHA_ENABLE_K3_FUSED_EPI=0` opts out.
export ESCHA_ENABLE_K3_FUSED_EPI=${ESCHA_ENABLE_K3_FUSED_EPI:-1}
# Other passthrough knobs: ESCHA_FORCE_SORTED_MOE=1 forces the slower validated fallback
# (debug/parity only); ESCHA_ROUTE=lovelace|blackwell overrides the kernel's auto-selected
# launch geometry (bit-identical either way; blackwell can win even on pre-sm_100 cards).
# ptxas discovery (rewritten 2026-07-25 after two field reports).
#
#  * Triton — which torch already pulls in — SHIPS a matching ptxas at
#    site-packages/triton/backends/nvidia/bin/ptxas. Probe that FIRST: a CUDA
#    toolkit is NOT a requirement of this runtime, and the old error message
#    sent users off to install one they did not need.
#  * Validate whatever the caller exported. A common shell profile sets
#    TRITON_PTXAS_PATH to the CUDA *bin directory*; `[[ -x <dir> ]]` is TRUE for
#    directories, so the old check accepted it and triton then died mid-startup
#    with `PermissionError: [Errno 13] Permission denied: /usr/local/cuda-12.8/bin`.
#    Require a regular file (-f) AND executable (-x); if a directory was given,
#    try <dir>/ptxas before discarding it.
_ptxas_ok() { [[ -n "${1:-}" && -f "$1" && -x "$1" ]]; }
if ! _ptxas_ok "${TRITON_PTXAS_PATH:-}"; then
  if [[ -n "${TRITON_PTXAS_PATH:-}" ]]; then
    if _ptxas_ok "${TRITON_PTXAS_PATH%/}/ptxas"; then
      TRITON_PTXAS_PATH="${TRITON_PTXAS_PATH%/}/ptxas"
      echo "[serve] TRITON_PTXAS_PATH was a directory; using $TRITON_PTXAS_PATH"
    else
      echo "[serve] ignoring TRITON_PTXAS_PATH='$TRITON_PTXAS_PATH' (not an executable file)"
      TRITON_PTXAS_PATH=""
    fi
  fi
  if [[ -z "${TRITON_PTXAS_PATH:-}" ]]; then
    for p in \
      "$(python -c 'import os,triton;print(os.path.join(os.path.dirname(triton.__file__),"backends","nvidia","bin","ptxas"))' 2>/dev/null)" \
      /usr/local/cuda*/bin/ptxas \
      "$(command -v ptxas 2>/dev/null)"; do
      _ptxas_ok "$p" && { TRITON_PTXAS_PATH="$p"; break; }
    done
  fi
  export TRITON_PTXAS_PATH
fi
: "${TRITON_PTXAS_PATH:?no usable ptxas found. It normally ships inside the venv with triton (a torch dependency) — check that \`python -c \"import triton\"\` works. Otherwise point TRITON_PTXAS_PATH at a ptxas BINARY (not its directory).}"

# Only if you hit it: the first time triton compiles a kernel (during warmup AND during CUDA-graph
# capture), the multi-GB scheduler must fork() to spawn ptxas. Under the default overcommit heuristic
# that fork can fail -> the server dies with a triton JIT "Cannot allocate memory" crash. Setting
# vm.overcommit_memory=1 fixes it. Many boxes never hit this (a warm triton cache, or enough free RAM
# to satisfy the fork), so it is advisory, not mandatory. In a container the sysctl is often read-only
# ("sysctl: setting key ...: Read-only file system") — then it is the HOST's job, not something you
# can set from inside; the server usually starts fine there anyway.
# (Advisory phrased WITHOUT the literal errno string so log-watchers grepping for the real
# crash message don't false-positive on this note — friction 2026-07-23 F4.)
if [[ "$(cat /proc/sys/vm/overcommit_memory 2>/dev/null)" == "0" ]]; then
  echo "[serve] note: if startup dies with a triton JIT memory-allocation failure at kernel-" >&2
  echo "[serve]       compile time, set 'sudo sysctl -w vm.overcommit_memory=1' (persist via" >&2
  echo "[serve]       /etc/sysctl.d/99-escha.conf). In a read-only container this must be set" >&2
  echo "[serve]       on the host; otherwise it's safe to ignore." >&2
fi

# Warn when the GPU is already substantially occupied (a foreign server / another model on a
# shared box is the #1 cause of an immediate, confusing pool-allocation OOM at startup).
_USED_MB=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null | head -1 | tr -dc '0-9')
if [[ -n "$_USED_MB" && "$_USED_MB" -gt 1024 ]]; then
  echo "[serve] WARNING: GPU already has ${_USED_MB} MiB in use by other processes (nvidia-smi)." >&2
  echo "[serve]          The static pool sizing assumes a free GPU — expect an allocation failure" >&2
  echo "[serve]          if another server is resident. Check: nvidia-smi" >&2
fi

# ---- assemble flags -------------------------------------------------------------------------
GRAPH_ARGS=(--disable-piecewise-cuda-graph)
[[ "$GRAPHS" == "1" ]] && GRAPH_ARGS+=(--cuda-graph-bs $CUDA_GRAPH_BS) || GRAPH_ARGS+=(--disable-cuda-graph)

POOL_ARGS=()
[[ "$MAXREQ"   != "0" ]] && POOL_ARGS+=(--max-running-requests "$MAXREQ")
[[ "$MAXMAMBA" != "0" ]] && POOL_ARGS+=(--max-mamba-cache-size "$MAXMAMBA")
[[ -n "$MAMBA_RATIO" ]] && POOL_ARGS+=(--mamba-full-memory-ratio "$MAMBA_RATIO")
RADIX_ARGS=(); [[ "$RADIX" == "0" ]] && RADIX_ARGS+=(--disable-radix-cache)
ATTN_ARGS=(); [[ -n "$ATTN_BACKEND" ]] && ATTN_ARGS+=(--attention-backend "$ATTN_BACKEND")
AUTH_ARGS=(); [[ -n "$API_KEY" ]] && AUTH_ARGS+=(--api-key "$API_KEY")

# speculative MTP/NEXTN draft (the model's own self-draft head)
SPEC_ARGS=()
if [[ "$SPEC" == "1" ]]; then
  : "${DRAFT:?SPEC=1 requires DRAFT=/path/to/<model>-MTP-draft (the fp16 NEXTN head, NOT a 2-bit quant)}"
  [[ -d "$DRAFT" ]] || { echo "[serve] ERROR: SPEC=1 but DRAFT dir not found: $DRAFT" >&2; exit 1; }
  SPEC_ARGS+=(--speculative-algorithm NEXTN --speculative-draft-model-path "$DRAFT"
              --speculative-num-steps "$STEPS" --speculative-eagle-topk "$TOPK"
              --speculative-num-draft-tokens "$DRAFT_TOKENS")
  if [[ "$GRAPHS" == "1" ]]; then   # enforce the escham_moe verify-op limit: max running batch * draft <= 64
    maxbs=0; for b in $CUDA_GRAPH_BS; do (( b > maxbs )) && maxbs=$b; done
    if (( maxbs * DRAFT_TOKENS > 64 )); then
      echo "[serve] WARNING: max cuda-graph bs ($maxbs) * draft_tokens ($DRAFT_TOKENS) > 64 — escham_moe only" >&2
      echo "[serve]          captures M<=64; trim CUDA_GRAPH_BS (e.g. \"1 2 4 8\" at draft=4)." >&2
    fi
  fi
  # radix cache + spec on this hybrid (mamba) arch needs the extra-buffer scheduler + Spec-V2.
  if [[ "$RADIX" != "0" ]]; then
    export SGLANG_ENABLE_SPEC_V2=1
    MAMBA_SCHED=${MAMBA_SCHED:-extra_buffer}
    echo "[serve] radix+spec: exporting SGLANG_ENABLE_SPEC_V2=1, --mamba-scheduler-strategy $MAMBA_SCHED" >&2
  fi
fi
SCHED_ARGS=(); [[ -n "$MAMBA_SCHED" ]] && SCHED_ARGS+=(--mamba-scheduler-strategy "$MAMBA_SCHED")
CLP_ARGS=(); [[ "$ENABLE_CLP" == "1" ]] && CLP_ARGS+=(--enable-custom-logit-processor)
if [[ "$HOST" != "127.0.0.1" && -z "$API_KEY" ]]; then
  echo "[serve] WARNING: HOST=$HOST is exposed but API_KEY is empty — no auth. Set API_KEY." >&2
fi

REASON_ARGS=(); TEMPLATE_ARGS=()
if [[ "$THINK" == "1" ]]; then
  REASON_ARGS+=(--reasoning-parser "$REASONING_PARSER")
else
  # thinking-OFF: generate a patched chat template (default-closed <think>) from the model's own.
  NOTHINK="${TMPDIR:-/tmp}/escha_nothink_$(basename "$MODEL").jinja"
  "$PYBIN" - "$MODEL" "$NOTHINK" <<'PY'
import json, os, sys
mdir, out = sys.argv[1], sys.argv[2]
src = os.path.join(mdir, "chat_template.jinja")
tmpl = open(src).read() if os.path.exists(src) else json.load(open(os.path.join(mdir,"tokenizer_config.json")))["chat_template"]
# Flip the qwen "enable_thinking" default from ON to OFF. Two common phrasings handled:
cands = [
 ("{%- if enable_thinking is defined and enable_thinking is false %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- else %}\n        {{- '<think>\\n' }}\n    {%- endif %}",
  "{%- if enable_thinking is defined and enable_thinking is true %}\n        {{- '<think>\\n' }}\n    {%- else %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- endif %}"),
 # froggeric/Qwen-Fixed-Chat-Templates (v21.x) hoists the toggle into a single
 # default-assignment near the top instead of an inline if/else at the generation
 # prompt, so the block above is absent and THINK=0 has nothing to patch. Flipping the
 # default is exactly equivalent here (ns_state.thinking is seeded from it).
 ("{%- set enable_thinking = enable_thinking if enable_thinking is defined else true %}",
  "{%- set enable_thinking = enable_thinking if enable_thinking is defined else false %}"),
]
patched = tmpl
for old, new in cands:
    if old in patched: patched = patched.replace(old, new); break
else:
    sys.exit("THINK=0: could not find the enable_thinking block to patch in this model's chat "
             "template. Serve with THINK=1, or hand-patch the template and pass it via --chat-template.")
open(out, "w").write(patched)
print(out)
PY
  TEMPLATE_ARGS+=(--chat-template "$NOTHINK")   # no reasoning parser when thinking is pre-closed
fi

echo "[serve] MODEL=$MODEL NAME=$SERVED_NAME $HOST:$PORT MEM=$MEM CTX=$CTXLEN GRAPHS=$GRAPHS THINK=$THINK RADIX=$RADIX SPEC=$SPEC"
echo "[serve] tool-call-parser=$TOOL_PARSER  ptxas=$TRITON_PTXAS_PATH"
[[ "$SPEC" == "1" ]] && echo "[serve] MTP/NEXTN: draft=$DRAFT steps=$STEPS topk=$TOPK draft_tokens=$DRAFT_TOKENS cuda_graph_bs='$CUDA_GRAPH_BS'"

JSON_ARGS=()
[[ "$JSON_WS" == "1" ]] && JSON_ARGS+=(--constrained-json-disable-any-whitespace)
DET_ARGS=()
[[ "$DETERMINISTIC" == "1" ]] && DET_ARGS+=(--enable-deterministic-inference)

exec python -m sglang.launch_server \
  --model-path "$MODEL" \
  --served-model-name "$SERVED_NAME" \
  --host "$HOST" --port "$PORT" \
  --dtype float16 \
  --mem-fraction-static "$MEM" \
  --context-length "$CTXLEN" \
  --chunked-prefill-size "$CHUNK" \
  --sampling-backend pytorch \
  --trust-remote-code \
  --tool-call-parser "$TOOL_PARSER" \
  "${REASON_ARGS[@]}" \
  "${TEMPLATE_ARGS[@]}" \
  "${POOL_ARGS[@]}" \
  "${RADIX_ARGS[@]}" \
  "${ATTN_ARGS[@]}" \
  "${AUTH_ARGS[@]}" \
  "${SPEC_ARGS[@]}" \
  "${SCHED_ARGS[@]}" \
  "${CLP_ARGS[@]}" \
  "${GRAPH_ARGS[@]}" \
  "${JSON_ARGS[@]}" \
  "${DET_ARGS[@]}" \
  --log-level info \
  "$@"          # forward any extra sglang flag verbatim