llm-server-mac

#!/bin/bash
#
# llm-server-mac — Smart launcher for llama.cpp on macOS (Apple Silicon)
# Auto-detects Metal GPU, unified memory, and launches with optimal settings.
#
# Usage:
#   llm-server-mac <repo/name> --download    # Download from HuggingFace
#   llm-server-mac -h                        # Show this help message
#   llm-server-mac <model.gguf>              # Launch with auto-tuning
#   llm-server-mac <model.gguf> --dry-run    # Print command without executing
#   llm-server-mac --port 8082 <model.gguf>  # Custom port
#   llm-server-mac <model.gguf> --benchmark  # Start, run quick benchmark, print tok/s
#   llm-server-mac <model.gguf> --cpu        # Force CPU-only (ignore Metal GPU)
#   llm-server-mac <model.gguf> --server-bin /path/to/llama-server
#   llm-server-mac <model.gguf> --kv-quality high  # KV cache: high(f16), mid(q8_0), low(q4_0)
#
# Environment variables:
#   LLAMA_SERVER    Path to llama-server binary (auto-detected if not set)
#   LLM_MODEL_DIR   Default model directory (default: ~/ai_models)
#   LLM_PORT        Server port (default: 8081)
#   LLM_CTX_SIZE    Context size (default: 32768)

set -euo pipefail

# ═══════════════════════════════════════════════════════════════
# Section 1: Configuration & Argument Parsing
# ═══════════════════════════════════════════════════════════════

VERSION="2.0.0"
PORT="${LLM_PORT:-8081}"
HOST="0.0.0.0"
# macOS default: 32K context (unified memory is shared, be conservative)
CTX_SIZE="${LLM_CTX_SIZE:-32768}"

# Persistent Configuration
CONFIG_DIR="$HOME/.config/llm-server"
CONFIG_FILE="$CONFIG_DIR/config.sh"
[[ -f "$CONFIG_FILE" ]] && source "$CONFIG_FILE"

MODEL_DIR="${LLM_MODEL_DIR:-$HOME/ai_models}"
mkdir -p "$MODEL_DIR"
CACHE_DIR="$HOME/.cache/llm-server"
SERVER_LOG="${TMPDIR:-/tmp}/llm-server.log"

RETUNE=0
DOWNLOAD=0
DRY_RUN=0
VERBOSE=0
BENCHMARK=0
CPU_ONLY=0
MODEL_ARG=""
USER_SERVER_BIN=""
KV_QUALITY="mid"  # high (f16), mid (q8_0, default), low (q4_0)

# Cleanup on exit
cleanup() { :; }
trap cleanup EXIT

# Auto-detect llama-server binary
find_server_binary() {
    if [[ -n "${LLAMA_SERVER:-}" && -x "$LLAMA_SERVER" ]]; then
        echo "$LLAMA_SERVER"
        return
    fi
    local candidates=(
        "$HOME/llama.cpp/build/bin/llama-server"
        "$HOME/ik_llama.cpp/build/bin/llama-server"
        "/opt/homebrew/bin/llama-server"
        "/usr/local/bin/llama-server"
        "$(command -v llama-server 2>/dev/null || true)"
    )
    for bin in "${candidates[@]}"; do
        [[ -n "$bin" && -x "$bin" ]] && echo "$bin" && return
    done
    echo ""
}

LLAMA_SERVER="$(find_server_binary)"

# Parse arguments
while [[ $# -gt 0 ]]; do
    case "$1" in
        --retune)    RETUNE=1; shift ;;
        --dry-run)   DRY_RUN=1; shift ;;
        --verbose)   VERBOSE=1; shift ;;
        --benchmark) BENCHMARK=1; shift ;;
        --cpu)       CPU_ONLY=1; shift ;;
        --port)      PORT="$2"; shift 2 ;;
        --ctx-size)  CTX_SIZE="$2"; shift 2 ;;
        --model-dir)   MODEL_DIR="$2"; shift 2 ;;
        --server-bin)  USER_SERVER_BIN="$2"; shift 2 ;;
        --kv-quality)  KV_QUALITY="$2"; shift 2 ;;
        --download)  DOWNLOAD=1; shift ;;
        --version)     echo "llm-server-mac v$VERSION"; exit 0 ;;
        --help|-h)
            sed -n '2,/^$/s/^# \?//p' "$0"
            exit 0
            ;;
        *)           MODEL_ARG="$1"; shift ;;
    esac
done

if [[ -z "$MODEL_ARG" ]]; then
    echo "Error: No model specified. Usage: llm-server-mac <model.gguf> or <repo/name> --download"
    exit 1
fi

