Add 2-replica TP=2 serving for 122B + FP8 KV cache throughput tuning

The single TP=4 server on 4x L40S (no NVLink) pays a per-layer all-reduce tax over PCIe. Since the A10B MoE fits in 2 cards at FP8, run two TP=2 replicas (GPUs 0,1 / 2,3) behind a streaming load balancer on the public port 7080 for better concurrent throughput. - 14_start_replica_122b.sh: one TP=2 replica pinned to a GPU pair - 15_start_replicas_122b.sh: launch both replicas + load balancer - 16_start_loadbalancer.sh + lb_proxy.py: least-in-flight streaming reverse proxy on 7080 -> replicas on 7091/7092 (clear of Open WebUI on 7081) - 17_stop_replicas_122b.sh: stop LB + both replicas - 11_start_server_122b.sh: add --kv-cache-dtype fp8 (~2x more 128k KV slots), --max-num-seqs 16, chunked prefill, gpu-util 0.95 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-09 15:54:38 +02:00 · 2026-06-09 15:54:38 +02:00 · 51726f9351
commit 51726f9351
parent b9eaf2df18
6 changed files with 434 additions and 3 deletions
--- a/11_start_server_122b.sh
+++ b/11_start_server_122b.sh
@ -17,9 +17,22 @@
 #   GPU_MEM_UTIL    - GPU memory utilization fraction (default: 0.92)
 #   API_KEY         - API key for authentication (default: none)
 #   TENSOR_PARALLEL - Number of GPUs for tensor parallelism (default: 4)
-#   MAX_NUM_SEQS    - Max concurrent sequences (default: 4). Raise this
+#   MAX_NUM_SEQS    - Max concurrent sequences (default: 16). Raise this
 #                     when more users need it AND the KV cache fits;
 #                     lower it if 128k context runs out of memory.
+#   MAX_BATCHED_TOK - Token budget per scheduler step (default: 8192).
+#                     With chunked prefill, long 128k prompts are split
+#                     into chunks of this size so they don't block other
+#                     users' decode steps. Key for concurrent throughput.
+#   KV_CACHE_DTYPE  - KV cache quantization (default: fp8). Halves KV
+#                     memory vs bf16 -> ~2x more concurrent 128k slots.
+#                     Set to "auto" to disable.
+#
+# NOTE: For best concurrent throughput on these 4x L40S (no NVLink),
+#       prefer the TWO-REPLICA setup (TP=2 x2 + load balancer):
+#         bash 14_start_replicas_122b.sh
+#       TP=4 pays a per-layer all-reduce tax over PCIe; two TP=2
+#       replicas avoid half of it and double the request slots.
 # ------------------------------------------------------------------
 set -euo pipefail

@ -29,10 +42,12 @@ SIF_FILE="${SCRIPT_DIR}/vllm_qwen.sif"
 MODEL_DIR="${MODEL_DIR:-$HOME/models/Qwen3.5-122B-A10B-FP8}"
 PORT="${PORT:-7080}"
 MAX_MODEL_LEN="${MAX_MODEL_LEN:-131072}"
-GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.92}"
+GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.95}"
 API_KEY="${API_KEY:-}"
 TENSOR_PARALLEL="${TENSOR_PARALLEL:-4}"
-MAX_NUM_SEQS="${MAX_NUM_SEQS:-4}"
+MAX_NUM_SEQS="${MAX_NUM_SEQS:-16}"
+MAX_BATCHED_TOK="${MAX_BATCHED_TOK:-8192}"
+KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-fp8}"

 if [ ! -f "$SIF_FILE" ]; then
    echo "ERROR: Container image not found at ${SIF_FILE}"
@ -61,7 +76,10 @@ VLLM_ARGS=(
    --reasoning-parser qwen3
    --served-model-name "qwen3.5-122b-a10b-fp8"
    --max-num-seqs "$MAX_NUM_SEQS"
+    --max-num-batched-tokens "$MAX_BATCHED_TOK"
+    --kv-cache-dtype "$KV_CACHE_DTYPE"
    --enable-prefix-caching
+    --enable-chunked-prefill
 )

 if [ -n "$API_KEY" ]; then
@ -78,6 +96,8 @@ echo " Context len: ${MAX_MODEL_LEN}"
 echo " GPU util:    ${GPU_MEM_UTIL}"
 echo " TP size:     ${TENSOR_PARALLEL}"
 echo " Max seqs:    ${MAX_NUM_SEQS}"
