#!/usr/bin/env bash
# Serve Qwen3.6-35B-A3B-FP8 with DFlash speculative decoding via vLLM on NVIDIA RTX 3090
# Uses Podman container with CUDA support
# 4 parallel streams, full context (262K) each
set -euo pipefail

# Configuration
MODEL="${MODEL:-Qwen/Qwen3.6-35B-A3B-FP8}"
DRAFT_MODEL="${DRAFT_MODEL:-z-lab/Qwen3.6-35B-A3B-DFlash}"
PORT="${PORT:-8000}"
GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.95}"
MAX_NUM_SEQS="${MAX_NUM_SEQS:-4}"
MAX_MODEL_LEN="${MAX_MODEL_LEN:-262144}"
KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-fp8}"
NUM_SPEC_TOKENS="${NUM_SPEC_TOKENS:-8}"
CONTAINER_NAME="${CONTAINER_NAME:-vllm-qwen36-dflash}"
HF_CACHE="${HF_CACHE:-/home/dank/.cache/huggingface}"
GPU_LIBS="${GPU_LIBS:-/usr/lib/x86_64-linux-gnu}"

echo "Starting vLLM server in Podman..."
echo "  Target model: ${MODEL}"
echo "  Draft model:  ${DRAFT_MODEL}"
echo "  Port:         ${PORT}"
echo "  GPU mem util: ${GPU_MEM_UTIL}"
echo "  Max seqs:     ${MAX_NUM_SEQS}"
echo "  Context len:  ${MAX_MODEL_LEN}"
echo "  KV cache:     ${KV_CACHE_DTYPE}"
echo "  Spec tokens:  ${NUM_SPEC_TOKENS}"
echo ""

# Remove existing container if present
podman rm -f "${CONTAINER_NAME}" 2>/dev/null || true

# Run vLLM in Podman container (foreground mode for systemd)
# Note: vllm/vllm-openai image has entrypoint "vllm serve", so we pass args directly
exec podman run --name "${CONTAINER_NAME}" \
  --device /dev/nvidia0 \
  --device /dev/nvidiactl \
  --device /dev/nvidia-uvm \
  --device /dev/nvidia-uvm-tools \
  --device /dev/nvidia-modeset \
  -v "${GPU_LIBS}:/usr/lib/x86_64-linux-gnu:ro" \
  -v "${HF_CACHE}:/root/.cache/huggingface" \
  -p "${PORT}:${PORT}" \
  docker.io/vllm/vllm-openai:latest \
  "${MODEL}" \
    --port "${PORT}" \
    --host "0.0.0.0" \
    --kv-cache-dtype "${KV_CACHE_DTYPE}" \
    --gpu-memory-utilization "${GPU_MEM_UTIL}" \
    --max-num-seqs "${MAX_NUM_SEQS}" \
    --max-model-len "${MAX_MODEL_LEN}" \
    --spec-method dflash \
    --spec-model "${DRAFT_MODEL}" \
    --spec-tokens "${NUM_SPEC_TOKENS}" \
    --trust-remote-code