# ── Download Logic ──
if (( DOWNLOAD )); then
    # macOS: Metal GPU memory = system unified memory
    TOTAL_VRAM_MB=0
    if (( ! CPU_ONLY )); then
        # On Apple Silicon, GPU uses unified memory — report total system RAM
        TOTAL_VRAM_MB=$(( $(sysctl -n hw.memsize) / 1048576 ))
    fi
    RAM_AVAIL_MB=$(vm_stat | awk '/Pages free/ {gsub(/\./,"",$3); free=$3} /Pages inactive/ {gsub(/\./,"",$3); inactive=$3} END {printf "%.0f", (free+inactive)*4096/1048576}')

    # Find the downloader
    SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
    DOWNLOADER="$SCRIPT_DIR/download_any_gguf.py"
    if [[ ! -f "$DOWNLOADER" ]]; then
        DOWNLOADER="$MODEL_DIR/download_any_gguf.py"
    fi

    if [[ ! -f "$DOWNLOADER" ]]; then
        echo "Error: Downloader script (download_any_gguf.py) not found."
        exit 1
    fi

    echo "Launching GGUF Downloader for: $MODEL_ARG"
    echo "Wait for the interactive session..."
    echo ""

    python3 "$DOWNLOADER" --repo "$MODEL_ARG" --dir "$MODEL_DIR" --vram "$TOTAL_VRAM_MB" --ram "$RAM_AVAIL_MB"

    echo ""
    echo "Download step complete."
    echo "To launch your new model, run: llm-server-mac <path_to_downloaded_file>"
    exit 0
fi

if [[ ! -f "$MODEL_ARG" ]]; then
    if [[ -f "$MODEL_DIR/$MODEL_ARG" ]]; then
        MODEL_ARG="$MODEL_DIR/$MODEL_ARG"
    else
        echo "Error: Model not found: $MODEL_ARG"
        exit 1
    fi
fi

MODEL_PATH="$MODEL_ARG"
MODEL_NAME="$(basename "$MODEL_PATH")"

# Apply --server-bin override
if [[ -n "$USER_SERVER_BIN" ]]; then
    if [[ ! -x "$USER_SERVER_BIN" ]]; then
        echo "Error: --server-bin binary not found or not executable: $USER_SERVER_BIN"
        exit 1
    fi
    LLAMA_SERVER="$USER_SERVER_BIN"
fi

if [[ -z "$LLAMA_SERVER" ]]; then
    echo "Error: llama-server binary not found."
    echo "Install llama.cpp: brew install llama.cpp"
    echo "Or build from source: https://github.com/ggml-org/llama.cpp"
    exit 1
fi

# Detect backend
IS_IK_LLAMA=0
if [[ "$LLAMA_SERVER" == *ik_llama* ]]; then
    IS_IK_LLAMA=1
elif timeout 5 "$LLAMA_SERVER" --help 2>&1 | grep -q 'graph' 2>/dev/null; then
    IS_IK_LLAMA=1
fi

log() { [[ "$VERBOSE" == "1" ]] && echo "[DEBUG] $*" >&2 || true; }

# ═══════════════════════════════════════════════════════════════
# Section 2: Hardware Detection (macOS)
# ═══════════════════════════════════════════════════════════════

echo "═══ llm-server v$VERSION (macOS) ═══"
echo "Binary: $LLAMA_SERVER"
(( IS_IK_LLAMA )) && echo "Backend: ik_llama.cpp" || echo "Backend: llama.cpp"

# Detect CPU
PHYSICAL_CORES=$(sysctl -n hw.physicalcpu 2>/dev/null || echo 4)
PERF_CORES=$(sysctl -n hw.perflevel0.physicalcpu 2>/dev/null || echo "$PHYSICAL_CORES")
# machdep.cpu.brand_string is Intel-only; Apple Silicon uses chip name from sysctl
CPU_BRAND=$(sysctl -n machdep.cpu.brand_string 2>/dev/null || sysctl -n hw.model 2>/dev/null || echo "Apple Silicon")

# Detect RAM (unified memory on Apple Silicon)
RAM_TOTAL_MB=$(( $(sysctl -n hw.memsize) / 1048576 ))
# Available memory: free + inactive pages (purgeable cache)
RAM_AVAIL_MB=$(vm_stat | awk '/Pages free/ {gsub(/\./,"",$3); free=$3} /Pages inactive/ {gsub(/\./,"",$3); inactive=$3} END {printf "%.0f", (free+inactive)*4096/1048576}')

echo "CPU: $CPU_BRAND (${PHYSICAL_CORES} cores, ${PERF_CORES} performance)"
echo "RAM: ${RAM_AVAIL_MB}MB available / ${RAM_TOTAL_MB}MB total (unified memory)"

