llm-server

#!/bin/bash
#
# llm-server — Universal smart launcher for ik_llama.cpp / llama.cpp
# Auto-detects GPUs, RAM, model architecture (dense/MoE), and optimizes placement.
# Works on any system: 0-8+ GPUs, any VRAM/RAM configuration.
#
# Usage:
#   llm-server <repo/name> --download    # Download from HuggingFace
#   llm-server -h                        # Show this help message
#   llm-server <model.gguf>              # Launch with auto-tuning
#   llm-server <model.gguf> --retune     # Force re-optimization
#   llm-server <model.gguf> --dry-run    # Print command without executing
#   llm-server --port 8082 <model.gguf>  # Custom port
#   llm-server <model.gguf> --benchmark  # Start, run quick benchmark, print tok/s
#   llm-server <model.gguf> --cpu        # Force CPU-only (ignore GPUs)
#   llm-server <model.gguf> --server-bin /path/to/llama-server
#   llm-server <model.gguf> --lib-path /path/to/libs
#   llm-server <model.gguf> --kv-quality high  # KV cache: high(f16), mid(q8_0), low(q4_0)
#   llm-server <model.gguf> --gpus 0,1         # Use only GPU 0 and 1
#   llm-server <model.gguf> --ram-budget 60G   # Cap RAM usage at 60GB
#
# 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: 65536)

set -euo pipefail

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

VERSION="2.0.0"
PORT="${LLM_PORT:-8081}"
HOST="0.0.0.0"
CTX_SIZE="${LLM_CTX_SIZE:-65536}"

# Tuning constants
MAX_RESTARTS=5
SYSTEM_HEADROOM_MB=5120          # RAM reserved for OS/system
COMPUTE_PER_GPU_MB=512           # CUDA context + compute buffers per GPU
MIN_CRAM_MB=512                  # minimum useful CPU RAM for expert offload
VRAM_OVERHEAD_PERCENT=130        # model size * this / 100 = estimated VRAM needed
SINGLE_GPU_HEADROOM_MB=4096      # extra VRAM headroom to fit single-GPU mode

# 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}"
# Ensure model directory exists
mkdir -p "$MODEL_DIR"
CACHE_DIR="$HOME/.cache/llm-server"
SERVER_LOG="/tmp/llm-server.log"
LIB_HUB_DIR=""

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

# Cleanup on exit
cleanup() {
    if [[ -n "${LIB_HUB_DIR:-}" && -d "$LIB_HUB_DIR" ]]; then
        rm -rf "$LIB_HUB_DIR"
    fi
}
trap cleanup EXIT

# Create a "Lib Hub" (symlink all found .so files into one temp dir)
# This solves the issue where libraries are scattered in build subfolders.
setup_lib_hub() {
    local bin_path="$1"
    local build_root
    build_root="$(dirname "$(dirname "$bin_path")")"
    
    if [[ ! -d "$build_root" ]]; then return; fi

    # Create fresh hub dir
    [[ -n "$LIB_HUB_DIR" && -d "$LIB_HUB_DIR" ]] && rm -rf "$LIB_HUB_DIR"
    LIB_HUB_DIR=$(mktemp -d /tmp/llm-lib-hub.XXXXXX)
    
    # Find and symlink all .so files
    find "$build_root" -name "*.so" -type f -exec ln -sf {} "$LIB_HUB_DIR/" \; 2>/dev/null
    
    # Export the hub to the library path
    export LD_LIBRARY_PATH="$LIB_HUB_DIR:${LD_LIBRARY_PATH:-}"
    
    # Also add the binary's directory for good measure
    export LD_LIBRARY_PATH="$(dirname "$bin_path"):$LD_LIBRARY_PATH"
}

# Auto-detect llama-server binary
find_server_binary() {
    if [[ -n "${LLAMA_SERVER:-}" && -x "$LLAMA_SERVER" ]]; then
        echo "$LLAMA_SERVER"
        return
    fi
    # Prefer ik_llama.cpp (faster multi-GPU)
    local candidates=(
        "$HOME/ik_llama.cpp/build/bin/llama-server"
        "$HOME/llama.cpp/build/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)"

# Initial lib hub setup
if [[ -n "$LLAMA_SERVER" ]]; then
    setup_lib_hub "$LLAMA_SERVER"
fi

# Note: --server-bin and --lib-path are applied after arg parsing (see below)

# 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 ;;
        --lib-path)    USER_LIB_PATH="$2"; shift 2 ;;
        --kv-quality)
            case "$2" in high|mid|low) ;; *) echo "Error: --kv-quality must be one of: high, mid, low"; exit 1 ;; esac
            KV_QUALITY="$2"; shift 2 ;;
        --gpus)        GPUS_FILTER="$2"; shift 2 ;;
        --ram-budget)  RAM_BUDGET_MB="$2"; shift 2 ;;
        --download)  DOWNLOAD=1; shift ;;
        --version)     echo "llm-server v$VERSION"; exit 0 ;;
        --help|-h)
            sed -n '2,/^$/s/^# \?//p' "$0"
            exit 0
            ;;
        *)           MODEL_ARG="$1"; shift ;;
    esac
done

# Parse --ram-budget human units (e.g. "60G" → 61440 MB)
if [[ -n "$RAM_BUDGET_MB" && "$RAM_BUDGET_MB" != "0" ]]; then
    case "$RAM_BUDGET_MB" in
        *[Gg]) RAM_BUDGET_MB=$(( ${RAM_BUDGET_MB%[Gg]} * 1024 )) ;;
        *[Mm]) RAM_BUDGET_MB=${RAM_BUDGET_MB%[Mm]} ;;
        *[0-9]) ;;  # already numeric (MB)
        *) echo "Error: --ram-budget must be a number with optional G/M suffix (e.g. 60G, 4096M, 4096)"; exit 1 ;;
    esac
    if ! [[ "$RAM_BUDGET_MB" =~ ^[0-9]+$ ]]; then
        echo "Error: --ram-budget must be a number with optional G/M suffix (e.g. 60G, 4096M, 4096)"; exit 1
    fi
fi

# Validate --port
if ! [[ "$PORT" =~ ^[0-9]+$ ]] || (( PORT < 1 || PORT > 65535 )); then
    echo "Error: --port must be a number between 1 and 65535 (got: $PORT)"
    exit 1
fi

# Validate --ctx-size
if ! [[ "$CTX_SIZE" =~ ^[0-9]+$ ]] || (( CTX_SIZE < 1 )); then
    echo "Error: --ctx-size must be a positive number (got: $CTX_SIZE)"
    exit 1
fi

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

# ── Download Logic ──
if (( DOWNLOAD )); then
    TOTAL_VRAM_MB=0
    if command -v nvidia-smi >/dev/null 2>&1; then
        TOTAL_VRAM_MB=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits | awk '{s+=$1} END {print s}')
    fi
    RAM_AVAIL_MB=$(grep MemAvailable /proc/meminfo | awk '{print int($2/1024)}')
    # Try finding the downloader in several places
    DOWNLOADER="$(dirname "$(readlink -f "$0")")/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."
        echo "Please ensure it exists in /home/mik/ai_models/ or the script directory."
        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 <path_to_downloaded_file>"
    exit 0
fi

if [[ ! -f "$MODEL_ARG" ]]; then
    # Try as relative to model dir
    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 (before binary check so it can provide the binary)
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"
    # Setup lib hub for the new binary
    setup_lib_hub "$LLAMA_SERVER"
fi

if [[ -z "$LLAMA_SERVER" ]]; then
    echo "Error: llama-server binary not found."
    echo "Set LLAMA_SERVER env var, use --server-bin, or install ik_llama.cpp / llama.cpp."
    exit 1
fi

# Apply --lib-path override (prepend to LD_LIBRARY_PATH)
if [[ -n "$USER_LIB_PATH" ]]; then
    export LD_LIBRARY_PATH="${USER_LIB_PATH}:${LD_LIBRARY_PATH:-}"
fi

# Detect if we're running ik_llama.cpp (supports --split-mode graph)
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

export CUDA_DEVICE_ORDER=PCI_BUS_ID

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

# ═══════════════════════════════════════════════════════════════
# Section 2: Hardware Detection
# ═══════════════════════════════════════════════════════════════

echo "═══ llm-server v$VERSION ═══"
echo "Binary: $LLAMA_SERVER"
(( IS_IK_LLAMA )) && echo "Backend: ik_llama.cpp (graph split enabled)"

