From f792537719416af582bf4d6c880a6799c400f306 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Sat, 16 May 2026 06:11:38 +0000
Subject: [PATCH] feat: auto-warmup in build_and_run.sh
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

After the container starts, the script waits for the API to come up,
then sends a warmup request to trigger all JIT compilation (Triton,
TileLang, CuTeDSL). This way the first real inference request is fast.

Also added tqdm for expert weight loading:
  Loading Native NVFP4 Expert Weights: 50%|██████████░░| 480/960
---
 build_and_run.sh | 20 ++++++++++++++++++++
 warmup.sh        | 27 +++++++++++++++++++++++++++
 2 files changed, 47 insertions(+)
 create mode 100755 warmup.sh

diff --git a/build_and_run.sh b/build_and_run.sh
index 4431cf04..41ca2baf 100755
--- a/build_and_run.sh
+++ b/build_and_run.sh
@@ -20,3 +20,23 @@ docker compose up -d
 sed -i -E "s/ARG CACHE_BUSTER=.*/ARG CACHE_BUSTER=\${TIMESTAMP}/" Dockerfile
 
 echo "=== Done. Container: $(docker compose ps -q) ==="
+echo ""
+echo "=== Waiting for API to come up, then sending warmup request... ==="
+echo "=== (This triggers Triton/TileLang/CuTeDSL JIT compilation) ==="
+
+for i in $(seq 1 120); do
+    if curl -s http://localhost:8000/health > /dev/null 2>&1; then
+        echo "=== API is up. Sending warmup request (expect 2-3 min)... ==="
+        START=$(date +%s)
+        curl -s -X POST http://localhost:8000/v1/completions \
+          -H "Content-Type: application/json" \
+          -d '{"model": "/model", "prompt": "The capital of France is", "max_tokens": 1}' \
+          2>&1 | head -5
+        END=$(date +%s)
+        ELAPSED=$((END - START))
+        echo ""
+        echo "=== Warmup complete (${ELAPSED}s). Model is ready for inference. ==="
+        break
+    fi
+    sleep 5
+done
diff --git a/warmup.sh b/warmup.sh
new file mode 100755
index 00000000..26e78ad7
--- /dev/null
+++ b/warmup.sh
@@ -0,0 +1,27 @@
+#!/bin/bash
+# Wait for vLLM API to be ready, then send a warmup request
+# to compile all JIT kernels (Triton, TileLang, CuTeDSL)
+# before the first real inference request.
+
+echo "=== Warmup: waiting for vLLM API... ==="
+
+# Wait for the API to respond
+for i in $(seq 1 120); do
+    if curl -s http://localhost:8000/health > /dev/null 2>&1; then
+        echo "=== Warmup: API is up, sending warmup request... ==="
+        break
+    fi
+    sleep 5
+done
+
+# Send a warmup request to trigger all JIT compilation
+# This will be slow (2-3 min) but subsequent requests will be fast
+START=$(date +%s)
+curl -s -X POST http://localhost:8000/v1/completions \
+  -H "Content-Type: application/json" \
+  -d '{"model": "/model", "prompt": "The capital of France is", "max_tokens": 1}' \
+  > /dev/null 2>&1
+END=$(date +%s)
+ELAPSED=$((END - START))
+
+echo "=== Warmup complete (${ELAPSED}s). Model is ready. ==="