# Detect Metal GPU
HAS_METAL=0
METAL_GPU_NAME=""
if (( ! CPU_ONLY )); then
    # Check if Metal is available via system_profiler
    # Works for all Apple Silicon: M1, M2, M3, M4 (and Pro/Max/Ultra variants)
    METAL_INFO=$(system_profiler SPDisplaysDataType 2>/dev/null || true)
    if echo "$METAL_INFO" | grep -qi "Metal Support\|Metal Family\|Apple M"; then
        HAS_METAL=1
        METAL_GPU_NAME=$(echo "$METAL_INFO" | awk -F': ' '/Chipset Model:/ {print $2; exit}')
        METAL_GPU_NAME=${METAL_GPU_NAME:-"Apple Silicon GPU"}
        echo "GPU: $METAL_GPU_NAME (Metal, unified memory)"
    fi
fi

if (( ! HAS_METAL )); then
    if (( CPU_ONLY )); then
        echo "GPU: skipped (--cpu flag)"
    else
        echo "GPU: none detected (CPU-only mode)"
    fi
fi

# On Apple Silicon, GPU and CPU share the same unified memory pool.
# The "VRAM" is the same as RAM — no separate budget needed.
# We treat the full system RAM as the memory budget.
GPU_COUNT=0
if (( HAS_METAL )); then
    GPU_COUNT=1
fi

# ═══════════════════════════════════════════════════════════════
# Section 3: Model Detection
# ═══════════════════════════════════════════════════════════════

# Get total model size (handles split GGUFs)
MODEL_DIR_PATH=$(dirname "$MODEL_PATH")
SPLIT_PATTERN=$(echo "$MODEL_NAME" | sed 's/-[0-9]*-of-[0-9]*\.gguf/.gguf/')

# macOS stat syntax differs from Linux
TOTAL_SIZE_BYTES=0
while IFS= read -r f; do
    sz=$(stat -f%z "$f" 2>/dev/null || stat --printf="%s" "$f" 2>/dev/null || echo 0)
    TOTAL_SIZE_BYTES=$(( TOTAL_SIZE_BYTES + sz ))
done < <(find "$MODEL_DIR_PATH" -name "${SPLIT_PATTERN%.gguf}*" -type f 2>/dev/null)

if (( TOTAL_SIZE_BYTES == 0 )); then
    TOTAL_SIZE_BYTES=$(stat -f%z "$MODEL_PATH" 2>/dev/null || stat --printf="%s" "$MODEL_PATH" 2>/dev/null || echo 0)
fi

TOTAL_SIZE_MB=$(awk "BEGIN {printf \"%.0f\", $TOTAL_SIZE_BYTES/1048576}")
TOTAL_SIZE_GB=$(awk "BEGIN {printf \"%.1f\", $TOTAL_SIZE_BYTES/1073741824}")