# Detect physical CPU cores
PHYSICAL_CORES=$(lscpu 2>/dev/null | awk '/^Core\(s\) per socket:/ {cores=$NF} /^Socket\(s\):/ {socks=$NF} END {print cores * socks}')
PHYSICAL_CORES=${PHYSICAL_CORES:-4}

# Detect RAM
RAM_AVAIL_MB=$(awk '/MemAvailable/ {printf "%.0f", $2/1024}' /proc/meminfo)
RAM_TOTAL_MB=$(awk '/MemTotal/ {printf "%.0f", $2/1024}' /proc/meminfo)

echo "CPU: ${PHYSICAL_CORES} physical cores"

# Apply --ram-budget cap
if (( RAM_BUDGET_MB > 0 && RAM_BUDGET_MB < RAM_AVAIL_MB )); then
    RAM_AVAIL_MB=$RAM_BUDGET_MB
    echo "RAM: ${RAM_AVAIL_MB}MB available (capped by --ram-budget) / ${RAM_TOTAL_MB}MB total"
else
    echo "RAM: ${RAM_AVAIL_MB}MB available / ${RAM_TOTAL_MB}MB total"
fi

# Detect GPUs
GPU_COUNT=0
declare -a GPU_INDEX GPU_NAME GPU_VRAM_TOTAL GPU_VRAM_FREE GPU_PCIE_WIDTH GPU_PCIE_GEN GPU_BANDWIDTH

if (( CPU_ONLY )); then
    echo "GPUs: skipped (--cpu flag)"
