#!/usr/bin/env bash
# Once: python3 patch_sglang_eagle3.py
# Run:  MODEL=Qwen/Qwen3.5-27B-FP8 DRAFTER=./compressed ./serve_sglang.sh [STEPS TOPK NUM_DRAFT]
set -uo pipefail

MODEL="${MODEL:-Qwen/Qwen3.5-27B-FP8}"
DRAFTER="${DRAFTER:-./compressed}"
PORT="${PORT:-30000}"
STEPS="${1:-3}"
TOPK="${2:-1}"
NUM_DRAFT="${3:-4}"

SPEC=(--speculative-algorithm EAGLE3
      --speculative-draft-model-path "$DRAFTER"
      --speculative-num-steps "$STEPS"
      --speculative-eagle-topk "$TOPK"
      --speculative-num-draft-tokens "$NUM_DRAFT"
      --mamba-scheduler-strategy extra_buffer)
[ "$TOPK" -gt 1 ] && SPEC+=(--attention-backend flashinfer)
[ -x /usr/local/cuda/bin/ptxas ] && export TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas
export SGLANG_ENABLE_SPEC_V2=1

exec python3 -m sglang.launch_server \
  --model-path "$MODEL" \
  "${SPEC[@]}" \
  --dtype bfloat16 \
  --trust-remote-code \
  --mm-attention-backend sdpa \
  --host 0.0.0.0 --port "$PORT"