# Read GGUF metadata (same Python parser — cross-platform)
read LAYER_COUNT EXPERT_COUNT HEAD_COUNT_KV KEY_LENGTH VALUE_LENGTH HAS_SSM HAS_FUSED EXPERT_BYTES NON_EXPERT_BYTES < <(python3 -c "
import struct, sys, glob, os, re

# GGUF type sizes: (bytes_per_block, elements_per_block)
TYPE_SIZES = {
    0:(4,1), 1:(2,1), 2:(2+16,32), 3:(2+16,32), 4:(2+16,32), 5:(2+16,32),
    6:(4+32,32), 7:(2+32,32), 8:(2+32,32), 9:(4+32,32), 10:(2+64,64),
    11:(2+16,16), 12:(2+64,64), 13:(4+64,64), 14:(2+256,256),
    15:(2+256,256), 16:(2+256,256), 17:(2+256,256), 18:(4+256,256),
    19:(4+256,256), 20:(2+256,256), 21:(4+256,256), 22:(2+128,256),
    23:(2+128,256), 24:(4+128,256), 25:(4+128,256),
    26:(2+8,8), 27:(2+16,16), 28:(4+1,1), 29:(4+1,1), 30:(2,1),
}

def read_gguf(path):
    r = {'fused': 0, 'expert_bytes': 0, 'non_expert_bytes': 0}
    with open(path, 'rb') as f:
        magic = f.read(4)
        if magic != b'GGUF':
            print('0 0 0 0 0 0 0 0 0'); return
        ver = struct.unpack('<I', f.read(4))[0]
        tensor_count = struct.unpack('<Q', f.read(8))[0]
        n_kv = struct.unpack('<Q', f.read(8))[0]
        for _ in range(n_kv):
            kl = struct.unpack('<Q', f.read(8))[0]
            key = f.read(kl).decode('utf-8', errors='replace')
            vtype = struct.unpack('<I', f.read(4))[0]
            if vtype == 4:
                val = struct.unpack('<I', f.read(4))[0]
                if 'block_count' in key or '.layer_count' in key: r['layers'] = val
                elif 'expert_count' in key and 'used' not in key: r['experts'] = val
                elif 'head_count_kv' in key or '.attention.head_count_kv' in key: r['hkv'] = val
                elif '.attention.key_length' in key: r['kl'] = val
                elif '.attention.value_length' in key: r['vl'] = val
                elif 'ssm.conv_kernel' in key or 'ssm_conv_kernel' in key: r['ssm'] = 1
            elif vtype == 5:
                struct.unpack('<i', f.read(4))
            elif vtype == 6:
                struct.unpack('<f', f.read(4))
            elif vtype == 7:
                f.read(1)
            elif vtype == 8:
                sl = struct.unpack('<Q', f.read(8))[0]
                f.read(sl)
            elif vtype == 10:
                struct.unpack('<Q', f.read(8))
            elif vtype == 12:
                struct.unpack('<H', f.read(2))
            elif vtype == 9:
                atype = struct.unpack('<I', f.read(4))[0]
                alen = struct.unpack('<Q', f.read(8))[0]
                sizes = {4:4, 5:4, 6:4, 7:1, 10:8, 12:2}
                if atype == 8:
                    for _ in range(alen):
                        sl = struct.unpack('<Q', f.read(8))[0]
                        f.read(sl)
                elif atype in sizes:
                    f.read(alen * sizes[atype])
            else:
                break
        # Read tensor info
        for _ in range(tensor_count):
            tl = struct.unpack('<Q', f.read(8))[0]
            tname = f.read(tl).decode('utf-8', errors='replace')
            if 'ffn_up_gate' in tname or 'ffn_gate_up' in tname: r['fused'] = 1
            n_dims = struct.unpack('<I', f.read(4))[0]
            dims = [struct.unpack('<Q', f.read(8))[0] for _ in range(n_dims)]
            ttype = struct.unpack('<I', f.read(4))[0]
            f.read(8)  # offset
            n_elements = 1
            for d in dims: n_elements *= d
            bpb, epb = TYPE_SIZES.get(ttype, (4,1))
            tensor_bytes = (n_elements // epb) * bpb
            is_expert = '_exps.' in tname or '_shexp.' in tname or 'experts.' in tname
            if is_expert:
                r['expert_bytes'] += tensor_bytes
            else:
                r['non_expert_bytes'] += tensor_bytes
    # Handle split GGUFs: if tensor_count==0, scan sibling shards
    if tensor_count == 0:
        d = os.path.dirname(path)
        base = os.path.basename(path)
        m = re.match(r'(.+)-0+1-of-(\d+)\.gguf', base)
        if m:
            prefix, total = m.group(1), int(m.group(2))
            for i in range(2, total+1):
                shard = os.path.join(d, f'{prefix}-{i:05d}-of-{total:05d}.gguf')
                if not os.path.exists(shard): continue
                with open(shard, 'rb') as sf:
                    sf.read(4); sf.read(4)
                    st_count = struct.unpack('<Q', sf.read(8))[0]
                    sn_kv = struct.unpack('<Q', sf.read(8))[0]
                    for _ in range(sn_kv):
                        skl = struct.unpack('<Q', sf.read(8))[0]
                        sf.read(skl)
                        svt = struct.unpack('<I', sf.read(4))[0]
                        if svt == 4: sf.read(4)
                        elif svt == 5: sf.read(4)
                        elif svt == 6: sf.read(4)
                        elif svt == 7: sf.read(1)
                        elif svt == 8:
                            ssl = struct.unpack('<Q', sf.read(8))[0]
                            sf.read(ssl)
                        elif svt == 10: sf.read(8)
                        elif svt == 12: sf.read(2)
                        elif svt == 9:
                            sat = struct.unpack('<I', sf.read(4))[0]
                            sal = struct.unpack('<Q', sf.read(8))[0]
                            ssizes = {4:4,5:4,6:4,7:1,10:8,12:2}
                            if sat == 8:
                                for _ in range(sal):
                                    ssl2 = struct.unpack('<Q', sf.read(8))[0]
                                    sf.read(ssl2)
                            elif sat in ssizes:
                                sf.read(sal * ssizes[sat])
                    for _ in range(st_count):
                        stl = struct.unpack('<Q', sf.read(8))[0]
                        stname = sf.read(stl).decode('utf-8', errors='replace')
                        if 'ffn_up_gate' in stname or 'ffn_gate_up' in stname: r['fused'] = 1
                        sn_dims = struct.unpack('<I', sf.read(4))[0]
                        sdims = [struct.unpack('<Q', sf.read(8))[0] for _ in range(sn_dims)]
                        sttype = struct.unpack('<I', sf.read(4))[0]
                        sf.read(8)
                        sn_el = 1
                        for sd in sdims: sn_el *= sd
                        sbpb, sepb = TYPE_SIZES.get(sttype, (4,1))
                        st_bytes = (sn_el // sepb) * sbpb
                        is_exp = '_exps.' in stname or '_shexp.' in stname or 'experts.' in stname
                        if is_exp: r['expert_bytes'] += st_bytes
                        else: r['non_expert_bytes'] += st_bytes
    print(r.get('layers',0), r.get('experts',0), r.get('hkv',0), r.get('kl',0), r.get('vl',0), r.get('ssm',0), r.get('fused',0), r.get('expert_bytes',0), r.get('non_expert_bytes',0))

read_gguf('$MODEL_PATH')
" 2>/dev/null || echo "0 0 0 0 0 0 0 0 0")

# Defaults if parsing failed
LAYER_COUNT=${LAYER_COUNT:-0}
EXPERT_COUNT=${EXPERT_COUNT:-0}
HEAD_COUNT_KV=${HEAD_COUNT_KV:-0}
KEY_LENGTH=${KEY_LENGTH:-128}
VALUE_LENGTH=${VALUE_LENGTH:-128}
HAS_SSM=${HAS_SSM:-0}
HAS_FUSED=${HAS_FUSED:-0}
EXPERT_BYTES=${EXPERT_BYTES:-0}
NON_EXPERT_BYTES=${NON_EXPERT_BYTES:-0}

IS_MOE=0
(( EXPERT_COUNT > 1 )) && IS_MOE=1

echo ""
echo "Model: $MODEL_NAME"
echo "Size: ${TOTAL_SIZE_GB}GB (${TOTAL_SIZE_MB}MB)"
echo "Architecture: ${LAYER_COUNT} layers, $([ "$IS_MOE" = "1" ] && echo "${EXPERT_COUNT} experts (MoE)" || echo "dense")$([ "$HAS_FUSED" = "1" ] && echo ", fused up|gate")"

# ── Fused model info ──
if (( HAS_FUSED == 1 )); then
    echo "Info: Fused up|gate model detected."
fi

# ═══════════════════════════════════════════════════════════════
# Section 4: Smart Flag Builder (macOS)
# ═══════════════════════════════════════════════════════════════

echo ""
echo "─── Configuring flags ───"

# ── Context shift: crashes on hybrid SSM/Mamba models ──
CONTEXT_SHIFT_FLAG=""
if (( HAS_SSM == 1 )); then
    CONTEXT_SHIFT_FLAG="--no-context-shift"
    echo "  SSM/Mamba hybrid → context-shift disabled"
fi

# ── Memory Analysis ──
# On Apple Silicon, GPU and CPU share unified memory
# Model + KV cache + overhead must all fit in system RAM
TOTAL_MEM_MB=$RAM_TOTAL_MB
KV_EST_MB=$(( CTX_SIZE * LAYER_COUNT / 1024 ))
ESTIMATED_TOTAL_NEEDED=$(( TOTAL_SIZE_MB + KV_EST_MB + 2048 ))

if (( TOTAL_SIZE_MB > TOTAL_MEM_MB )); then
    echo "⚠️  WARNING: Model (${TOTAL_SIZE_GB}GB) is larger than system memory (${RAM_TOTAL_MB}MB)."
    echo "   Heavy swapping likely. Consider a smaller quantization."
elif (( ESTIMATED_TOTAL_NEEDED > TOTAL_MEM_MB )); then
    echo "ℹ️  INFO: Model + context (~${ESTIMATED_TOTAL_NEEDED}MB) is close to system memory (${RAM_TOTAL_MB}MB)."
    echo "   Consider reducing context: --ctx-size 16384"
else
    echo "✓ Memory: Model and context fit within unified memory."
fi

# ── Batch sizes: scale with available memory ──
MEM_HEADROOM=$(( RAM_AVAIL_MB - TOTAL_SIZE_MB ))
if (( MEM_HEADROOM > 8192 )); then
    BATCH=4096; UBATCH=512
    echo "  Memory headroom large → batch=$BATCH ubatch=$UBATCH"
elif (( MEM_HEADROOM > 2048 )); then
    BATCH=2048; UBATCH=512
    echo "  Moderate headroom → batch=$BATCH ubatch=$UBATCH"
else
    BATCH=512; UBATCH=256
    echo "  Tight memory → batch=$BATCH ubatch=$UBATCH"
fi

# ── KV cache type ──
if (( HEAD_COUNT_KV > 0 && KEY_LENGTH > 0 && VALUE_LENGTH > 0 )); then
    KV_Q4_MB=$(awk "BEGIN {printf \"%.0f\", $CTX_SIZE * $LAYER_COUNT * $HEAD_COUNT_KV * ($KEY_LENGTH + $VALUE_LENGTH) * 0.5 / 1048576}")
    KV_Q8_MB=$(awk "BEGIN {printf \"%.0f\", $CTX_SIZE * $LAYER_COUNT * $HEAD_COUNT_KV * ($KEY_LENGTH + $VALUE_LENGTH) * 1.0 / 1048576}")
    KV_F16_MB=$(awk "BEGIN {printf \"%.0f\", $CTX_SIZE * $LAYER_COUNT * $HEAD_COUNT_KV * ($KEY_LENGTH + $VALUE_LENGTH) * 2.0 / 1048576}")
else
    KV_Q4_MB=$(( LAYER_COUNT * 70 ))
    KV_Q8_MB=$(( LAYER_COUNT * 140 ))
    KV_F16_MB=$(( LAYER_COUNT * 280 ))
fi

case "$KV_QUALITY" in
    high) KV_TYPE="f16";  echo "  KV cache: f16 (${KV_F16_MB}MB) — highest quality" ;;
    mid)  KV_TYPE="q8_0"; echo "  KV cache: q8_0 (${KV_Q8_MB}MB) — balanced (default)" ;;
    low)  KV_TYPE="q4_0"; echo "  KV cache: q4_0 (${KV_Q4_MB}MB) — minimum memory" ;;
    *)    echo "Error: --kv-quality must be one of: high, mid, low"; exit 1 ;;