elif command -v nvidia-smi &>/dev/null; then
    while IFS= read -r line; do
        # Parse comma-separated fields from nvidia-smi
        IFS=',' read -ra fields <<< "$line"
        [[ ${#fields[@]} -lt 6 ]] && continue
        idx=$(echo "${fields[0]}" | tr -d ' ')
        name=$(echo "${fields[1]}" | sed 's/^ //')
        vram_total=$(echo "${fields[2]}" | tr -d ' ')
        vram_free=$(echo "${fields[3]}" | tr -d ' ')
        pcie_width=$(echo "${fields[4]}" | tr -d ' ')
        pcie_gen=$(echo "${fields[5]}" | tr -d ' ')

        # Skip GPUs with less than 500MB free
        if (( vram_free < 500 )); then
            log "Skipping GPU $idx ($name): only ${vram_free}MB free"
            continue
        fi
        GPU_INDEX+=("$idx")
        GPU_NAME+=("$name")
        GPU_VRAM_TOTAL+=("$vram_total")
        GPU_VRAM_FREE+=("$vram_free")
        GPU_PCIE_WIDTH+=("$pcie_width")
        GPU_PCIE_GEN+=("$pcie_gen")
        # Effective bandwidth score: width * gen (higher = faster)
        GPU_BANDWIDTH+=( $(( pcie_width * pcie_gen )) )
        GPU_COUNT=$(( GPU_COUNT + 1 ))
    done < <(nvidia-smi --query-gpu=index,name,memory.total,memory.free,pcie.link.width.current,pcie.link.gen.current --format=csv,noheader,nounits 2>/dev/null)
fi

# Apply --gpus filter (restrict to specified CUDA indices)
if [[ -n "$GPUS_FILTER" && $GPU_COUNT -gt 0 ]]; then
    IFS=',' read -ra ALLOWED_GPUS <<< "$GPUS_FILTER"
    declare -a NEW_GPU_INDEX NEW_GPU_NAME NEW_GPU_VRAM_TOTAL NEW_GPU_VRAM_FREE NEW_GPU_PCIE_WIDTH NEW_GPU_PCIE_GEN NEW_GPU_BANDWIDTH
    NEW_GPU_COUNT=0
    for i in $(seq 0 $(( GPU_COUNT - 1 ))); do
        for allowed in "${ALLOWED_GPUS[@]}"; do
            if [[ "${GPU_INDEX[$i]}" == "$allowed" ]]; then
                NEW_GPU_INDEX+=("${GPU_INDEX[$i]}")
                NEW_GPU_NAME+=("${GPU_NAME[$i]}")
                NEW_GPU_VRAM_TOTAL+=("${GPU_VRAM_TOTAL[$i]}")
                NEW_GPU_VRAM_FREE+=("${GPU_VRAM_FREE[$i]}")
                NEW_GPU_PCIE_WIDTH+=("${GPU_PCIE_WIDTH[$i]}")
                NEW_GPU_PCIE_GEN+=("${GPU_PCIE_GEN[$i]}")
                NEW_GPU_BANDWIDTH+=("${GPU_BANDWIDTH[$i]}")
                (( NEW_GPU_COUNT++ )) || true
                break
            fi
        done
    done
    if (( NEW_GPU_COUNT == 0 )); then
        echo "Error: --gpus filter '$GPUS_FILTER' matched no detected GPUs"
        exit 1
    fi
    GPU_INDEX=("${NEW_GPU_INDEX[@]}")
    GPU_NAME=("${NEW_GPU_NAME[@]}")
    GPU_VRAM_TOTAL=("${NEW_GPU_VRAM_TOTAL[@]}")
    GPU_VRAM_FREE=("${NEW_GPU_VRAM_FREE[@]}")
    GPU_PCIE_WIDTH=("${NEW_GPU_PCIE_WIDTH[@]}")
    GPU_PCIE_GEN=("${NEW_GPU_PCIE_GEN[@]}")
    GPU_BANDWIDTH=("${NEW_GPU_BANDWIDTH[@]}")
    GPU_COUNT=$NEW_GPU_COUNT
    echo "GPU filter: using only GPU(s) $GPUS_FILTER ($GPU_COUNT matched)"
fi

if (( GPU_COUNT == 0 )); then
    echo "GPUs: none detected (CPU-only mode)"
else
    echo "GPUs: $GPU_COUNT detected"
    # Sort GPUs by bandwidth (descending) — build priority order
    declare -a GPU_ORDER
    GPU_ORDER=($(
        for i in $(seq 0 $(( GPU_COUNT - 1 ))); do
            echo "${GPU_BANDWIDTH[$i]} $i"
        done | sort -rn | awk '{print $2}'
    ))

    for i in $(seq 0 $(( GPU_COUNT - 1 ))); do
        gi=${GPU_ORDER[$i]}
        echo "  GPU${GPU_INDEX[$gi]}: ${GPU_NAME[$gi]} ${GPU_VRAM_FREE[$gi]}MB free / ${GPU_VRAM_TOTAL[$gi]}MB total (PCIe x${GPU_PCIE_WIDTH[$gi]} gen${GPU_PCIE_GEN[$gi]})"
    done
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/')
TOTAL_SIZE_BYTES=$(find "$MODEL_DIR_PATH" -name "${SPLIT_PATTERN%.gguf}*" -type f 2>/dev/null | xargs du -scb 2>/dev/null | tail -1 | awk '{print $1}')
if [[ -z "$TOTAL_SIZE_BYTES" || "$TOTAL_SIZE_BYTES" == "0" ]]; then
    TOTAL_SIZE_BYTES=$(du -sb "$MODEL_PATH" | awk '{print $1}')
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
read LAYER_COUNT EXPERT_COUNT HEAD_COUNT_KV KEY_LENGTH VALUE_LENGTH HAS_SSM HAS_FUSED EXPERT_BYTES NON_EXPERT_BYTES < <(MODEL_PATH="$MODEL_PATH" python3 -c "
import struct, sys, glob, os, re

# GGUF type -> (bytes_per_block, elements_per_block) from ggml.h struct sizes
GGUF_TYPE_SIZE = {
    0:  (4, 1),       # F32
    1:  (2, 1),       # F16
    2:  (18, 32),     # Q4_0
    3:  (20, 32),     # Q4_1
    6:  (22, 32),     # Q5_0
    7:  (24, 32),     # Q5_1
    8:  (34, 32),     # Q8_0
    9:  (36, 32),     # Q8_1
    10: (84, 256),    # Q2_K
    11: (110, 256),   # Q3_K
    12: (144, 256),   # Q4_K
    13: (176, 256),   # Q5_K
    14: (210, 256),   # Q6_K
    15: (292, 256),   # Q8_K
    16: (66, 256),    # IQ2_XXS
    17: (74, 256),    # IQ2_XS
    18: (98, 256),    # IQ3_XXS
    19: (50, 256),    # IQ1_S
    20: (18, 32),     # IQ4_NL
    21: (110, 256),   # IQ3_S
    22: (82, 256),    # IQ2_S
    23: (136, 256),   # IQ4_XS
    24: (56, 256),    # IQ1_M
    25: (2, 1),       # BF16
    26: (18, 32),     # Q4_0_4_4
    27: (18, 32),     # Q4_0_4_8
    28: (18, 32),     # Q4_0_8_8
    29: (40, 256),    # TQ1_0
    30: (54, 256),    # TQ2_0
    31: (1, 1),       # I8
}

r = {'fused': 0, 'expert_bytes': 0, 'non_expert_bytes': 0}

def skip_kv(f, kv_count):
    # KV type sizes: 0=u8(1) 1=i8(1) 2=u16(2) 3=i16(2) 4=u32(4) 5=i32(4)
    # 6=f32(4) 7=bool(1) 8=string 9=array 10=u64(8) 11=i64(8) 12=f64(8)
    KV_FIXED = {0:1,1:1,2:2,3:2,4:4,5:4,6:4,7:1,10:8,11:8,12:8}
    for _ in range(kv_count):
        kl = struct.unpack('<Q', f.read(8))[0]
        key = f.read(kl).decode('utf-8', errors='replace')
        vt = struct.unpack('<I', f.read(4))[0]
        if vt == 4:
            val = struct.unpack('<I', f.read(4))[0]
            if key.endswith('.block_count'): r['layers'] = val
            if 'expert_count' in key and 'used' not in key: r['experts'] = val
            if 'head_count_kv' in key: r['hkv'] = val
            if 'key_length' in key: r['kl'] = val
            if 'value_length' in key: r['vl'] = val
            if 'ssm.state_size' in key: r['ssm'] = 1
        elif vt == 8: f.read(struct.unpack('<Q', f.read(8))[0])
        elif vt == 9:
            at = struct.unpack('<I', f.read(4))[0]; al = struct.unpack('<Q', f.read(8))[0]
            if at in KV_FIXED: f.read(al * KV_FIXED[at])
            elif at == 8:
                for _ in range(al): f.read(struct.unpack('<Q', f.read(8))[0])
            elif at == 9: break  # nested arrays: bail
            else: break
        elif vt in KV_FIXED: f.read(KV_FIXED[vt])
        else: break

def read_tensors(f, tensor_count):
    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
        if ttype in GGUF_TYPE_SIZE:
            bpb, epb = GGUF_TYPE_SIZE[ttype]
            n_blocks = (n_elements + epb - 1) // epb
            tbytes = n_blocks * bpb
        else:
            tbytes = n_elements * 2
        is_expert = '_exps.' in tname or '_shexp.' in tname or 'experts.' in tname
        if is_expert:
            r['expert_bytes'] += tbytes
        else:
            r['non_expert_bytes'] += tbytes

try:
    model_path = os.environ['MODEL_PATH']
    with open(model_path, 'rb') as f:
        magic = f.read(4)
        if magic != b'GGUF': sys.exit(1)
        f.read(4)  # version
        tensor_count = struct.unpack('<Q', f.read(8))[0]
        kv_count = struct.unpack('<Q', f.read(8))[0]
        skip_kv(f, kv_count)
        read_tensors(f, tensor_count)

    # For split GGUFs: always scan all sibling shards (tensors may be spread across shards)
    m = re.search(r'-(\d+)-of-(\d+)\.gguf$', model_path)
    if m:
        total_shards = int(m.group(2))
        base = model_path[:m.start()]
        for shard_num in range(2, total_shards + 1):
            shard_path = f'{base}-{shard_num:05d}-of-{total_shards:05d}.gguf'
            if not os.path.exists(shard_path): continue
            with open(shard_path, 'rb') as f:
                magic = f.read(4)
                if magic != b'GGUF': continue
                f.read(4)
                tc = struct.unpack('<Q', f.read(8))[0]
                kvc = struct.unpack('<Q', f.read(8))[0]
                skip_kv(f, kvc)
                read_tensors(f, tc)
except: pass
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))
" 2>/dev/null || echo "0 0 0 0 0 0 0 0 0")

LAYER_COUNT=${LAYER_COUNT:-0}
EXPERT_COUNT=${EXPERT_COUNT:-0}
HEAD_COUNT_KV=${HEAD_COUNT_KV:-0}
KEY_LENGTH=${KEY_LENGTH:-0}
VALUE_LENGTH=${VALUE_LENGTH:-0}
HAS_SSM=${HAS_SSM:-0}
HAS_FUSED=${HAS_FUSED:-0}
EXPERT_BYTES=${EXPERT_BYTES:-0}
NON_EXPERT_BYTES=${NON_EXPERT_BYTES:-0}

[[ "$LAYER_COUNT" == "0" ]] && LAYER_COUNT=48 && echo "Warning: Could not detect layer count, using default (48)"

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")"

(( HAS_FUSED && IS_IK_LLAMA )) && echo "Info: Fused up|gate model detected. Optimized ik_llama.cpp kernels enabled."

# Health check timeout: 240s base, +60s per 100GB of model size
HEALTH_TIMEOUT=$(( 240 + TOTAL_SIZE_MB / 1700 ))

# ═══════════════════════════════════════════════════════════════
# Section 4: Smart Flag Builder
# ═══════════════════════════════════════════════════════════════

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

# ── Compute totals for decisions ──
TOTAL_VRAM_MB=0
BEST_GPU_VRAM=0
for i in $(seq 0 $(( GPU_COUNT - 1 ))); do
    (( TOTAL_VRAM_MB += GPU_VRAM_FREE[$i] )) || true
    (( GPU_VRAM_FREE[$i] > BEST_GPU_VRAM )) && BEST_GPU_VRAM=${GPU_VRAM_FREE[$i]}
done

# ── Memory Analysis (Informational) ──
TOTAL_MEM_MB=$(( TOTAL_VRAM_MB + RAM_AVAIL_MB ))
ESTIMATED_TOTAL_NEEDED=$(( TOTAL_SIZE_MB + TOTAL_SIZE_MB / 10 + SYSTEM_HEADROOM_MB )) # Model + ~10% overhead + system headroom

if (( TOTAL_SIZE_MB > TOTAL_MEM_MB )); then
    echo "⚠️  WARNING: Model (${TOTAL_SIZE_GB}GB) is larger than your total available memory (${TOTAL_MEM_MB}MB)."
    echo "   A system crash or heavy swapping is likely. Proceeding anyway..."
elif (( ESTIMATED_TOTAL_NEEDED > TOTAL_MEM_MB )); then
    echo "ℹ️  INFO: Model + KV Cache (~${ESTIMATED_TOTAL_NEEDED}MB) is very close to your total memory (${TOTAL_MEM_MB}MB)."
    echo "   Expect potential instability or slow performance if memory fills up."
else
    echo "✓ Memory: Model and context estimated to fit within available RAM/VRAM."
fi

FITS_ON_GPU=0
(( TOTAL_SIZE_MB * VRAM_OVERHEAD_PERCENT / 100 <= TOTAL_VRAM_MB )) && FITS_ON_GPU=1

# Free RAM after model is loaded (estimate: model goes to VRAM if it fits, else spills)
if (( FITS_ON_GPU )); then
    RAM_AFTER_LOAD=$RAM_AVAIL_MB
else
    RAM_ON_CPU=$(( TOTAL_SIZE_MB - TOTAL_VRAM_MB ))
    (( RAM_ON_CPU < 0 )) && RAM_ON_CPU=0
    RAM_AFTER_LOAD=$(( RAM_AVAIL_MB - RAM_ON_CPU ))
    (( RAM_AFTER_LOAD < 0 )) && RAM_AFTER_LOAD=0
fi

# ── Batch sizes: scale with available VRAM ──
# Larger batches = faster prompt processing but use more VRAM
# Small models with lots of VRAM headroom → max batches
# Large models barely fitting → conservative batches
if (( FITS_ON_GPU && BEST_GPU_VRAM > TOTAL_SIZE_MB + SINGLE_GPU_HEADROOM_MB )); then
    # Tons of VRAM headroom
    BATCH=8192; UBATCH=1024
    echo "  VRAM headroom large → batch=$BATCH ubatch=$UBATCH"
elif (( FITS_ON_GPU )); then
    BATCH=4096; UBATCH=512
    echo "  Model fits on GPU → batch=$BATCH ubatch=$UBATCH"
else
    # MoE offload / CPU spill — GPU VRAM is tight
    BATCH=2048; UBATCH=512
    echo "  GPU+CPU split → batch=$BATCH ubatch=$UBATCH"
fi

# ── KV cache type: high=f16, mid=q8_0 (default), low=q4_0, auto=fit-based ──
# KV cache size estimates per 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 VRAM"
        ;;
    *)
        echo "Error: --kv-quality must be one of: high, mid, low"
        exit 1
        ;;
esac

# Set KV_TOTAL_MB based on selected type (needed for OOM checks in all strategies)
case "$KV_TYPE" in
    f16)   KV_TOTAL_MB=$KV_F16_MB ;;
    q8_0)  KV_TOTAL_MB=$KV_Q8_MB ;;
    q4_0)  KV_TOTAL_MB=$KV_Q4_MB ;;
