#!/usr/bin/env bash
# Purpose: start one vLLM instance with a host-memory tier under its key-value cache, so
#          that blocks evicted from device memory are recalled instead of recomputed, and
#          print the one distinction that matters here: this offloads the CACHE, not the
#          weights
# Platform: spark, nvidia (vLLM's GPU path; Track X only where your ROCm build works)
# Minimum memory: 16 GB of device memory for Qwen3-8B at a four-bit or eight-bit
#          checkpoint, plus KV_OFFLOAD_GB of HOST memory that other programs will not get
# Assumes: vLLM installed as in Part 9 and on PATH; a .env copied from env-example.txt;
#          nothing else listening on DECODE_PORT. Run it once with OFFLOAD=0 for the
#          baseline and once with OFFLOAD=1, and compare the two notebook lines.
#
# Usage: OFFLOAD=0 bash serve-offload.sh     the baseline: device memory only
#        OFFLOAD=1 bash serve-offload.sh     with a host-memory tier of KV_OFFLOAD_GB
set -euo pipefail

HERE="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
if [ -f "${HERE}/.env" ]; then
    set -a
    # shellcheck disable=SC1091  # written by the reader from env-example.txt
    . "${HERE}/.env"
    set +a
fi

MODEL="${MODEL:-Qwen/Qwen3-8B}"
SERVED_NAME="${SERVED_NAME:-local-chat}"
SERVE_HOST="${SERVE_HOST:-127.0.0.1}"
PORT="${DECODE_PORT:-8200}"
CTX="${CTX:-8192}"
MAX_SEQS="${MAX_SEQS:-8}"
MEM_FRACTION="${MEM_FRACTION:-0.85}"
KV_OFFLOAD_GB="${KV_OFFLOAD_GB:-8}"
KV_OFFLOAD_BACKEND="${KV_OFFLOAD_BACKEND:-native}"
OFFLOAD="${OFFLOAD:-1}"

fail() { printf '%s\n' "$*" >&2; exit 1; }

command -v vllm >/dev/null 2>&1 || fail "vllm is not on PATH. Install it as Part 9 describes."

case "$KV_OFFLOAD_BACKEND" in
    native|lmcache) ;;
    *) fail "KV_OFFLOAD_BACKEND must be 'native' or 'lmcache'. Got '${KV_OFFLOAD_BACKEND}'." ;;
esac

ARGS=(
    "$MODEL"
    --host "$SERVE_HOST"
    --port "$PORT"
    --served-model-name "$SERVED_NAME"
    --max-model-len "$CTX"
    --max-num-seqs "$MAX_SEQS"
    --gpu-memory-utilization "$MEM_FRACTION"
    --enable-prefix-caching
)

if [ "$OFFLOAD" = "1" ]; then
    ARGS+=(--kv-offloading-backend "$KV_OFFLOAD_BACKEND" --kv-offloading-size "$KV_OFFLOAD_GB")
    TIER="host memory, ${KV_OFFLOAD_GB} GiB, backend ${KV_OFFLOAD_BACKEND}"
    LABEL_HINT="offload-on"
else
    TIER="none: device memory only"
    LABEL_HINT="offload-off"
fi

cat <<INFO
==> vLLM with a key-value cache tier
    model             ${MODEL}
    served as         ${SERVED_NAME}
    listening on      http://${SERVE_HOST}:${PORT}
    max model length  ${CTX}
    memory fraction   ${MEM_FRACTION}
    prefix caching    enabled
    cache tier        ${TIER}

    Suggested label for the measurement:  ${LABEL_HINT}

    The distinction to keep straight. --kv-offloading-backend and --kv-offloading-size
    move KEY-VALUE BLOCKS into host memory. --cpu-offload-gb is a different option that
    moves MODEL WEIGHTS there, and vLLM's engine argument documentation describes it as a
    virtual way to increase the GPU memory size. They solve different problems and this
    lab is about the first one. If your build does not accept these two options, run
    "vllm serve --help | grep -i offload" and record what it does accept.

    The tier is host memory taken away from everything else on this machine. Start with a
    small KV_OFFLOAD_GB and raise it only if the measurement says it helped.

INFO

exec vllm serve "${ARGS[@]}"