esac

# ── Auto-reduce context if model barely fits ──
ACTUAL_CTX=$CTX_SIZE
MODEL_PLUS_KV=$ESTIMATED_TOTAL_NEEDED
if (( MODEL_PLUS_KV > RAM_TOTAL_MB * 90 / 100 )); then
    # Over 90% of RAM — reduce context to fit
    for try_ctx in 16384 8192 4096; do
        kv_at_ctx=$(( KV_EST_MB * try_ctx / CTX_SIZE ))
        needed=$(( TOTAL_SIZE_MB + kv_at_ctx + 2048 ))
        if (( needed < RAM_TOTAL_MB * 85 / 100 )); then
            echo "  Auto-reducing context from $CTX_SIZE to $try_ctx to fit in memory"
            ACTUAL_CTX=$try_ctx
            CTX_SIZE=$try_ctx
            break
        fi
    done
fi

# ── Thread count ──
# On Apple Silicon: use performance cores for generation
THREADS_GEN=$PERF_CORES
THREADS_BATCH=$PHYSICAL_CORES
echo "  Threads: gen=$THREADS_GEN (P-cores) batch=$THREADS_BATCH (all cores)"

# ── Build the flags array ──
COMMON_FLAGS=(
    -m "$MODEL_PATH"
    --host "$HOST"
    --port "$PORT"
    --ctx-size "$CTX_SIZE"
    --flash-attn on
    -b "$BATCH"
    -ub "$UBATCH"
    --cache-type-k "$KV_TYPE"
    --cache-type-v "$KV_TYPE"
    --jinja
    --threads "$THREADS_GEN"
    --threads-batch "$THREADS_BATCH"
)