esac

# ── Hadamard K-cache: only useful with quantized KV cache ──
KHAD_FLAG=""
if [[ "$KV_TYPE" == "q4_0" || "$KV_TYPE" == "q8_0" ]]; then
    KHAD_FLAG="-khad"
    echo "  Quantized KV cache → Hadamard K-transform enabled"
fi

# ── Prompt cache: scale with free RAM, cap at 10% ──
CRAM_MB=$(( RAM_AFTER_LOAD / 10 ))
(( CRAM_MB > 16384 )) && CRAM_MB=16384  # cap at 16GB
(( CRAM_MB < MIN_CRAM_MB )) && CRAM_MB=0  # disable if too small
if (( GPU_COUNT <= 1 )); then
    if (( CRAM_MB > 0 )); then
        echo "  Prompt cache: ${CRAM_MB}MB (10% of free RAM)"
    else
        echo "  Prompt cache: disabled (not enough free RAM)"
    fi
fi

# ── Thread count ──
# threads: physical cores for generation (hyperthreads hurt)
# threads-batch: can use all cores for prompt processing (parallel work)
THREADS_GEN=$PHYSICAL_CORES
THREADS_BATCH=$PHYSICAL_CORES
echo "  Threads: gen=$THREADS_GEN batch=$THREADS_BATCH (${PHYSICAL_CORES} physical cores)"

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

# Conditional flags
if [[ "$IS_IK_LLAMA" == "1" ]]; then
    COMMON_FLAGS+=(--run-time-repack)
    [[ -n "$KHAD_FLAG" ]] && COMMON_FLAGS+=("$KHAD_FLAG")
    [[ -n "$CONTEXT_SHIFT_FLAG" ]] && COMMON_FLAGS+=("$CONTEXT_SHIFT_FLAG")
    COMMON_FLAGS+=(--defrag-thold 0.1)

    # MoE-specific optimizations
    if (( IS_MOE == 1 )); then
        COMMON_FLAGS+=(-muge)           # merge up/gate expert tensors
        COMMON_FLAGS+=(-ger)            # grouped expert routing
        echo "  MoE model → -muge -ger enabled"
    fi

    if (( GPU_COUNT > 0 )); then
        COMMON_FLAGS+=(-mqkv)
    fi

    # ── Prompt cache budget ──
    # Multi-GPU: checkpoints can OOM smaller GPUs (cuBLAS workspace fails silently).
    # Cap based on VRAM headroom after model+KV+compute.
    # Single-GPU / CPU-only: use RAM-based default (CRAM_MB from above).
    if (( GPU_COUNT > 1 )); then
        MODEL_ON_GPU_MB=$(( TOTAL_SIZE_MB * VRAM_OVERHEAD_PERCENT / 100 ))
        (( MODEL_ON_GPU_MB > TOTAL_VRAM_MB )) && MODEL_ON_GPU_MB=$TOTAL_VRAM_MB
        VRAM_HEADROOM=$(( TOTAL_VRAM_MB - MODEL_ON_GPU_MB - KV_TOTAL_MB - COMPUTE_PER_GPU_MB * GPU_COUNT ))
        (( VRAM_HEADROOM < 0 )) && VRAM_HEADROOM=0
        CACHE_RAM_MB=$(( VRAM_HEADROOM / 2 ))
        (( CACHE_RAM_MB > 4096 )) && CACHE_RAM_MB=4096
        if (( CACHE_RAM_MB < 256 )); then
            CACHE_RAM_MB=0
            MAX_CHECKPOINTS=0
            COMMON_FLAGS+=(-cram 0 --ctx-checkpoints 0)
        else
            MAX_CHECKPOINTS=$(( CACHE_RAM_MB / 200 ))
            (( MAX_CHECKPOINTS < 2 )) && MAX_CHECKPOINTS=2
            (( MAX_CHECKPOINTS > 16 )) && MAX_CHECKPOINTS=16
            COMMON_FLAGS+=(-cram "$CACHE_RAM_MB" --ctx-checkpoints "$MAX_CHECKPOINTS")
        fi
        echo "  Prompt cache: ${CACHE_RAM_MB}MB (VRAM headroom: ${VRAM_HEADROOM}MB), checkpoints: ${MAX_CHECKPOINTS}"
    else
        (( CRAM_MB > 0 )) && COMMON_FLAGS+=(-cram "$CRAM_MB")
    fi
else
    # Mainline-specific or compatible flags
    [[ -n "$CONTEXT_SHIFT_FLAG" ]] && COMMON_FLAGS+=("$CONTEXT_SHIFT_FLAG")
fi

# GPU offloading
if (( GPU_COUNT > 0 )); then
    COMMON_FLAGS+=(-ngl 999 -mg "${GPU_INDEX[${GPU_ORDER[0]}]}")
fi

# ═══════════════════════════════════════════════════════════════
# Section 5: Strategy Selection
# ═══════════════════════════════════════════════════════════════

# OOM guard: refuse to launch if model+KV+compute don't fit in the given memory pool
# Args: pool_mb pool_label
check_memory_or_die() {
    local pool_mb=$1
    local pool_label=$2
    local model_overhead_mb=$(( TOTAL_SIZE_MB * VRAM_OVERHEAD_PERCENT / 100 ))
    local needed_mb=$(( model_overhead_mb + KV_TOTAL_MB + COMPUTE_PER_GPU_MB ))

    if (( needed_mb > pool_mb )); then
        # Back-solve max safe context: max_kv = pool - model_overhead - compute
        local max_kv_mb=$(( pool_mb - model_overhead_mb - COMPUTE_PER_GPU_MB ))
        (( max_kv_mb < 0 )) && max_kv_mb=0
        local max_ctx=0
        if (( KV_TOTAL_MB > 0 )); then
            max_ctx=$(( max_kv_mb * CTX_SIZE / KV_TOTAL_MB ))
        fi
        echo ""
        echo "ERROR: Model does not fit in ${pool_label}."
        echo "  Model (with overhead): ${model_overhead_mb}MB"
        echo "  KV cache (ctx=${CTX_SIZE}): ${KV_TOTAL_MB}MB"
        echo "  Compute buffers:       ${COMPUTE_PER_GPU_MB}MB"
        echo "  ─────────────────────────────"
        echo "  Total needed:          ${needed_mb}MB"
        echo "  Available (${pool_label}): ${pool_mb}MB"
        echo "  Shortfall:             $(( needed_mb - pool_mb ))MB"
        echo ""
        if (( max_ctx > 0 )); then
            echo "  Max safe context at this memory: --ctx-size $max_ctx"
        else
            echo "  Model weights alone exceed available memory."
        fi
        echo "  Or use a smaller quantization / model."
        exit 1
    fi
}

VRAM_NEEDED_MB=$(( TOTAL_SIZE_MB * VRAM_OVERHEAD_PERCENT / 100 + KV_TOTAL_MB + COMPUTE_PER_GPU_MB ))

# Global OOM check: model must fit somewhere (all GPUs + RAM combined)
# Skip for MoE — expert offload has its own fine-grained budgeting (Phase 1+2)
if (( IS_MOE == 0 )); then
    TOTAL_POOL_MB=$(( TOTAL_VRAM_MB + RAM_AVAIL_MB - SYSTEM_HEADROOM_MB ))
    check_memory_or_die "$TOTAL_POOL_MB" "total GPU+RAM"
fi

choose_strategy() {
    # CPU-only
    if (( GPU_COUNT == 0 )); then
        echo "cpu_only"
        return
    fi

    # Single GPU: model + overhead fits in best GPU
    # Use tighter estimate (110%) — quantized weights don't expand much in VRAM,
    # and 130% global overhead is too conservative for single-GPU feasibility
    local best=${GPU_ORDER[0]}
    local single_gpu_needed=$(( TOTAL_SIZE_MB * 110 / 100 + KV_TOTAL_MB + COMPUTE_PER_GPU_MB ))
    if (( single_gpu_needed <= GPU_VRAM_FREE[best] )); then
        echo "single_gpu"
        return
    fi

    # Multi-GPU: model fits across all GPUs
    if (( IS_MOE == 0 && VRAM_NEEDED_MB <= TOTAL_VRAM_MB )); then
        echo "multi_gpu_dense"
        return
    fi

    # MoE expert offload
    if (( IS_MOE == 1 )); then
        echo "moe_offload"
        return
    fi

    # Dense model with CPU spill
    echo "dense_cpu_offload"
}

