unified_v2: selective per-request PD-sep via Mooncake (E3+E4)
Adds a sixth routing policy --policy unified_v2 that wraps the
existing unified hybrid picker with a selective PD-sep branch.
When all of the following hold, a request is split prefill-on-src,
decode-on-chosen via Mooncake kv_role=kv_both transfer:
1. new_local = input_length - chosen.cache_hit > 16k
(B2 microbench shows same-worker TTFT idx >= 3x from this size up)
2. chosen has live decodes worth protecting (>= 2 in-flight)
3. some other instance holds materially more cache for this prefix
(>= 8k tokens, and >= 4k more than chosen)
4. cost(src_interference + RDMA xfer) + 0.2s margin < cost(chosen_interference)
The cost model is the audit-blessed shape from E1's post-mortem:
- gate on new_tokens (post-cache), NOT input_length (the old PUSH gate)
- bind to a single transfer mechanism (kv_both peer-to-peer pull)
- realistic RDMA cost as a function of bytes: 0.3s base +
bytes / 2.7 GB/s (calibrated against contention_16s_elastic p50)
- both source and target decode counts considered
E2 mechanism-level patches not yet applied (this commit is policy-only).
Patches 6.2 / 6.3 / 6.5 remain on the table. Patch 6.6 (per-request
xfer timeout, 60s default) is implemented on the proxy side as an
httpx per-chunk read timeout on the dst streaming call, so a stuck
KV transfer fails the request instead of hanging for 600s.
cache_aware_proxy.py:
- Settings: kv_bytes_per_token, prefill_throughput_kv_both,
rdma_base_overhead_s, rdma_effective_gb_per_s, pd_sep_* gating knobs
- estimate_transfer_cost(bytes) replaces the constant rdma_overhead_s
- estimate_same_worker_interference_s(new_tokens, num_decodes) reads off
the B2 penalty curve in 4 bins
- pick_instance_unified_v2: inherits unified, returns extra
(src_inst, src_idx) tuple when PD-sep wins the cost compare
- _handle_combined_pd_sep_v2: prefill on src (do_remote_decode=True,
max_tokens=1), Mooncake xfer, decode-stream on dst with httpx
Timeout(read=pd_sep_xfer_timeout_s)
- --policy unified_v2 added to argparse choices
- lifespan auto-runs init_prefill_bootstrap when policy is unified_v2
b3_isolated_policy.sh:
- ENABLE_KV_BOTH env var, auto-set when POLICY=unified_v2, threads
kv_role=kv_both + VLLM_MOONCAKE_BOOTSTRAP_PORT to vllm and
--bootstrap-ports to the proxy
Tests: 8 new unit tests cover the gating predicates and the cost
estimators; all 32 proxy tests still pass.
Refs: E1 (PUSH post-mortem) + E2 (Mooncake audit) reports.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -19,11 +19,22 @@ BASE_PORT="${BASE_PORT:-8000}"
|
||||
GPU_INDICES="${GPU_INDICES:-0 1 2 3 4 5 6 7}"
|
||||
EXTRA_VLLM_ARGS="${EXTRA_VLLM_ARGS:---enable-prompt-tokens-details}"
|
||||
N_INSTANCES=$(echo $GPU_INDICES | wc -w)
|
||||
# When ENABLE_KV_BOTH=1, vLLM launches with the Mooncake KV connector in
|
||||
# kv_both role and the proxy is given bootstrap ports. This is required
|
||||
# for --policy unified_v2 (per-request PD-sep) but disabled by default
|
||||
# because it adds always-on KV-transfer overhead even when not triggered.
|
||||
ENABLE_KV_BOTH="${ENABLE_KV_BOTH:-0}"
|
||||
BOOTSTRAP_BASE_PORT="${BOOTSTRAP_BASE_PORT:-8998}"
|
||||
|
||||
POLICY="${1:?usage: $0 <policy> <trace> <rundir>}"
|
||||
TRACE="${2:?usage: $0 <policy> <trace> <rundir>}"
|
||||
RUNDIR="${3:?usage: $0 <policy> <trace> <rundir>}"
|
||||
|
||||
# Auto-enable kv_both when the policy requires it.
|
||||
if [ "$POLICY" = "unified_v2" ]; then
|
||||
ENABLE_KV_BOTH=1
|
||||
fi
|
||||
|
||||
mkdir -p "$RUNDIR/engine_state" "$RUNDIR/logs"
|
||||
echo "[isolated] policy=$POLICY trace=$(basename $TRACE) rundir=$RUNDIR"
|
||||
|
||||
@@ -38,23 +49,46 @@ trap cleanup EXIT
|
||||
# Hard reset first
|
||||
cleanup
|
||||
|
||||
echo "[isolated] launching $N_INSTANCES vLLM on GPUs $GPU_INDICES ..."
|
||||
echo "[isolated] launching $N_INSTANCES vLLM on GPUs $GPU_INDICES ENABLE_KV_BOTH=$ENABLE_KV_BOTH ..."
|
||||
i=0
|
||||
kv_both_extra=""
|
||||
if [ "$ENABLE_KV_BOTH" = "1" ]; then
|
||||
kv_both_extra="--kv-transfer-config {\"kv_connector\":\"MooncakeConnector\",\"kv_role\":\"kv_both\"}"
|
||||
fi
|
||||
for gpu in $GPU_INDICES; do
|
||||
port=$((BASE_PORT + i))
|
||||
master=$((29500 + i))
|
||||
AGENTIC_STEP_LOG_PATH="$RUNDIR/engine_state/engine_${i}.jsonl" \
|
||||
AGENTIC_WORKER_ID="engine_${i}" \
|
||||
CUDA_VISIBLE_DEVICES=$gpu \
|
||||
MASTER_PORT=$master \
|
||||
nohup "$VENV/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 \
|
||||
$EXTRA_VLLM_ARGS \
|
||||
> "$RUNDIR/logs/vllm_inst_${i}_gpu${gpu}.log" 2>&1 &
|
||||
bp=$((BOOTSTRAP_BASE_PORT + i))
|
||||
if [ "$ENABLE_KV_BOTH" = "1" ]; then
|
||||
PYTHONHASHSEED=42 \
|
||||
VLLM_MOONCAKE_BOOTSTRAP_PORT=$bp \
|
||||
AGENTIC_STEP_LOG_PATH="$RUNDIR/engine_state/engine_${i}.jsonl" \
|
||||
AGENTIC_WORKER_ID="engine_${i}" \
|
||||
CUDA_VISIBLE_DEVICES=$gpu \
|
||||
MASTER_PORT=$master \
|
||||
nohup "$VENV/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 \
|
||||
--kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_both"}' \
|
||||
$EXTRA_VLLM_ARGS \
|
||||
> "$RUNDIR/logs/vllm_inst_${i}_gpu${gpu}.log" 2>&1 &
|
||||
else
|
||||
AGENTIC_STEP_LOG_PATH="$RUNDIR/engine_state/engine_${i}.jsonl" \
|
||||
AGENTIC_WORKER_ID="engine_${i}" \
|
||||
CUDA_VISIBLE_DEVICES=$gpu \
|
||||
MASTER_PORT=$master \
|
||||
nohup "$VENV/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 \
|
||||
$EXTRA_VLLM_ARGS \
|
||||
> "$RUNDIR/logs/vllm_inst_${i}_gpu${gpu}.log" 2>&1 &
|
||||
fi
|
||||
disown
|
||||
sleep 2
|
||||
i=$((i + 1))
|
||||
@@ -80,10 +114,23 @@ combined_args=""
|
||||
for i in $(seq 0 $((N_INSTANCES - 1))); do
|
||||
combined_args="$combined_args http://127.0.0.1:$((BASE_PORT + i))"
|
||||
done
|
||||
proxy_extra=""
|
||||
if [ "$ENABLE_KV_BOTH" = "1" ]; then
|
||||
bp_list=""
|
||||
for i in $(seq 0 $((N_INSTANCES - 1))); do
|
||||
if [ -z "$bp_list" ]; then
|
||||
bp_list="$((BOOTSTRAP_BASE_PORT + i))"
|
||||
else
|
||||
bp_list="$bp_list,$((BOOTSTRAP_BASE_PORT + i))"
|
||||
fi
|
||||
done
|
||||
proxy_extra="--bootstrap-ports $bp_list"
|
||||
fi
|
||||
nohup "$VENV/python" "$ROOT/scripts/cache_aware_proxy.py" \
|
||||
--port "$PROXY_PORT" \
|
||||
--combined $combined_args \
|
||||
--policy "$POLICY" \
|
||||
$proxy_extra \
|
||||
> "$RUNDIR/proxy.log" 2>&1 &
|
||||
disown
|
||||
tries=0
|
||||
|
||||
Reference in New Issue
Block a user