# Conditional flags
if [[ "$IS_IK_LLAMA" == "1" ]]; then
    COMMON_FLAGS+=(--run-time-repack)
    if [[ "$KV_TYPE" == "q4_0" || "$KV_TYPE" == "q8_0" ]]; then
        COMMON_FLAGS+=(-khad)
        echo "  Quantized KV cache → Hadamard K-transform enabled"
    fi
    [[ -n "$CONTEXT_SHIFT_FLAG" ]] && COMMON_FLAGS+=("$CONTEXT_SHIFT_FLAG")
    COMMON_FLAGS+=(--defrag-thold 0.1)

    if (( IS_MOE == 1 )); then
        COMMON_FLAGS+=(-muge -ger)
        echo "  MoE model → -muge -ger enabled"
    fi
else
    [[ -n "$CONTEXT_SHIFT_FLAG" ]] && COMMON_FLAGS+=("$CONTEXT_SHIFT_FLAG")
fi

# Metal GPU offloading — offload all layers to Metal
if (( HAS_METAL )); then
    COMMON_FLAGS+=(-ngl 999)
    echo "  Metal GPU → all layers offloaded (-ngl 999)"
fi

echo ""

# ═══════════════════════════════════════════════════════════════
# Section 5: Strategy Selection (macOS — simplified)
# ═══════════════════════════════════════════════════════════════

# On Apple Silicon, there's only one strategy decision:
# - Model fits in unified memory → GPU offload all layers
# - Model doesn't fit → warn and try anyway (macOS handles swap via compressed memory)
# - CPU-only mode → no GPU offload