STRATEGY=$(choose_strategy)
echo ""
echo "Strategy: $STRATEGY"

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

# Build -ot string from array of (cuda_device, layer_start, layer_count) assignments
# Args: pairs of "CUDA_INDEX LAYER_START LAYER_COUNT" followed by end marker
build_ot_string() {
    local parts=()
    # Support both fused (gate_up/up_gate) and unfused (gate|up|down) expert tensor names
    # Longer alternatives first so regex engine matches them before shorter prefixes
    local expert_pattern="ffn_(gate_up|up_gate|gate|up|down)_exps"

    while [[ $# -ge 3 ]]; do
        local cuda_idx=$1 layer_start=$2 layer_count=$3
        shift 3
        if (( layer_count > 0 )); then
            local last=$(( layer_start + layer_count - 1 ))
            local re=$(seq "$layer_start" "$last" | tr '\n' '|' | sed 's/|$//')
            parts+=("blk\\.(${re})\\.${expert_pattern}.*=CUDA${cuda_idx}")
        fi
    done
    parts+=("exps=CPU")
    local IFS=,
    echo "${parts[*]}"
}

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
    # Aggressive 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
    # Wait for port to free (max 10s)
    local w=0
    while (( w < 10 )) && lsof -i:"${PORT}" >/dev/null 2>&1; do
        sleep 1; (( w++ ))
    done
}

# Try starting server, wait for health. Sets RUNNING_PID on success.
RUNNING_PID=""
try_start() {
    local extra_flags=("$@")
    RUNNING_PID=""
    # Use process substitution to get the REAL pid of the server, not the tee pid
    "$LLAMA_SERVER" "${COMMON_FLAGS[@]}" "${extra_flags[@]}" > >(tee -a "$SERVER_LOG") 2>&1 &
    local pid=$!
    local i=0
    while (( i < HEALTH_TIMEOUT )); 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]}"
        # Pair --flag with its value (next arg that doesn't start with -)
        if [[ "$arg" == -* ]] && (( i + 1 < ${#args[@]} )) && [[ "${args[$((i+1))]}" != -* ]]; then
            local pair="$arg ${args[$((i+1))]}"
            i=$(( i + 1 ))
        else
            local pair="$arg"
        fi
        # Wrap lines at ~80 chars
        if (( ${#line} + ${#pair} + 1 > 80 )); then
            echo "$line \\"
            line="    $pair"
        else
            line="$line $pair"
        fi
    done
    echo "$line"
}

# Run benchmark: send a prompt, measure tok/s
run_benchmark() {
    local url="http://127.0.0.1:${PORT}"
    echo ""
    echo "═══ Benchmark ═══"

    # Prompt processing benchmark (short prompt)
    local pp_prompt="Explain the theory of relativity in simple terms. Cover special and general relativity, time dilation, and gravitational effects."
    local result
    result=$(curl -sf "$url/v1/chat/completions" \
        -H "Content-Type: application/json" \
        -d "{\"model\":\"test\",\"messages\":[{\"role\":\"user\",\"content\":\"$pp_prompt\"}],\"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 pp_time tg_tokens tg_time
    pp_tokens=$(echo "$result" | python3 -c "import sys,json; d=json.load(sys.stdin); u=d.get('usage',{}); print(u.get('prompt_tokens',0))" 2>/dev/null)
    tg_tokens=$(echo "$result" | python3 -c "import sys,json; d=json.load(sys.stdin); u=d.get('usage',{}); print(u.get('completion_tokens',0))" 2>/dev/null)

    # Use timings endpoint if available (ik_llama.cpp / llama.cpp specific)
    local timings
    timings=$(curl -sf "$url/slots" 2>/dev/null)
    if [[ -n "$timings" ]]; then
        local pp_rate tg_rate
        # Improved extraction from slot timings
        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:-?}"
        echo "  (Install ik_llama.cpp for detailed timing)"
    fi
    echo ""
    # Auto-exit after benchmark to prevent "hanging"
    echo "Benchmark complete. Shutting down server..."
    kill_server ""
    exit 0
}

# Suggest alternatives when model doesn't fit
# Run server with automatic crash restart and backoff
# Args: all flags to pass to llama-server
run_with_restart() {
    local flags=("$@")
    local restarts=0
    local last_start=0

    while true; do
        last_start=$(date +%s)
        local log_start
        log_start=$(wc -c < "$SERVER_LOG" 2>/dev/null || echo 0)
        echo "Starting server..."
        "$LLAMA_SERVER" "${flags[@]}" > >(tee -a "$SERVER_LOG" | grep -v --line-buffered "record AND rewind is invalid") 2>&1 &
        local pid=$!
        echo "Server PID: $pid"

        # Wait for health before declaring success
        local healthy=0
        for i in $(seq 1 $HEALTH_TIMEOUT); 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  # only run once
            fi
            restarts=0  # reset counter on successful start

            # Monitor log for CUDA errors — server stays "alive" but braindead after these
            local cuda_killed=0
            local log_pos=$log_start
            while kill -0 "$pid" 2>/dev/null; do
                sleep 5
                if tail -c +"$log_pos" "$SERVER_LOG" 2>/dev/null | grep -q "CUDA error"; then
                    echo ""
                    echo "═══ CUDA error detected — server is braindead, killing ═══"
                    tail -5 "$SERVER_LOG" 2>/dev/null
                    kill -9 "$pid" 2>/dev/null; wait "$pid" 2>/dev/null || true
                    cuda_killed=1
                    break
                fi
                log_pos=$(wc -c < "$SERVER_LOG" 2>/dev/null || echo "$log_pos")
            done

            local exit_code=0
            if (( cuda_killed )); then
                # CUDA kill is a recovery, not a repeated crash — don't count against restarts
                echo "Restarting after CUDA error recovery..."
                kill_server ""
                # Wait for GPU driver to fully reclaim CUDA contexts (up to 60s)
                local gpu_wait=0
                echo "Waiting for GPU memory to be released..."
                while (( gpu_wait < 60 )); do
                    local gpu_procs
                    gpu_procs=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader 2>/dev/null | grep -c '[0-9]') || gpu_procs=0
                    if (( gpu_procs == 0 )); then
                        echo "GPU memory free after ${gpu_wait}s"
                        break
                    fi
                    sleep 2
                    (( gpu_wait += 2 ))
                done
                if (( gpu_wait >= 60 )); then
                    echo "Warning: GPU processes still present after 60s, attempting restart anyway"
                fi
                sleep 3  # extra settle time for driver cleanup
                restarts=-3  # extra patience: 8 total attempts (3 extra + 5 normal)
                continue
            fi
            wait "$pid" 2>/dev/null && exit_code=$? || 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

        # Server exited — decide whether to restart
        local uptime=$(( $(date +%s) - last_start ))

        # Clean exit (signal 0 or SIGTERM)
        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

        # Backoff: 2s, 4s, 8s, 16s, 32s
        local wait_sec=$(( 2 ** restarts ))
        (( wait_sec > 30 )) && wait_sec=30

        # If server ran for >60s, it was a runtime crash (not startup OOM)
        # Reset backoff since the config itself works
        if (( uptime > 60 )); then
            wait_sec=2
            # Still count runtime crashes — don't loop forever
            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"

        # Clean up port before restart
        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 GPUs detected)"
        fi
        check_memory_or_die "$RAM_AVAIL_MB" "RAM"
        # Build CPU-optimized flags from scratch (no GPU flags)
        LAUNCH_FLAGS=(
            -m "$MODEL_PATH" --host "$HOST" --port "$PORT"
            --ctx-size "$CTX_SIZE" --flash-attn on
            -b 512 -ub 256 -ngl 0 --no-mmap --mlock
            --cache-type-k "$KV_TYPE" --cache-type-v "$KV_TYPE"
            --jinja --threads "$THREADS_GEN" --threads-batch "$THREADS_BATCH"
        )
        [[ -n "$CONTEXT_SHIFT_FLAG" ]] && LAUNCH_FLAGS+=("$CONTEXT_SHIFT_FLAG")
        if (( DRY_RUN )); then
            print_cmd "${LAUNCH_FLAGS[@]}"
            exit 0
        fi
        run_with_restart "${LAUNCH_FLAGS[@]}"
        ;;

    single_gpu)
        local_best=${GPU_ORDER[0]}
        echo "Model fits on GPU${GPU_INDEX[$local_best]} (${GPU_NAME[$local_best]})"
        LAUNCH_FLAGS=("${COMMON_FLAGS[@]}" --device "CUDA${GPU_INDEX[$local_best]}")
        if (( DRY_RUN )); then
            print_cmd "${LAUNCH_FLAGS[@]}"
            exit 0
        fi
        run_with_restart "${LAUNCH_FLAGS[@]}"
        ;;

    multi_gpu_dense)
        echo "Dense model split across $GPU_COUNT GPUs"
        # Build tensor-split proportional to free VRAM weighted by bandwidth
        total_weighted=0
        for i in $(seq 0 $(( GPU_COUNT - 1 ))); do
            (( total_weighted += GPU_VRAM_FREE[$i] * GPU_BANDWIDTH[$i] )) || true
        done
        TS=""
        for i in $(seq 0 $(( GPU_COUNT - 1 ))); do
            w=$(awk "BEGIN {printf \"%.2f\", ${GPU_VRAM_FREE[$i]} * ${GPU_BANDWIDTH[$i]} / $total_weighted}")
            TS="${TS:+$TS,}$w"
        done
        echo "Tensor split: $TS"
        
        # Prefer graph split for both backends if more than 1 GPU
        GRAPH_EXTRA=(--split-mode graph --tensor-split "$TS")
        
        if (( IS_IK_LLAMA )); then
            LAYER_EXTRA=(--split-mode layer --tensor-split "$TS")
            if (( DRY_RUN )); then
                print_cmd "${COMMON_FLAGS[@]}" "${GRAPH_EXTRA[@]}"
                exit 0
            fi
            echo "Trying graph split (fastest)..."
            if try_start "${GRAPH_EXTRA[@]}"; then
                echo "  Graph split works — restarting with crash recovery..."
                kill_server "$RUNNING_PID"
                run_with_restart "${COMMON_FLAGS[@]}" "${GRAPH_EXTRA[@]}"
            else
                echo "  Graph split failed — retrying with NCCL_P2P_DISABLE=1..."
                kill_server ""
                export NCCL_P2P_DISABLE=1
                if try_start "${GRAPH_EXTRA[@]}"; then
                    echo "  Graph split works with P2P disabled — restarting with crash recovery..."
                    kill_server "$RUNNING_PID"
                    run_with_restart "${COMMON_FLAGS[@]}" "${GRAPH_EXTRA[@]}"
                else
                    echo "  Graph split unavailable — falling back to layer split"
                    kill_server ""
                    unset NCCL_P2P_DISABLE
                    run_with_restart "${COMMON_FLAGS[@]}" "${LAYER_EXTRA[@]}"
                fi
            fi
        else
            # Mainline llama.cpp
            if (( DRY_RUN )); then
                print_cmd "${COMMON_FLAGS[@]}" "${GRAPH_EXTRA[@]}"
                exit 0
            fi
            # Test if mainline supports graph mode
            if timeout 5 "$LLAMA_SERVER" --help 2>&1 | grep -q "graph" 2>/dev/null; then
                echo "Mainline supports graph mode — attempting launch..."
                run_with_restart "${COMMON_FLAGS[@]}" "${GRAPH_EXTRA[@]}"
            else
                # Fallback to row split (classic mainline multi-GPU)
                LAUNCH_FLAGS=("${COMMON_FLAGS[@]}" --split-mode row --tensor-split "$TS")
                run_with_restart "${LAUNCH_FLAGS[@]}"
            fi
        fi
        ;;

    dense_cpu_offload)
        echo "Dense model with CPU offload"
        check_memory_or_die "$(( TOTAL_VRAM_MB + RAM_AVAIL_MB - SYSTEM_HEADROOM_MB ))" "GPU+RAM"
        if (( GPU_COUNT > 1 )); then
            GRAPH_EXTRA=(--split-mode graph --no-mmap)
            if (( IS_IK_LLAMA )); then
                # (ik_llama specific graph logic already handled above or can be simplified here)
                if (( DRY_RUN )); then print_cmd "${COMMON_FLAGS[@]}" "${GRAPH_EXTRA[@]}"; exit 0; fi
                run_with_restart "${COMMON_FLAGS[@]}" "${GRAPH_EXTRA[@]}"
            else
                # Mainline
                if timeout 5 "$LLAMA_SERVER" --help 2>&1 | grep -q "graph" 2>/dev/null; then
                    run_with_restart "${COMMON_FLAGS[@]}" "${GRAPH_EXTRA[@]}"
                else
                    run_with_restart "${COMMON_FLAGS[@]}" --no-mmap
                fi
            fi
        else
            LAUNCH_FLAGS=("${COMMON_FLAGS[@]}" --no-mmap)
            if (( DRY_RUN )); then
                print_cmd "${LAUNCH_FLAGS[@]}"
                exit 0
            fi
            run_with_restart "${LAUNCH_FLAGS[@]}"
        fi
        ;;

    moe_offload)
        # MoE expert placement and auto-tuning follows
        ;;
    *)
        echo "Error: Unknown strategy: $STRATEGY"
        exit 1
        ;;
