scripts/serve_vllm_optimized.sh

#!/usr/bin/env bash
# ============================================================
# serve_vllm_optimized.sh — vLLM with every RTX 3090 optimization
#
# This script applies ALL known optimizations for Ampere GPUs:
#   1. FP8 KV Cache           → -30% KV cache VRAM
#   2. FlashInfer backend     → faster prefill on SM 8.6
#   3. Chunked prefill        → better scheduling under load
#   4. Block size tuning      → larger blocks for long context
#   5. CUDA graph capture     → matches max_model_len
#   6. Swap space             → 4GB CPU overflow buffer
#   7. AWQ/GPTQ Marlin        → INT4 tensor core kernels
#   8. Prefix caching         → avoids recomputation
#
# Usage:
#   bash scripts/serve_vllm_optimized.sh [config_name]
#
# Configs:
#   awq-8k       — AWQ 4-bit, 8192 ctx, flashinfer, FP8 KV cache (★ recommended)
#   awq-16k      — AWQ 4-bit, 16384 ctx
#   awq-8k-fast  — AWQ 4-bit, 8192 ctx, max throughput tuning
#   fp8-8k       — FP8 weight, 8192 ctx
#   fp8-8k-kvc   — FP8 weight + FP8 KV cache (VRAM saving)
#   gptq-8k      — GPTQ Marlin, 8192 ctx
#   spec-8k      — speculative decoding with 0.6B draft model
# ============================================================
set -euo pipefail

CONFIG="${1:-awq-8k}"
PORT="${PORT:-8000}"
HOST="${HOST:-0.0.0.0}"

# --- Model paths ---
MODEL_AWQ="QuantTrio/Qwen3.6-27B-AWQ"
MODEL_FP8="Qwen/Qwen3.6-27B-FP8"
MODEL_GPTQ="groxaxo/Qwen3.6-27B-GPTQ-Pro-4bit"
MODEL_DRAFT="Qwen/Qwen2.5-0.5B-Instruct"  # for speculative decoding

# --- Config matrix ---
case "$CONFIG" in
    # ============================================================
    # AWQ — best performance/quality on Ampere (★ recommended)
    # ============================================================
    awq-8k)
        MODEL="$MODEL_AWQ"
        MAX_MODEL_LEN=8192
        GPU_MEM=0.90
        MAX_SEQS=8
        MAX_BATCHED=32768
        DTYPE="auto"
        ATTENTION="flashinfer"
        KV_CACHE_DTYPE="fp8"
        CHUNKED_PREFILL=true
        BLOCK_SIZE=16
        SWAP_SPACE=4
        SPEC_DRAFT=""
        ;;
    awq-8k-fast)
        MODEL="$MODEL_AWQ"
        MAX_MODEL_LEN=8192
        GPU_MEM=0.90
        MAX_SEQS=16
        MAX_BATCHED=65536
        DTYPE="auto"
        ATTENTION="flashinfer"
        KV_CACHE_DTYPE="fp8"
        CHUNKED_PREFILL=true
        BLOCK_SIZE=32
        SWAP_SPACE=8
        SPEC_DRAFT=""
        ;;
    awq-16k)
        MODEL="$MODEL_AWQ"
        MAX_MODEL_LEN=16384
        GPU_MEM=0.90
        MAX_SEQS=4
        MAX_BATCHED=32768
        DTYPE="auto"
        ATTENTION="flashinfer"
        KV_CACHE_DTYPE="fp8"
        CHUNKED_PREFILL=true
        BLOCK_SIZE=32
        SWAP_SPACE=8
        SPEC_DRAFT=""
        ;;
    # ============================================================
    # FP8 — lighter weight quant, tight on 24GB
    # ============================================================
    fp8-8k)
        MODEL="$MODEL_FP8"
        MAX_MODEL_LEN=8192
        GPU_MEM=0.88
        MAX_SEQS=4
        MAX_BATCHED=16384
        DTYPE="auto"
        ATTENTION="flashinfer"
        KV_CACHE_DTYPE="auto"
        CHUNKED_PREFILL=true
        BLOCK_SIZE=16
        SWAP_SPACE=4
        SPEC_DRAFT=""
        ;;
    fp8-8k-kvc)
        # FP8 weights + FP8 KV cache = tight fit but works
        MODEL="$MODEL_FP8"
        MAX_MODEL_LEN=8192
        GPU_MEM=0.85
        MAX_SEQS=6
        MAX_BATCHED=16384
        DTYPE="auto"
        ATTENTION="flashinfer"
        KV_CACHE_DTYPE="fp8"
        CHUNKED_PREFILL=true
        BLOCK_SIZE=16
        SWAP_SPACE=4
        SPEC_DRAFT=""
        ;;
    # ============================================================
    # GPTQ — INT4 via Marlin kernels (fast on Ampere)
    # ============================================================
    gptq-8k)
        MODEL="$MODEL_GPTQ"
        MAX_MODEL_LEN=8192
        GPU_MEM=0.90
        MAX_SEQS=8
        MAX_BATCHED=32768
        DTYPE="auto"
        ATTENTION="flashinfer"
        KV_CACHE_DTYPE="fp8"
        CHUNKED_PREFILL=true
        BLOCK_SIZE=16
        SWAP_SPACE=4
        SPEC_DRAFT=""
        ;;
    # ============================================================
    # Speculative decoding — Qwen 27B target + 0.6B draft
    # Spec can 1.5x-2x throughput on batch workloads
    # ============================================================
    spec-8k)
        MODEL="$MODEL_AWQ"
        MAX_MODEL_LEN=8192
        GPU_MEM=0.88
        MAX_SEQS=4
        MAX_BATCHED=16384
        DTYPE="auto"
        ATTENTION="flashinfer"
        KV_CACHE_DTYPE="fp8"
        CHUNKED_PREFILL=false
        BLOCK_SIZE=16
        SWAP_SPACE=4
        SPEC_DRAFT="$MODEL_DRAFT"
        ;;
    *)
        echo "❌ Unknown config: $CONFIG"
        echo ""
        echo "Available optimized configs:"
        echo "  awq-8k        ★ AWQ 4-bit, 8k ctx, all optimizations (recommended)"
        echo "  awq-8k-fast     AWQ 4-bit, 8k ctx, max throughput"
        echo "  awq-16k         AWQ 4-bit, 16k ctx"
        echo "  fp8-8k          FP8 weight, 8k ctx"
        echo "  fp8-8k-kvc      FP8 weight + FP8 KV cache"
        echo "  gptq-8k         GPTQ Marlin, 8k ctx"
        echo "  spec-8k         Speculative decoding (AWQ + 0.6B draft)"
        exit 1
        ;;