if (( CPU_ONLY )); then
    STRATEGY="cpu_only"
elif (( HAS_METAL )); then
    STRATEGY="metal_gpu"
else
    STRATEGY="cpu_only"
fi

echo "Strategy: $STRATEGY"

# ═══════════════════════════════════════════════════════════════
# Section 6: Helper Functions
# ═══════════════════════════════════════════════════════════════

kill_server() {
    local pid=${1:-}
    if [[ -n "$pid" ]] && kill -0 "$pid" 2>/dev/null; then
        kill -9 "$pid" 2>/dev/null
        wait "$pid" 2>/dev/null || true
    fi
    # Port cleanup
    local port_pids
    port_pids=$(lsof -t -i:"${PORT}" 2>/dev/null) || true
    if [[ -n "$port_pids" ]]; then
        kill -9 $port_pids 2>/dev/null || true
    fi
    local w=0
    while (( w < 10 )); do
        lsof -i:"${PORT}" >/dev/null 2>&1 || break
        sleep 1; (( w++ ))
    done
    sleep 1
}

RUNNING_PID=""
try_start() {
    local extra_flags=("$@")
    RUNNING_PID=""
    "$LLAMA_SERVER" "${COMMON_FLAGS[@]}" "${extra_flags[@]}" > >(tee -a "$SERVER_LOG") 2>&1 &
    local pid=$!
    local i=0
    while (( i < 240 )); do
        if curl -sf "http://127.0.0.1:${PORT}/health" >/dev/null 2>&1; then
            RUNNING_PID=$pid
            return 0
        fi
        if ! kill -0 "$pid" 2>/dev/null; then
            return 1
        fi
        sleep 1; (( i++ ))
    done
    kill -9 "$pid" 2>/dev/null; wait "$pid" 2>/dev/null || true
    return 1
}