esac

# Only MoE offload reaches here — all other strategies call run_with_restart (never returns)
[[ "$STRATEGY" != "moe_offload" ]] && exit 0

# ═══════════════════════════════════════════════════════════════
# Section 8: MoE Expert Placement & Auto-Tuning
# ═══════════════════════════════════════════════════════════════

mkdir -p "$CACHE_DIR"
CACHE_KEY_INPUT="${MODEL_NAME}_${TOTAL_SIZE_BYTES}_ctx${CTX_SIZE}_kv${KV_QUALITY}_gpus${GPU_COUNT}"
[[ -n "$GPUS_FILTER" ]] && CACHE_KEY_INPUT="${CACHE_KEY_INPUT}_filter${GPUS_FILTER}"
CACHE_KEY=$(echo "$CACHE_KEY_INPUT" | md5sum | awk '{print $1}')
CACHE_FILE="$CACHE_DIR/${CACHE_KEY}.conf"

# Per-layer expert size from GGUF tensor data (fallback: estimate 90%)
if (( EXPERT_BYTES > 0 )); then
    EXPERT_TOTAL_MB=$(( EXPERT_BYTES / 1048576 ))
    NON_EXPERT_TOTAL_MB=$(( NON_EXPERT_BYTES / 1048576 ))
    EXPERT_PER_LAYER_MB=$(( EXPERT_TOTAL_MB / LAYER_COUNT ))
    EXPERT_PCT=$(( EXPERT_TOTAL_MB * 100 / TOTAL_SIZE_MB ))
    echo "Expert size per layer: ${EXPERT_PER_LAYER_MB}MB (${EXPERT_PCT}% expert, from GGUF tensors)"
else
    EXPERT_TOTAL_MB=$(awk "BEGIN {printf \"%.0f\", $TOTAL_SIZE_MB * 0.90}")
    EXPERT_PER_LAYER_MB=$(awk "BEGIN {printf \"%.0f\", $EXPERT_TOTAL_MB / $LAYER_COUNT}")
    echo "Expert size per layer: ~${EXPERT_PER_LAYER_MB}MB (estimated 90%, no tensor data)"
fi

# KV_TOTAL_MB already set in Section 4 after KV type selection

# Check cached config
if [[ -f "$CACHE_FILE" && "$RETUNE" == "0" ]]; then
    echo ""
    echo "Using cached config: $CACHE_FILE"
    # Validate cache file contains only expected content (comments and CACHED_GPU_ASSIGNMENTS)
    if grep -qvE '^(#.*|CACHED_GPU_ASSIGNMENTS="[0-9: ]*"|[[:space:]]*)$' "$CACHE_FILE"; then
        echo "WARNING: Cache file contains unexpected content, ignoring."
        rm -f "$CACHE_FILE"
    else
        source "$CACHE_FILE"
    fi
fi

