Layerwise KV transfer on Mooncake: PoC + microbench (worktree exploration)

Implements per-layer KV push during prefill (write mode) on vLLM's
MooncakeConnector, env-gated by MOONCAKE_LAYERWISE=1. 2-instance microbench
(mb7) shows correctness (KV lands, cached==prompt) and that the transfer is
hidden behind prefill compute: critical-path overhead drops from O(KV size)
(123/202/529ms for 8k/16k/32k) to a flat ~58ms (2-9x), with no prefill
slowdown, on idle instances. Caveats: idle-only, chunked-prefill disabled,
single concurrent transfer — see DESIGN.md.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

microbench/connector_tax/layerwise/run_mb7.sh

@@ -0,0 +1,111 @@
+#!/usr/bin/env bash
+# MB7 launcher (runs on dash0). Two 2-instance modes selected by MODE env:
+#   MODE=baseline  : restore stock connector, no layerwise env
+#   MODE=layerwise : deploy mooncake_connector.LAYERWISE.py + MOONCAKE_LAYERWISE=1
+#
+# Chunked prefill is DISABLED (max-num-batched-tokens >= max prompt) so the
+# producer prefill is a single forward and save_kv_layer fires once per layer
+# in order — the layer-wise counter assumes this.
+#
+# The connector is always restored from .ORIG_BACKUP on exit.
+#
+# Usage (on dash0):
+#   MODE=baseline  bash run_mb7.sh
+#   MODE=layerwise bash run_mb7.sh
+
+set -uo pipefail
+
+MODE="${MODE:-baseline}"
+PROJ_DIR="${PROJ_DIR:-/home/admin/cpfs/wjh/agentic-kv}"
+VENV="${VENV:-$PROJ_DIR/.venv}"
+MODEL="${MODEL:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-Coder-30B-A3B-Instruct}"
+GPUS=(${GPUS:-0 1})
+SIZES="${SIZES:-8192,16384,32768}"
+REPEATS="${REPEATS:-3}"
+MAX_BATCHED="${MAX_BATCHED:-40960}"   # >= max prompt => no chunked prefill
+DATE="$(date +%Y%m%d_%H%M)"
+OUTDIR="${OUTDIR:-$PROJ_DIR/outputs/mb7_${MODE}_${DATE}}"
+PYTHON="$VENV/bin/python"
+MC_FILE="$VENV/lib/python3.12/site-packages/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py"
+LW_SRC="${LW_SRC:-/tmp/mooncake_connector.LAYERWISE.py}"
+DRIVER="$PROJ_DIR/microbench/connector_tax/layerwise/mb7_layerwise.py"
+
+mkdir -p "$OUTDIR/logs"
+PORTS=(8000 8001); BPS=(8998 8999)
+
+echo "=== MB7 ($MODE) ==="
+echo "Out: $OUTDIR ; connector: $MC_FILE"
+
+restore_connector() {
+    if [ -f "$MC_FILE.ORIG_BACKUP" ]; then
+        cp -f "$MC_FILE.ORIG_BACKUP" "$MC_FILE"
+        echo "[restore] connector reset to ORIG"
+    fi
+}
+cleanup() {
+    pkill -9 -f "vllm serve" 2>/dev/null || true
+    pkill -9 -f "EngineCore" 2>/dev/null || true
+    sleep 4
+    restore_connector
+}
+trap cleanup EXIT
+pkill -9 -f "vllm serve" 2>/dev/null || true; sleep 3
+
+# Deploy the connector for the chosen mode.
+if [ "$MODE" = "layerwise" ]; then
+    if [ ! -f "$LW_SRC" ]; then echo "FATAL: $LW_SRC not found (scp it first)"; exit 1; fi
+    cp -f "$LW_SRC" "$MC_FILE"
+    "$PYTHON" -c "import ast; ast.parse(open('$MC_FILE').read()); print('[deploy] LAYERWISE connector AST OK')" || exit 1
+    LW_ENV="MOONCAKE_LAYERWISE=1"
+else
+    restore_connector
+    LW_ENV=""
+fi
+
+echo "[launch] 2 instances (max-num-batched-tokens=$MAX_BATCHED, chunked-prefill off)"
+i=0
+for gpu in "${GPUS[@]:0:2}"; do
+    port=${PORTS[$i]}; bp=${BPS[$i]}; master=$((29700 + i))
+    env $LW_ENV \
+        PYTHONHASHSEED=42 VLLM_MOONCAKE_BOOTSTRAP_PORT=$bp \
+        CUDA_VISIBLE_DEVICES=$gpu MASTER_PORT=$master \
+        nohup "$VENV/bin/vllm" serve "$MODEL" \
+        --host 0.0.0.0 --port "$port" --tensor-parallel-size 1 \
+        --trust-remote-code --enable-prefix-caching --dtype auto \
+        --gpu-memory-utilization 0.9 --max-model-len 200000 \
+        --max-num-batched-tokens "$MAX_BATCHED" \
+        --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_both"}' \
+        --enable-prompt-tokens-details \
+        > "$OUTDIR/logs/vllm_${i}_gpu${gpu}.log" 2>&1 &
+    disown; sleep 2; i=$((i + 1))
+done
+
+echo "[health] waiting ..."
+for i in 0 1; do
+    port=${PORTS[$i]}; tries=0
+    while ! curl -sf "http://127.0.0.1:$port/health" >/dev/null 2>&1; do
+        tries=$((tries + 1)); [ $tries -gt 180 ] && { echo "FATAL inst_$i"; exit 1; }
+        sleep 2
+    done
+    echo "  inst_$i ready"
+done
+for i in 0 1; do
+    bp=${BPS[$i]}; tries=0
+    while ! curl -sf "http://127.0.0.1:$bp/query" >/dev/null 2>&1; do
+        tries=$((tries+1)); [ $tries -gt 60 ] && { echo "WARN bp $bp"; break; }; sleep 2
+    done
+done
+
+echo "[run] mb7 --mode $MODE"
+"$PYTHON" "$DRIVER" --mode "$MODE" \
+    --src-port "${PORTS[0]}" --dst-port "${PORTS[1]}" \
+    --src-bp "${BPS[0]}" --dst-bp "${BPS[1]}" \
+    --sizes "$SIZES" --repeats "$REPEATS" --out "$OUTDIR/mb7_result.json" \
+    2>&1 | tee "$OUTDIR/mb7_run.txt"
+
+echo "[done] $OUTDIR"
+# grep layerwise transfer logs from the producer (gpu0) for sanity
+if [ "$MODE" = "layerwise" ]; then
+    echo "=== producer layerwise log lines ==="
+    grep -i "layerwise" "$OUTDIR/logs/vllm_0_gpu${GPUS[0]}.log" | tail -10 || true
+fi