print_cmd() {
    echo ""
    echo "Command:"
    local line="  $LLAMA_SERVER"
    local args=("$@")
    for (( i=0; i<${#args[@]}; i++ )); do
        local arg="${args[$i]}"
        if [[ "$arg" == -* ]] && (( i + 1 < ${#args[@]} )) && [[ "${args[$((i+1))]}" != -* ]]; then
            local pair="$arg ${args[$((i+1))]}"
            i=$(( i + 1 ))
        else
            local pair="$arg"
        fi
        if (( ${#line} + ${#pair} + 1 > 80 )); then
            echo "$line \\"
            line="    $pair"
        else
            line="$line $pair"
        fi
    done
    echo "$line"
}

run_benchmark() {
    local url="http://127.0.0.1:${PORT}"
    echo ""
    echo "═══ Benchmark ═══"

    local result
    result=$(curl -sf "$url/v1/chat/completions" \
        -H "Content-Type: application/json" \
        -d '{"model":"test","messages":[{"role":"user","content":"Explain the theory of relativity in simple terms. Cover special and general relativity, time dilation, and gravitational effects."}],"max_tokens":200,"temperature":0.1}' 2>/dev/null)

    if [[ -z "$result" ]]; then
        echo "Benchmark failed — server not responding"
        return 1
    fi

    local pp_tokens tg_tokens
    pp_tokens=$(echo "$result" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('usage',{}).get('prompt_tokens',0))" 2>/dev/null)
    tg_tokens=$(echo "$result" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('usage',{}).get('completion_tokens',0))" 2>/dev/null)

    local timings
    timings=$(curl -sf "$url/slots" 2>/dev/null)
    if [[ -n "$timings" ]]; then
        local pp_rate tg_rate
        pp_rate=$(echo "$timings" | python3 -c "
import sys,json
d=json.load(sys.stdin)
s=d[0] if isinstance(d,list) else d
t_pp = s.get('t_prompt_processing', 0)
n_pp = s.get('n_prompt_tokens_processed', 1)
if t_pp > 0: print(f\"{n_pp / (t_pp / 1000):.1f}\")
else: print('?')
" 2>/dev/null)
        tg_rate=$(echo "$timings" | python3 -c "
import sys,json
d=json.load(sys.stdin)
s=d[0] if isinstance(d,list) else d
t_gen = s.get('t_token_generation', 0)
n_gen = s.get('n_decoded', 1)
if t_gen > 0: print(f\"{n_gen / (t_gen / 1000):.1f}\")
else: print('?')
" 2>/dev/null)
        echo "  Prompt processing: ${pp_tokens:-?} tokens @ ${pp_rate:-?} tok/s"
        echo "  Generation:        ${tg_tokens:-?} tokens @ ${tg_rate:-?} tok/s"
    else
        echo "  Prompt tokens: ${pp_tokens:-?}"
        echo "  Generated tokens: ${tg_tokens:-?}"
    fi
    echo ""
    echo "Benchmark complete. Shutting down server..."
    kill_server ""
    exit 0
}

suggest_alternatives() {
    local model_mb=$1
    local available_mb=$2
    echo ""
    echo "═══ Model doesn't fit ═══"
    echo "  Model needs: ~${model_mb}MB (with overhead)"
    echo "  Available:   ~${available_mb}MB (unified memory)"
    echo ""
    echo "  Options:"
    echo "  1. Use a smaller quantization (e.g., Q4 → IQ3, IQ3 → IQ2)"
    echo "  2. Reduce context size: --ctx-size 16384 or --ctx-size 8192"
    echo "  3. Use a smaller model"
    echo "  4. Upgrade to a Mac with more unified memory (M4 Pro 48GB, M4 Max 128GB)"
    echo ""
}

MAX_RESTARTS=5
run_with_restart() {
    local flags=("$@")
    local restarts=0
    local last_start=0

    while true; do
        last_start=$(date +%s)
        echo "Starting server (PID will follow)..."
        "$LLAMA_SERVER" "${flags[@]}" 2>&1 | tee "$SERVER_LOG" &
        local pid=$!
        echo "Server PID: $pid"

        local healthy=0
        for i in $(seq 1 240); do
            if curl -sf "http://127.0.0.1:${PORT}/health" >/dev/null 2>&1; then
                healthy=1
                break
            fi
            if ! kill -0 "$pid" 2>/dev/null; then
                break
            fi
            sleep 1
        done

        if (( healthy )); then
            echo "Server healthy."
            if (( BENCHMARK )); then
                run_benchmark
                BENCHMARK=0
            fi
            restarts=0
            wait "$pid" 2>/dev/null || true
            local exit_code=$?
        else
            echo "Server failed to become healthy."
            kill -9 "$pid" 2>/dev/null; wait "$pid" 2>/dev/null || true
            local exit_code=1
        fi

        local uptime=$(( $(date +%s) - last_start ))

        if (( exit_code == 0 || exit_code == 143 )); then
            echo "Server exited cleanly (code $exit_code). Not restarting."
            exit 0
        fi

        restarts=$(( restarts + 1 ))
        echo ""
        echo "═══ Server crashed (exit code $exit_code, uptime ${uptime}s) ═══"
        tail -5 "$SERVER_LOG" 2>/dev/null
        echo ""

        if (( restarts > MAX_RESTARTS )); then
            echo "Too many restarts ($MAX_RESTARTS). Giving up."
            echo "Check $SERVER_LOG for details."
            exit 1
        fi

        local wait_sec=$(( 2 ** restarts ))
        (( wait_sec > 30 )) && wait_sec=30

        if (( uptime > 60 )); then
            wait_sec=2
            echo "Runtime crash (ran ${uptime}s). Restart $restarts/$MAX_RESTARTS in ${wait_sec}s..."
        else
            echo "Startup crash. Restart $restarts/$MAX_RESTARTS in ${wait_sec}s..."
        fi

        sleep "$wait_sec"
        kill_server ""
    done
}

# ═══════════════════════════════════════════════════════════════
# Section 7: Execute Strategy
# ═══════════════════════════════════════════════════════════════

# Kill any existing server on this port
kill_server ""

case "$STRATEGY" in
    cpu_only)
        if (( CPU_ONLY )); then
            echo "Running on CPU only (forced via --cpu)"
        else
            echo "Running on CPU only (no Metal GPU detected)"
        fi
        LAUNCH_FLAGS=("${COMMON_FLAGS[@]}" -ngl 0 --mlock -b 512 -ub 256)
        # Strip GPU flags
        clean_flags=()
        skip_next=0
        for f in "${LAUNCH_FLAGS[@]}"; do
            if (( skip_next )); then skip_next=0; continue; fi
            case "$f" in
                -ngl|-mg|-b|-ub) skip_next=1; continue ;;
                --run-time-repack|-muge|-ger) continue ;;
                *) clean_flags+=("$f") ;;
            esac
        done
        LAUNCH_FLAGS=("${clean_flags[@]}" -ngl 0 -b 512 -ub 256)
        if (( DRY_RUN )); then
            print_cmd "${LAUNCH_FLAGS[@]}"
            exit 0
        fi
        run_with_restart "${LAUNCH_FLAGS[@]}"
        ;;

    metal_gpu)
        echo "Metal GPU: all layers offloaded to $METAL_GPU_NAME"
        LAUNCH_FLAGS=("${COMMON_FLAGS[@]}")
        if (( DRY_RUN )); then
            print_cmd "${LAUNCH_FLAGS[@]}"
            exit 0
        fi
        run_with_restart "${LAUNCH_FLAGS[@]}"
        ;;
esac