+echo " Batched tok: ${MAX_BATCHED_TOK} (chunked prefill)"
+echo " KV dtype:    ${KV_CACHE_DTYPE}"
 echo " dtype:       auto (FP8)"
 echo " API key:     ${API_KEY:-<none>}"
 echo "=============================================="
--- a/14_start_replica_122b.sh
+++ b/14_start_replica_122b.sh
@ -0,0 +1,96 @@
+#!/usr/bin/env bash
+# ------------------------------------------------------------------
+# 14_start_replica_122b.sh
+# Launches ONE vLLM replica of Qwen3.5-122B-A10B-FP8 with TP=2,
+# pinned to a specific pair of GPUs, on its own internal port.
+#
+# Two of these (GPUs 0,1 and GPUs 2,3) behind the load balancer
+# (16_start_loadbalancer.sh) beat a single TP=4 server on these
+# L40S cards: TP=4 does a per-layer all-reduce over PCIe (no NVLink),
+# which is pure overhead. Two TP=2 replicas halve that comm cost per
+# token AND double the number of concurrent request slots.
+#
+# This is normally invoked by 15_start_replicas_122b.sh, not directly.
+#
+# Required environment variables:
+#   GPUS            - GPU pair, e.g. "0,1" or "2,3"
+#   PORT            - Internal port for this replica, e.g. 7081 / 7082
+#
+# Optional (sane defaults below):
+#   MODEL_DIR, MAX_MODEL_LEN, GPU_MEM_UTIL, API_KEY,
+#   MAX_NUM_SEQS, MAX_BATCHED_TOK, KV_CACHE_DTYPE
+# ------------------------------------------------------------------
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+SIF_FILE="${SCRIPT_DIR}/vllm_qwen.sif"
+
+: "${GPUS:?Set GPUS, e.g. GPUS=0,1}"
+: "${PORT:?Set PORT, e.g. PORT=7081}"
+
+MODEL_DIR="${MODEL_DIR:-$HOME/models/Qwen3.5-122B-A10B-FP8}"
+MAX_MODEL_LEN="${MAX_MODEL_LEN:-131072}"
+GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.95}"
+API_KEY="${API_KEY:-}"
+# Per-replica concurrency. Two replicas => 2x this many total slots.
+MAX_NUM_SEQS="${MAX_NUM_SEQS:-16}"
+MAX_BATCHED_TOK="${MAX_BATCHED_TOK:-8192}"
+KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-fp8}"
+
+if [ ! -f "$SIF_FILE" ]; then
+    echo "ERROR: Container image not found at ${SIF_FILE}"
+    echo "       Run 01_build_container.sh first."
+    exit 1
+fi
+
+if [ ! -d "$MODEL_DIR" ]; then
+    echo "ERROR: Model directory not found at ${MODEL_DIR}"
+    echo "       Run 10_download_model_122b.sh first."
+    exit 1
+fi
+
+MODEL_PARENT="$(dirname "$MODEL_DIR")"
+MODEL_NAME="$(basename "$MODEL_DIR")"
+
+VLLM_ARGS=(
+    --model "/models/${MODEL_NAME}"
+    --port "$PORT"
+    --host 0.0.0.0
+    --tensor-parallel-size 2
+    --max-model-len "$MAX_MODEL_LEN"
+    --gpu-memory-utilization "$GPU_MEM_UTIL"
+    --dtype auto
+    --trust-remote-code
+    --reasoning-parser qwen3
+    --served-model-name "qwen3.5-122b-a10b-fp8"
+    --max-num-seqs "$MAX_NUM_SEQS"
+    --max-num-batched-tokens "$MAX_BATCHED_TOK"
+    --kv-cache-dtype "$KV_CACHE_DTYPE"
+    --enable-prefix-caching
+    --enable-chunked-prefill
+)
+
+if [ -n "$API_KEY" ]; then
+    VLLM_ARGS+=(--api-key "$API_KEY")
+fi
+
+echo "=============================================="
+echo " vLLM REPLICA — Qwen3.5-122B-A10B-FP8 (TP=2)"
+echo "=============================================="
+echo " GPUs:        ${GPUS}"
+echo " Port:        ${PORT}"
+echo " Context len: ${MAX_MODEL_LEN}"
+echo " GPU util:    ${GPU_MEM_UTIL}"
+echo " Max seqs:    ${MAX_NUM_SEQS}"
+echo " Batched tok: ${MAX_BATCHED_TOK} (chunked prefill)"
+echo " KV dtype:    ${KV_CACHE_DTYPE}"
+echo "=============================================="
+
+# CUDA_VISIBLE_DEVICES pins this replica to its GPU pair. vLLM then
+# sees them as logical devices 0 and 1 for its TP=2.
+apptainer exec --nv \
+    --writable-tmpfs \
+    --env CUDA_VISIBLE_DEVICES="${GPUS}" \
+    --bind "${MODEL_PARENT}:/models" \
+    "$SIF_FILE" \
+    python3 -m vllm.entrypoints.openai.api_server "${VLLM_ARGS[@]}"
--- a/15_start_replicas_122b.sh
+++ b/15_start_replicas_122b.sh
@ -0,0 +1,77 @@
+#!/usr/bin/env bash
+# ------------------------------------------------------------------
+# 15_start_replicas_122b.sh
+# Starts TWO vLLM replicas of Qwen3.5-122B-A10B-FP8 in the background:
+#   replica A: GPUs 0,1 -> internal port 7091
+#   replica B: GPUs 2,3 -> internal port 7092
+# Then starts the load balancer on the public port 7080.
+#
+# Ports 7091/7092 are used (not 7081/7082) because Open WebUI runs on
+# 7081 — the replicas must not collide with it.
+#
+# This replaces the single TP=4 server (11/12_*.sh) for better
+# concurrent throughput on the 4x L40S box. Clients keep using
+# http://<host>:7080/v1 exactly as before.
+#
+# Usage:
+#   bash 15_start_replicas_122b.sh
+#
+# Stop everything:  bash 17_stop_replicas_122b.sh
+# ------------------------------------------------------------------
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+LOG_DIR="${SCRIPT_DIR}/logs"
+mkdir -p "$LOG_DIR"
+
+TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
+
+REPLICA_A_PORT="${REPLICA_A_PORT:-7091}"
+REPLICA_B_PORT="${REPLICA_B_PORT:-7092}"
+PUBLIC_PORT="${PUBLIC_PORT:-7080}"
+export API_KEY="${API_KEY:-}"
+
+# Refuse to start if anything is already bound to the public port.
+if [ -f "${LOG_DIR}/vllm_server.pid" ]; then
+    OLD_PID=$(cat "${LOG_DIR}/vllm_server.pid")
+    if kill -0 "$OLD_PID" 2>/dev/null; then
+        echo "A single-server vLLM (PID ${OLD_PID}) is still running."
+        echo "Stop it first:  bash 05_stop_server.sh"
+        exit 1
+    fi
+fi
+
+start_replica () {
+    local name="$1" gpus="$2" port="$3"
+    local log="${LOG_DIR}/vllm_replica_${name}_${TIMESTAMP}.log"
+    local pidfile="${LOG_DIR}/vllm_replica_${name}.pid"
+
+    if [ -f "$pidfile" ] && kill -0 "$(cat "$pidfile")" 2>/dev/null; then
+        echo "Replica ${name} already running (PID $(cat "$pidfile"))."
+        return
+    fi
+
+    echo "Starting replica ${name}: GPUs ${gpus} -> port ${port}"
+    echo "  log: ${log}"
+    GPUS="$gpus" PORT="$port" \
+        nohup bash "${SCRIPT_DIR}/14_start_replica_122b.sh" > "$log" 2>&1 &
+    echo "$!" > "$pidfile"
+    echo "  PID: $(cat "$pidfile")"
+}
+
+start_replica "a" "0,1" "$REPLICA_A_PORT"
+start_replica "b" "2,3" "$REPLICA_B_PORT"
+
+echo ""
+echo "Both replicas launching (model load takes 5-10 min each, in parallel)."
+echo "Watch:  tail -f ${LOG_DIR}/vllm_replica_*_${TIMESTAMP}.log"
+echo ""
+
+# Start the load balancer on the public port.
+REPLICA_A_PORT="$REPLICA_A_PORT" REPLICA_B_PORT="$REPLICA_B_PORT" \
+PUBLIC_PORT="$PUBLIC_PORT" \
+    bash "${SCRIPT_DIR}/16_start_loadbalancer.sh"
+
+echo ""
+echo "Public API:  http://$(hostname):${PUBLIC_PORT}/v1"
+echo "Stop all:    bash 17_stop_replicas_122b.sh"
--- a/16_start_loadbalancer.sh
+++ b/16_start_loadbalancer.sh
@ -0,0 +1,50 @@
+#!/usr/bin/env bash
+# ------------------------------------------------------------------
+# 16_start_loadbalancer.sh
+# Starts the streaming load balancer (lb_proxy.py) on the public port,
+# fanning requests out to the two 122B replicas by least-in-flight.
+#
+# Runs inside the vLLM container (Python + httpx/uvicorn already there).
+# Normally invoked by 15_start_replicas_122b.sh, but can run alone.
+#
+# Env vars:
+#   PUBLIC_PORT     public port            (default 7080)
+#   REPLICA_A_PORT  replica A port         (default 7091)
+#   REPLICA_B_PORT  replica B port         (default 7092)
+# ------------------------------------------------------------------
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+SIF_FILE="${SCRIPT_DIR}/vllm_qwen.sif"
+LOG_DIR="${SCRIPT_DIR}/logs"
+mkdir -p "$LOG_DIR"
+
+PUBLIC_PORT="${PUBLIC_PORT:-7080}"
+REPLICA_A_PORT="${REPLICA_A_PORT:-7091}"
+REPLICA_B_PORT="${REPLICA_B_PORT:-7092}"
+TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
+LOG_FILE="${LOG_DIR}/vllm_lb_${TIMESTAMP}.log"
+PID_FILE="${LOG_DIR}/vllm_lb.pid"
+
+if [ -f "$PID_FILE" ] && kill -0 "$(cat "$PID_FILE")" 2>/dev/null; then
+    echo "Load balancer already running (PID $(cat "$PID_FILE"))."
+    exit 0
+fi
+
+BACKENDS="http://127.0.0.1:${REPLICA_A_PORT},http://127.0.0.1:${REPLICA_B_PORT}"
+
+echo "Starting load balancer on port ${PUBLIC_PORT} -> ${BACKENDS}"
+echo "  log: ${LOG_FILE}"
+
+# --net not needed: the proxy and replicas share the host network.
+# Install starlette/uvicorn/httpx into a tmpfs if the image lacks them.
+nohup apptainer exec --writable-tmpfs \
+    --env PUBLIC_PORT="$PUBLIC_PORT" \
+    --env BACKENDS="$BACKENDS" \
+    --bind "${SCRIPT_DIR}:${SCRIPT_DIR}" \
+    "$SIF_FILE" \
+    bash -c "python3 -c 'import starlette,uvicorn,httpx' 2>/dev/null || pip install --quiet starlette uvicorn httpx; exec python3 '${SCRIPT_DIR}/lb_proxy.py'" \
+    > "$LOG_FILE" 2>&1 &
+
+echo "$!" > "$PID_FILE"
+echo "  PID: $(cat "$PID_FILE")"
--- a/17_stop_replicas_122b.sh
+++ b/17_stop_replicas_122b.sh
@ -0,0 +1,46 @@
+#!/usr/bin/env bash
+# ------------------------------------------------------------------
+# 17_stop_replicas_122b.sh
+# Stops the load balancer and both 122B replicas started by
+# 15_start_replicas_122b.sh.
+# ------------------------------------------------------------------
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+LOG_DIR="${SCRIPT_DIR}/logs"
+
+stop_one () {
+    local name="$1" pidfile="$2"
+    if [ ! -f "$pidfile" ]; then
+        echo "${name}: no PID file, skipping."
+        return
+    fi
+    local pid
+    pid=$(cat "$pidfile")
+    if kill -0 "$pid" 2>/dev/null; then
+        echo "Stopping ${name} (PID ${pid})..."
+        kill "$pid" 2>/dev/null || true
+        for _ in 1 2 3 4 5; do
+            kill -0 "$pid" 2>/dev/null || break
+            sleep 1
+        done
+        if kill -0 "$pid" 2>/dev/null; then
+            echo "  still alive, SIGKILL..."
+            kill -9 "$pid" 2>/dev/null || true
+        fi
+        echo "  ${name} stopped."
+    else
+        echo "${name} (PID ${pid}) not running."
+    fi
+    rm -f "$pidfile"
+}
+
+# Stop LB first so no new requests get routed to dying replicas.
+stop_one "load balancer" "${LOG_DIR}/vllm_lb.pid"
+stop_one "replica a"     "${LOG_DIR}/vllm_replica_a.pid"
+stop_one "replica b"     "${LOG_DIR}/vllm_replica_b.pid"
+
+echo ""
+echo "All replicas and load balancer stopped."
+echo "Note: vLLM worker subprocesses may take a few seconds to release GPU memory."
+echo "Verify with:  nvidia-smi"
--- a/lb_proxy.py
+++ b/lb_proxy.py
@ -0,0 +1,142 @@
+#!/usr/bin/env python3
+"""
+lb_proxy.py — minimal streaming reverse proxy / load balancer for the
+two vLLM 122B replicas.
+
+Why not nginx: this needs to stream Server-Sent Events (token-by-token
+chat completions) without buffering, route only to replicas that are
+actually up, and have zero install footprint. A ~120-line asyncio proxy
+does exactly that and runs inside the existing vLLM container (which has
+Python + the stdlib; we also use httpx, which ships in the vLLM image).
+
+Routing: least-busy of the healthy replicas. We track in-flight requests
+per backend and send each new request to the one with the fewest. This
+beats blind round-robin when one user fires a huge 128k prompt — the
+other replica keeps serving short requests.
+
+Health: a backend that fails to connect is marked down and skipped;
+we re-probe it on the next request after a short cooldown.
+
+Env vars:
+  PUBLIC_PORT     public port to listen on            (default 7080)
+  BACKENDS        comma-separated backend base URLs
+                  (default http://127.0.0.1:7081,http://127.0.0.1:7082)
+"""
+import asyncio
+import os
+import time
+
+import httpx
+from starlette.applications import Starlette
+from starlette.requests import Request
+from starlette.responses import Response, StreamingResponse, JSONResponse
+from starlette.routing import Route
+import uvicorn
+
+BACKENDS = [
+    b.strip()
+    for b in os.environ.get(
+        "BACKENDS", "http://127.0.0.1:7091,http://127.0.0.1:7092"
+    ).split(",")
+    if b.strip()
+]
+PUBLIC_PORT = int(os.environ.get("PUBLIC_PORT", "7080"))
+COOLDOWN_S = 10.0  # how long to keep a backend marked down before retry
+
+# Per-backend state: in-flight count + "down until" timestamp.
+inflight = {b: 0 for b in BACKENDS}
+down_until = {b: 0.0 for b in BACKENDS}
+
+# Long timeout: 128k-context generations can run for minutes.
+client = httpx.AsyncClient(timeout=httpx.Timeout(None, connect=10.0))
+
+
+def pick_backend() -> str | None:
+    now = time.monotonic()
+    healthy = [b for b in BACKENDS if down_until[b] <= now]
+    if not healthy:
+        # All cooling down — try the soonest-available one anyway.
+        healthy = BACKENDS
+    return min(healthy, key=lambda b: inflight[b])
+
+
+async def proxy(request: Request) -> Response:
+    backend = pick_backend()
+    if backend is None:
+        return JSONResponse({"error": "no backend available"}, status_code=503)
+
+    url = backend + request.url.path
+    if request.url.query:
+        url += "?" + request.url.query
+
+    # Drop hop-by-hop / host headers; pass the rest (incl. Authorization).
+    headers = {
+        k: v
+        for k, v in request.headers.items()
+        if k.lower() not in ("host", "content-length", "connection")
+    }
+    body = await request.body()
+
+    inflight[backend] += 1
+    try:
+        req = client.build_request(
+            request.method, url, headers=headers, content=body
+        )
+        upstream = await client.send(req, stream=True)
+    except (httpx.ConnectError, httpx.ConnectTimeout):
+        down_until[backend] = time.monotonic() + COOLDOWN_S
+        inflight[backend] -= 1
+        # Retry once on the other backend.
+        other = pick_backend()
+        if other and other != backend:
+            return await proxy(request)
+        return JSONResponse({"error": "backend unreachable"}, status_code=502)
+
+    resp_headers = {
+        k: v
+        for k, v in upstream.headers.items()
+        if k.lower() not in ("content-length", "transfer-encoding", "connection")
+    }
+
+    async def stream():
+        try:
+            async for chunk in upstream.aiter_raw():
+                yield chunk
+        finally:
+            await upstream.aclose()
+            inflight[backend] -= 1
+
+    return StreamingResponse(
+        stream(),
+        status_code=upstream.status_code,
+        headers=resp_headers,
+        media_type=upstream.headers.get("content-type"),
+    )
+
+
+async def health(request: Request) -> Response:
+    now = time.monotonic()
+    return JSONResponse(
+        {
+            "backends": [
+                {
+                    "url": b,
+                    "inflight": inflight[b],
+                    "up": down_until[b] <= now,
+                }
+                for b in BACKENDS
+            ]
+        }
+    )
+
+
+app = Starlette(
+    routes=[
+        Route("/_lb_health", health, methods=["GET"]),
+        Route("/{path:path}", proxy, methods=["GET", "POST", "PUT", "DELETE"]),
+    ]
+)
+
+if __name__ == "__main__":
+    print(f"Load balancer on :{PUBLIC_PORT} -> {BACKENDS}", flush=True)
+    uvicorn.run(app, host="0.0.0.0", port=PUBLIC_PORT, log_level="warning")