if [[ -f "$CACHE_FILE" && "$RETUNE" == "0" && -n "${CACHED_GPU_ASSIGNMENTS:-}" ]]; then
    # Validate cached config against current VRAM
    CACHE_VALID=1
    OT_ARGS=()
    for assignment in $CACHED_GPU_ASSIGNMENTS; do
        IFS=: read -r ci cs cc <<< "$assignment"
        OT_ARGS+=("$ci" "$cs" "$cc")
        # Check this GPU has enough VRAM for assigned layers + its share of KV cache
        layer_mb=$(( cc * EXPERT_PER_LAYER_MB ))
        kv_share_mb=$(( KV_TOTAL_MB * cc / LAYER_COUNT + COMPUTE_PER_GPU_MB ))  # proportional KV + compute
        needed=$(( layer_mb + kv_share_mb ))
        # Find free VRAM for this CUDA index
        for gi in $(seq 0 $(( GPU_COUNT - 1 ))); do
            if [[ "${GPU_INDEX[$gi]}" == "$ci" ]]; then
                if (( needed > GPU_VRAM_FREE[$gi] )); then
                    echo "  WARNING: GPU${ci} needs ${needed}MB but only ${GPU_VRAM_FREE[$gi]}MB free"
                    CACHE_VALID=0
                fi
                break
            fi
        done
        echo "  GPU${ci}: ${cc} layers (starting at block ${cs})"
    done

    if (( CACHE_VALID )); then
        OT_STRING=$(build_ot_string "${OT_ARGS[@]}")
        LAUNCH_FLAGS=("${COMMON_FLAGS[@]}" --no-mmap -ot "$OT_STRING")
        if (( DRY_RUN )); then
            print_cmd "${LAUNCH_FLAGS[@]}"
            exit 0
        fi
        run_with_restart "${LAUNCH_FLAGS[@]}"
    else
        echo "  Cached config invalid for current VRAM — retuning..."
        rm -f "$CACHE_FILE"
    fi
fi

# Calculate initial conservative placement
# Budget: GPUs get decreasing % based on priority (fastest gets most headroom for compute)
# Reserve VRAM on primary GPU for KV cache + batch buffers + compute scratch
declare -a LAYERS_PER_GPU
BUDGET_PCTS=(70 55 50 45 40 35 30 25)  # Decreasing budget for lower-priority GPUs
NEXT_LAYER=0
TOTAL_GPU_LAYERS=0

# Per-GPU VRAM budget: each layer on a GPU costs expert + non-expert + KV
# With -ngl 999, non-expert tensors (attention, norms) also go to GPUs
# Primary GPU (-mg 0) also gets global tensors (embeddings, output head)
KV_PER_LAYER_MB=$(( KV_TOTAL_MB / LAYER_COUNT ))
NON_EXPERT_RAM_MB=${NON_EXPERT_TOTAL_MB:-$(( TOTAL_SIZE_MB - EXPERT_TOTAL_MB ))}
NON_EXPERT_PER_LAYER_MB=$(( NON_EXPERT_RAM_MB / LAYER_COUNT ))
# COMPUTE_PER_GPU_MB defined in constants section at top
MAIN_GPU_EXTRA_MB=2048  # embeddings + output head on primary GPU (-mg 0)
COST_PER_LAYER_MB=$(( EXPERT_PER_LAYER_MB + KV_PER_LAYER_MB + NON_EXPERT_PER_LAYER_MB ))
echo "  Per-layer cost: ${EXPERT_PER_LAYER_MB}MB experts + ${NON_EXPERT_PER_LAYER_MB}MB attn/norm + ${KV_PER_LAYER_MB}MB KV = ${COST_PER_LAYER_MB}MB"

for i in $(seq 0 $(( GPU_COUNT - 1 ))); do
    gi=${GPU_ORDER[$i]}
    pct=${BUDGET_PCTS[$i]:-20}
    # Primary GPU needs extra headroom for global tensors (embeddings, output)
    local_overhead=$COMPUTE_PER_GPU_MB
    (( i == 0 )) && local_overhead=$(( COMPUTE_PER_GPU_MB + MAIN_GPU_EXTRA_MB ))
    available_mb=$(( GPU_VRAM_FREE[$gi] - local_overhead ))
    (( available_mb < 0 )) && available_mb=0
    budget_mb=$(( available_mb * pct / 100 ))
    layers=0
    if (( COST_PER_LAYER_MB > 0 )); then
        layers=$(( budget_mb / COST_PER_LAYER_MB ))
    fi
    # Cap to remaining layers
    remain=$(( LAYER_COUNT - NEXT_LAYER ))
    (( layers > remain )) && layers=$remain
    LAYERS_PER_GPU[$gi]=$layers
    NEXT_LAYER=$(( NEXT_LAYER + layers ))
    TOTAL_GPU_LAYERS=$(( TOTAL_GPU_LAYERS + layers ))
done

LAYERS_CPU=$(( LAYER_COUNT - TOTAL_GPU_LAYERS ))

# RAM safety check
# With -ngl 999: non-expert tensors (attention, norms, embeddings) go to GPUs, not RAM
# Only expert tensors for CPU layers stay in RAM
# CUDA_Host pinned memory ~1GB for output + compute buffers
CPU_EXPERT_MB=$(( LAYERS_CPU * EXPERT_PER_LAYER_MB ))
# 15% padding covers alignment, tensor metadata, graph structures, mmap overhead
TENSOR_RAM_PADDED=$(( CPU_EXPERT_MB * 115 / 100 ))
CUDA_HOST_MB=1024   # pinned memory for output + compute transfers
RAM_NEEDED=$(( TENSOR_RAM_PADDED + CUDA_HOST_MB ))

# Phase 1: If CPU layers would OOM, push layers to GPUs one at a time (least disruption)
if (( RAM_NEEDED > RAM_AVAIL_MB && LAYERS_CPU > 0 && GPU_COUNT > 0 )); then
    echo "WARNING: Conservative placement needs ${RAM_NEEDED}MB RAM but only ${RAM_AVAIL_MB}MB available"
    echo "  Redistributing layers to GPUs to fit in RAM..."

    # Recalculate with increasing GPU budgets until CPU fits in RAM
    for bump in 80 85 90 95; do
        NEXT_LAYER=0
        TOTAL_GPU_LAYERS=0
        # Build budget array: primary GPU gets bump%, others get slightly less
        for i in $(seq 0 $(( GPU_COUNT - 1 ))); do
            gi=${GPU_ORDER[$i]}
            pct=$(( bump - i * 5 ))
            (( pct < 30 )) && pct=30
            local_overhead=$COMPUTE_PER_GPU_MB
            (( i == 0 )) && local_overhead=$(( COMPUTE_PER_GPU_MB + MAIN_GPU_EXTRA_MB ))
            available_mb=$(( GPU_VRAM_FREE[$gi] - local_overhead ))
            (( available_mb < 0 )) && available_mb=0
            budget_mb=$(( available_mb * pct / 100 ))
            layers=0
            if (( COST_PER_LAYER_MB > 0 )); then
                layers=$(( budget_mb / COST_PER_LAYER_MB ))
            fi
            remain=$(( LAYER_COUNT - NEXT_LAYER ))
            (( layers > remain )) && layers=$remain
            LAYERS_PER_GPU[$gi]=$layers
            NEXT_LAYER=$(( NEXT_LAYER + layers ))
            TOTAL_GPU_LAYERS=$(( TOTAL_GPU_LAYERS + layers ))
        done
        LAYERS_CPU=$(( LAYER_COUNT - TOTAL_GPU_LAYERS ))
        CPU_EXPERT_MB=$(( LAYERS_CPU * EXPERT_PER_LAYER_MB ))
        TENSOR_RAM_PADDED=$(( CPU_EXPERT_MB * 115 / 100 ))
        RAM_NEEDED=$(( TENSOR_RAM_PADDED + CUDA_HOST_MB ))
        if (( RAM_NEEDED <= RAM_AVAIL_MB )); then
            echo "  Adjusted GPU budget to ${bump}% — ${LAYERS_CPU} CPU layers (${RAM_NEEDED}MB RAM)"
            break
        fi
    done
fi

# Phase 2: If still over RAM after maxing GPU budgets, refuse to launch
# KV cache lives on GPUs so reducing context doesn't help RAM
if (( RAM_NEEDED > RAM_AVAIL_MB )); then
    echo ""
    echo "ERROR: MoE model does not fit in RAM after maxing GPU budgets."
    echo "  CPU expert layers: ${LAYERS_CPU} layers (${CPU_EXPERT_MB}MB)"
    echo "  CUDA host buffers: ${CUDA_HOST_MB}MB"
    echo "  RAM padding (15%): $(( TENSOR_RAM_PADDED - CPU_EXPERT_MB ))MB"
    echo "  ─────────────────────────────"
    echo "  Total RAM needed:  ${RAM_NEEDED}MB"
    echo "  RAM available:     ${RAM_AVAIL_MB}MB"
    echo "  Shortfall:         $(( RAM_NEEDED - RAM_AVAIL_MB ))MB"
    echo ""
    echo "  Options:"
    echo "  1. Add more system RAM"
    echo "  2. Use a smaller quantization (IQ3 → IQ2)"
    echo "  3. Use a smaller model"
    exit 1