esac

echo "╔══════════════════════════════════════════════════════════╗"
echo "║  vLLM OPTIMIZED — Qwen SpeedLab                         ║"
echo "║  GPU: RTX 3090 (24 GB, SM 8.6)                          ║"
echo "╚══════════════════════════════════════════════════════════╝"
echo ""
echo "  Config:          $CONFIG"
echo "  Model:           $MODEL"
echo "  Max model len:   $MAX_MODEL_LEN"
echo "  GPU mem util:    $GPU_MEM"
echo "  Max seqs:        $MAX_SEQS"
echo "  Max batched:     $MAX_BATCHED"
echo "  Dtype:           $DTYPE"
echo "  Attention:       $ATTENTION"
echo "  KV cache dtype:  $KV_CACHE_DTYPE"
echo "  Chunked prefill: $CHUNKED_PREFILL"
echo "  Block size:      $BLOCK_SIZE"
echo "  Swap space:      ${SWAP_SPACE}GB"
if [[ -n "$SPEC_DRAFT" ]]; then
    echo "  Draft model:     $SPEC_DRAFT"
fi
echo ""

# --- Build vllm command ---
CMD=(
    vllm serve "$MODEL"
    --host "$HOST"
    --port "$PORT"
    --max-model-len "$MAX_MODEL_LEN"
    --gpu-memory-utilization "$GPU_MEM"
    --dtype "$DTYPE"
    --max-num-seqs "$MAX_SEQS"
    --max-num-batched-tokens "$MAX_BATCHED"
    --tensor-parallel-size 1
    --trust-remote-code
    --enable-prefix-caching
)

# Attention backend
if [[ "$ATTENTION" == "flashinfer" ]]; then
    CMD+=(--attention-backend flashinfer)
fi

# FP8 KV cache (biggest VRAM saving on Ampere)
if [[ "$KV_CACHE_DTYPE" == "fp8" ]]; then
    CMD+=(--kv-cache-dtype fp8)
fi

# Chunked prefill
if [[ "$CHUNKED_PREFILL" == "true" ]]; then
    CMD+=(--enable-chunked-prefill)
fi

# Block size
CMD+=(--block-size "$BLOCK_SIZE")

# Swap space (CPU offload safety net)
CMD+=(--swap-space "$SWAP_SPACE")

# CUDA graph: capture up to max_model_len (set equal to avoid recompilation)
CMD+=(--max-seq-len-to-capture "$MAX_MODEL_LEN")

# Speculative decoding
if [[ -n "$SPEC_DRAFT" ]]; then
    CMD+=(--speculative-model "$SPEC_DRAFT")
    # num_speculative_tokens: 5 is a good default for Qwen
    CMD+=(--num-speculative-tokens 5)
fi

echo "  Command:"
echo "    ${CMD[*]}"
echo ""
echo "  Log: /tmp/vllm_optimized_${CONFIG}.log"
echo ""

# Launch
"${CMD[@]}" 2>&1 | tee "/tmp/vllm_optimized_${CONFIG}.log"