fi

# Build -ot string
OT_ARGS=()
NEXT_LAYER=0
echo ""
echo "Expert placement (conservative):"
for i in $(seq 0 $(( GPU_COUNT - 1 ))); do
    gi=${GPU_ORDER[$i]}
    layers=${LAYERS_PER_GPU[$gi]}
    if (( layers > 0 )); then
        OT_ARGS+=("${GPU_INDEX[$gi]}" "$NEXT_LAYER" "$layers")
        echo "  GPU${GPU_INDEX[$gi]} (${GPU_NAME[$gi]}): ${layers} layers (blk ${NEXT_LAYER}-$((NEXT_LAYER + layers - 1)))"
        NEXT_LAYER=$(( NEXT_LAYER + layers ))
    fi
done
echo "  CPU (RAM): ${LAYERS_CPU} layers (~${CPU_EXPERT_MB}MB)"

OT_STRING=$(build_ot_string "${OT_ARGS[@]}")

if (( DRY_RUN )); then
    print_cmd "${COMMON_FLAGS[@]}" --no-mmap -ot "$OT_STRING"
    exit 0
fi

# ── Auto-tuning loop ──
echo ""
echo "Starting server (conservative)..."
if try_start --no-mmap -ot "$OT_STRING"; then
    echo "  Server healthy (PID $RUNNING_PID)"

    # Measure actual VRAM after load (server is already healthy, VRAM is settled)
    echo "  Measuring VRAM headroom..."
    declare -a ACTUAL_FREE
    for i in $(seq 0 $(( GPU_COUNT - 1 ))); do
        gi=${GPU_ORDER[$i]}
        ACTUAL_FREE[$gi]=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits -i "${GPU_INDEX[$gi]}" 2>/dev/null | tr -d ' ')
        log "  GPU${GPU_INDEX[$gi]}: ${ACTUAL_FREE[$gi]}MB free after load"
    done

    # Calculate extra layers per GPU
    # IMPORTANT: ACTUAL_FREE doesn't include KV cache yet (allocated on first prompt)
    # Subtract current KV share per GPU before calculating extra capacity
    SAFETY_MB=$COMPUTE_PER_GPU_MB
    TOTAL_EXTRA=0
    declare -a EXTRA_PER_GPU
    for i in $(seq 0 $(( GPU_COUNT - 1 ))); do
        gi=${GPU_ORDER[$i]}
        extra=0
        # Subtract KV cache that will be allocated for layers already on this GPU
        current_kv=$(( LAYERS_PER_GPU[$gi] * KV_PER_LAYER_MB ))
        usable_free=$(( ACTUAL_FREE[$gi] - current_kv - SAFETY_MB ))
        (( usable_free < 0 )) && usable_free=0
        if (( COST_PER_LAYER_MB > 0 )); then
            extra=$(( usable_free / COST_PER_LAYER_MB ))
        fi
        (( extra < 0 )) && extra=0
        EXTRA_PER_GPU[$gi]=$extra
        (( TOTAL_EXTRA += extra ))
    done

    # Apply extras (priority order), cap to remaining CPU layers
    if (( TOTAL_EXTRA > 0 && LAYERS_CPU > 0 )); then
        remaining_cpu=$LAYERS_CPU
        for i in $(seq 0 $(( GPU_COUNT - 1 ))); do
            gi=${GPU_ORDER[$i]}
            add=$(( EXTRA_PER_GPU[$gi] < remaining_cpu ? EXTRA_PER_GPU[$gi] : remaining_cpu ))
            LAYERS_PER_GPU[$gi]=$(( LAYERS_PER_GPU[$gi] + add ))
            remaining_cpu=$(( remaining_cpu - add ))
            (( remaining_cpu <= 0 )) && break
        done
        LAYERS_CPU=$remaining_cpu
        TOTAL_GPU_LAYERS=$(( LAYER_COUNT - LAYERS_CPU ))

        # Rebuild -ot string with optimized placement
        OT_ARGS=()
        NEXT_LAYER=0
        echo ""
        echo "Optimized placement:"
        for i in $(seq 0 $(( GPU_COUNT - 1 ))); do
            gi=${GPU_ORDER[$i]}
            layers=${LAYERS_PER_GPU[$gi]}
            if (( layers > 0 )); then
                OT_ARGS+=("${GPU_INDEX[$gi]}" "$NEXT_LAYER" "$layers")
                echo "  GPU${GPU_INDEX[$gi]} (${GPU_NAME[$gi]}): ${layers} layers (blk ${NEXT_LAYER}-$((NEXT_LAYER + layers - 1)))"
                NEXT_LAYER=$(( NEXT_LAYER + layers ))
            fi
        done
        echo "  CPU (RAM): ${LAYERS_CPU} layers"

        OT_STRING=$(build_ot_string "${OT_ARGS[@]}")

        # Restart with optimized config
        kill_server "$RUNNING_PID"
        echo ""
        echo "Restarting with optimized config (Gentle Load enabled)..."
        # --no-mmap is critical for 397B models to prevent virtual memory blowout
        # --mlock ensures the 100GB+ in RAM doesn't get swapped out
        if ! try_start --no-mmap --mlock -ot "$OT_STRING"; then
            echo "  Optimized config failed — retrying without mlock..."
            kill_server ""
            if ! try_start --no-mmap -ot "$OT_STRING"; then
                echo "  FAILED: Could not start server. Check $SERVER_LOG"
                exit 1
            fi
        fi
    else
        echo "  Conservative config is already optimal (GPUs full)"
    fi

    # Save to cache
    CACHED_GPU_ASSIGNMENTS=""
    NEXT_LAYER=0
    for i in $(seq 0 $(( GPU_COUNT - 1 ))); do
        gi=${GPU_ORDER[$i]}
        layers=${LAYERS_PER_GPU[$gi]}
        if (( layers > 0 )); then
            CACHED_GPU_ASSIGNMENTS+="${GPU_INDEX[$gi]}:${NEXT_LAYER}:${layers} "
            NEXT_LAYER=$(( NEXT_LAYER + layers ))
        fi
    done

    cat > "$CACHE_FILE" <<CEOF
# Auto-tuned config for $MODEL_NAME
# Generated: $(date)
# Total layers: $LAYER_COUNT, Experts: $EXPERT_COUNT
CACHED_GPU_ASSIGNMENTS="$CACHED_GPU_ASSIGNMENTS"
CEOF

    TOTAL_GPU_LAYERS=$(( LAYER_COUNT - LAYERS_CPU ))
    echo ""
    echo "═══════════════════════════════════════════"
    echo "  Model: $MODEL_NAME"
    echo "  Size: ${TOTAL_SIZE_GB}GB | Layers: $LAYER_COUNT | Experts: $EXPERT_COUNT"
    echo "  GPU layers: $TOTAL_GPU_LAYERS | CPU layers: $LAYERS_CPU"
    echo "  Context: $CTX_SIZE | Port: $PORT"
    echo "  Config cached: $CACHE_FILE"
    echo "═══════════════════════════════════════════"

    # Kill the tuning server, restart via the crash-safe loop
    kill_server "$RUNNING_PID"
    run_with_restart "${COMMON_FLAGS[@]}" --no-mmap -ot "$OT_STRING"
else
    echo "FAILED: Could not start server even with conservative config."
    echo "Check $SERVER_LOG for details."
    tail -20 "$SERVER_LOG" 2>/dev/null
    echo ""
    echo "═══ Model doesn't fit ═══"
    echo "  Model needs: ~${VRAM_NEEDED_MB}MB (with overhead)"
    echo "  Available:   ~$(( TOTAL_VRAM_MB + RAM_AVAIL_MB ))MB (GPU + RAM)"
    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 (fewer parameters or MoE with fewer active)"
    if (( GPU_COUNT > 0 )); then
        echo "  4. Add more RAM (cheapest upgrade for MoE expert offload)"
    else
        echo "  4. Add a GPU (even a used RTX 3060 12GB helps significantly)"
    fi
    echo ""
    exit 1
fi