Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
306 changes: 221 additions & 85 deletions benchmarks/single_node/agentic/dsv4_fp4_b300_vllm.sh
Original file line number Diff line number Diff line change
@@ -1,60 +1,47 @@
#!/usr/bin/env bash
set -euo pipefail
set -eo pipefail

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Line 2 weakens set -euo pipefail to set -eo pipefail, dropping nounset — this is the only script among 23 in benchmarks/single_node/agentic/ that does so (its B200 sibling dsv4_fp4_b200_vllm.sh from #2231 handles the exact same optionally-unset vars with ${VAR:-} instead). Recommend restoring set -u and using ${KV_OFFLOAD_BACKEND:-}, ${SLURM_JOB_ID:-}, and ${MODEL_PATH:-} in the new unguarded expansions instead of disabling nounset for the whole ~350-line script.

Extended reasoning...

What changed: Line 2 of dsv4_fp4_b300_vllm.sh changes set -euo pipefail to set -eo pipefail, dropping the -u (nounset) flag. I verified this is the only script among all 23 scripts in benchmarks/single_node/agentic/ that lacks nounset — every other script, including the direct B200 counterpart dsv4_fp4_b200_vllm.sh introduced in #2231 (which this PR explicitly says it mirrors), keeps set -euo pipefail.

Why it was dropped: The rewrite replaces the old declare -p VAR >/dev/null 2>&1 && [ -n "$VAR" ] guards for MODEL_PATH/SLURM_JOB_ID with plain [[ -n "$VAR" ]], and adds a new case "$KV_OFFLOAD_BACKEND" in "") branch that relies on KV_OFFLOAD_BACKEND expanding to empty when unset. None of KV_OFFLOAD_BACKEND, SLURM_JOB_ID, or MODEL_PATH are in the check_env_vars allowlist, so under nounset these bare references would abort with "unbound variable" whenever the var isn't exported (e.g. a stand-alone run without SLURM, or the vllm-simple/mooncake KV_OFFLOAD_BACKEND path). Dropping -u script-wide is the blunt fix for that.

Why this is unnecessary: The B200 sibling script that this PR is explicitly mirroring solves the identical problem without sacrificing nounset, using the standard ${VAR:-} parameter-expansion idiom: [[ -n "${SLURM_JOB_ID:-}" ]], ${SLURMD_NODENAME:-unknown}, [[ -n "${MODEL_PATH:-}" ]]. benchmark_lib.sh itself uses this idiom throughout. So the intent (let three specific optional vars expand to empty) can be fully achieved without disabling the safety net for the other ~350 lines of the script — the fix conflated "these three vars are allowed to be unset" with "disable unbound-variable checking everywhere."

Impact: With nounset off, any future typo in a variable reference anywhere in this script (e.g. a misspelled $LMCACHE_PORT, $TOTAL_CPU_DRAM_GB, or $GPU_COUNT) will silently expand to an empty string instead of aborting immediately with "unbound variable." Concretely: GPU_COUNT=$TP further down is guarded by a regex check so a typo there would still be caught, but a hypothetical typo inside one of the arithmetic expressions (e.g. LMCACHE_L1_SIZE_GB=$((TOTAL_CPU_DRAM_GB * 3 / 4)) misspelled as TOTAL_CPU_DRM_GB) would silently evaluate to 0 rather than failing fast, producing a confusing downstream failure (or a nonsensical LMCache pool size) instead of an immediate, obvious error at the point of the typo.

Step-by-step proof of the regression mechanism (not a live failure today, since check_env_vars still validates required vars and the arithmetic guards catch $GPU_COUNT):

  1. Suppose a future edit misspells $TOTAL_CPU_DRAM_GB as $TOTAL_CPU_DRAM_G inside the lmcache branch's LMCACHE_L1_SIZE_GB=$((TOTAL_CPU_DRAM_GB * 3 / 4)) line.
  2. Under set -u (the previous/sibling-script behavior), bash would immediately abort with TOTAL_CPU_DRAM_G: unbound variable, pointing directly at the typo.
  3. Under this PR's set -eo pipefail (nounset disabled), the misspelled var silently expands to empty string, the arithmetic becomes $(( * 3 / 4 )) → bash arithmetic treats an empty operand as 0, so LMCACHE_L1_SIZE_GB=0.
  4. The script proceeds to start lmcache server --l1-size-gb 0 ..., which either fails deep inside the LMCache server startup logs (unclear connection to the root cause) or silently runs a degenerate 0-size cache — either way, a much harder failure to root-cause than an immediate bash error at the typo site.

Fix: Restore set -euo pipefail on line 2, and change the three unguarded references to use parameter expansion consistent with the B200 sibling: [[ -n "${SLURM_JOB_ID:-}" ]], ${SLURMD_NODENAME:-unknown}, [[ -n "${MODEL_PATH:-}" ]], and case "${KV_OFFLOAD_BACKEND:-}" in "").

On severity: All four verifiers independently confirmed the factual claim (only script in the directory without nounset, unnecessary given the ${VAR:-} idiom is used successfully in the sibling script) but converged on nit rather than normal, and I agree. check_env_vars still validates every genuinely required variable, GPU_COUNT is separately regex-validated, and the script runs correctly today as written on every current sweep point — the harm is a hypothetical future maintenance hazard, not a concrete failure in this PR. That's a real, worth-fixing consistency/robustness regression, but not a merge blocker.

set -x

# Agentic trace replay benchmark for DeepSeek-V4-Pro FP4 on B300 using vLLM.
# Mirrors the fixed-seq-len parallelism options (pure TP and DEP) so the
# agentic sweep can probe both interactivity and throughput regimes:
# pure TP (DP_ATTENTION=false, EP_SIZE=1): attention TP-sharded across
# all $TP GPUs in a single engine. Lower TPOT, lower batch.
# TP+EP (DP_ATTENTION=false, EP_SIZE>1): attention TP-sharded, MoE
# experts EP-sharded within the TP group.
# DEP (DP_ATTENTION=true, EP_SIZE>1): per-DP-rank attention with
# experts EP-sharded across DP ranks (per the vLLM blog recipe).
# Highest aggregate throughput at large CONC.
# v4pro-b300.yaml TP4 and DEP4 recipe. SimpleCPUOffload / MooncakeStore /
# LMCache
#
# Image is vllm/vllm-openai:v0.20.0-cu130. block_size=256, kv-cache-dtype=fp8,
# FP4 indexer cache enabled, FULL_AND_PIECEWISE cudagraph capture with
# custom_ops=all (per the vLLM blog recipe at https://vllm.ai/blog/deepseek-v4).
# Image is configured in nvidia-master.yaml. The recipe uses FP8 KV cache,
# sparse DeepSeek-V4 FlashInfer attention with an FP4 indexer cache, mega-MoE,
# and FULL_DECODE_ONLY CUDA graphs with every batch size captured explicitly.
#
# Required env vars:
# MODEL, TP, CONC, KV_OFFLOADING, TOTAL_CPU_DRAM_GB, RESULT_DIR
#
# KV_OFFLOADING=dram requires KV_OFFLOAD_BACKEND=mooncake.
# KV_OFFLOADING=none is used by TP4. DEP4 uses KV_OFFLOADING=dram with
# KV_OFFLOAD_BACKEND=vllm-simple, mooncake, or lmcache.

source "$(dirname "$0")/../../benchmark_lib.sh"

check_env_vars MODEL TP CONC KV_OFFLOADING TOTAL_CPU_DRAM_GB RESULT_DIR DURATION EP_SIZE DP_ATTENTION

DCP_SIZE="${DCP_SIZE:-1}"
PCP_SIZE="${PCP_SIZE:-1}"
VLLM_CP_ARGS=()
if [ "$DCP_SIZE" -gt 1 ]; then
VLLM_CP_ARGS+=(--decode-context-parallel-size "$DCP_SIZE")
fi
if [ "$PCP_SIZE" -gt 1 ]; then
VLLM_CP_ARGS+=(--prefill-context-parallel-size "$PCP_SIZE")
fi

GPU_COUNT="${GPU_COUNT:-$((TP * PCP_SIZE))}"
GPU_COUNT=$TP
if [[ ! "$GPU_COUNT" =~ ^[1-9][0-9]*$ ]]; then
echo "Error: GPU_COUNT must be a positive integer, got '$GPU_COUNT'" >&2
exit 1
fi
export GPU_COUNT

if declare -p SLURM_JOB_ID >/dev/null 2>&1 && [ -n "$SLURM_JOB_ID" ]; then
SLURM_NODE=unknown
if declare -p SLURMD_NODENAME >/dev/null 2>&1 && [ -n "$SLURMD_NODENAME" ]; then
SLURM_NODE="$SLURMD_NODENAME"
fi
echo "JOB $SLURM_JOB_ID running on $SLURM_NODE"
# Under DP-attention the DP world size equals TP, and the DEP recipe sizes
# per-rank batch as MAX_NUM_SEQS = 2*CONC/TP, which must be an integer.
if [ "$DP_ATTENTION" = "true" ] && [ $((2 * CONC % TP)) -ne 0 ]; then
echo "Error: DEP requires 2*CONC divisible by TP, got CONC='$CONC' and TP='$TP'" >&2
exit 1
fi

if [[ -n "$SLURM_JOB_ID" ]]; then
echo "JOB $SLURM_JOB_ID running on $SLURMD_NODENAME"
fi

# `hf download` creates the target dir if missing and is itself idempotent.
# When MODEL_PATH is unset (stand-alone runs), fall back to the HF_HUB_CACHE
# When MODEL_PATH is unset (stand-alone runs), fall back to the HF_HUB_CACHE.
# Either way, MODEL_PATH is what the server is launched with.
if declare -p MODEL_PATH >/dev/null 2>&1 && [ -n "$MODEL_PATH" ]; then
if [[ -n "$MODEL_PATH" ]]; then
if [[ ! -d "$MODEL_PATH" || -z "$(ls -A "$MODEL_PATH" 2>/dev/null)" ]]; then
hf download "$MODEL" --local-dir "$MODEL_PATH"
fi
Expand All @@ -68,17 +55,9 @@ nvidia-smi
resolve_trace_source
install_agentic_deps

# vLLM v0.22.1 can ship CUTLASS DSL 4.5.2 with stale native MLIR bindings,
# which fails DSV4 indexer compilation with mlir_global_dtors(..., data).
# Reinstall the matching native wheel until NVIDIA/cutlass#3259 is resolved.
agentic_pip_install --quiet --force-reinstall --no-deps \
'nvidia-cutlass-dsl-libs-cu13==4.5.2'

# vllm-project/router expands the one HTTP backend into one logical worker per
# DP rank and sends X-data-parallel-rank on forwarded requests. aiperf's
# X-Correlation-ID is stable for every turn of a conversation; alias it to the
# router's preferred X-Session-ID header. This also keeps affinity correct when
# testing older wheels that prioritize per-request X-Request-ID.
# DP rank. Bind every turn of a conversation to the same rank by mapping
# AIPerf's stable correlation ID to the router's X-Session-ID header.
USE_VLLM_ROUTER=false
VLLM_BACKEND_PORT="$PORT"
if [ "$DP_ATTENTION" = "true" ]; then
Expand All @@ -91,31 +70,60 @@ if [ "$DP_ATTENTION" = "true" ]; then
agentic_pip_install --quiet "vllm-router==$VLLM_ROUTER_VERSION"
fi

# DeepSeek-V4-Pro weights are large; engine startup can exceed default 600s.
# Match the environment used by v4pro-b300.yaml.
export VLLM_USE_V2_MODEL_RUNNER=1
export VLLM_ENGINE_READY_TIMEOUT_S=3600

# vllm-project/vllm#43447 keeps local SWA prefix-cache tails sparsely, while
# vllm-project/vllm#44774 applies the same reachability policy to Mooncake's
# store mask. 32k matches the trace-replay tuning validated for this workload.
export VLLM_PREFIX_CACHE_RETENTION_INTERVAL=32768
export VLLM_DSV4_MEGA_FP8_COMBINE=1
export NCCL_NVLS_ENABLE=1
export VLLM_USE_RUST_FRONTEND=1

# ---- Server config ----------------------------------------------------------
SERVER_LOG="$RESULT_DIR/server.log"
ROUTER_LOG="$RESULT_DIR/router.log"
MOONCAKE_MASTER_LOG="$RESULT_DIR/mooncake_master.log"
LMCACHE_SERVER_LOG="$RESULT_DIR/lmcache_server.log"
mkdir -p "$RESULT_DIR"

SERVER_PID=""
ROUTER_PID=""
MOONCAKE_MASTER_PID=""
LMCACHE_SERVER_PID=""

# The generated TOTAL_CPU_DRAM_GB budget is proportional to allocated GPUs.
# On cluster:b300-nv, dram-utilization=0.80 and DEP4 resolve to roughly the
# source recipe's 280 GiB per DP rank. TP4 remains GPU-resident.
OFFLOAD_ARGS=()
if require_agentic_kv_offload_backend mooncake; then
# Mooncake embedded mode contributes one global segment per GPU rank to
# a shared distributed store. Pre-divide the aggregate host budget
# across those rank-contributed segments.
case "$KV_OFFLOAD_BACKEND" in
"")
require_agentic_kv_offload_none
;;
vllm-simple)
require_agentic_kv_offload_backend vllm-simple
CPU_BYTES_PER_RANK=$(( TOTAL_CPU_DRAM_GB * 1000 * 1000 * 1000 / GPU_COUNT ))
# Identical prefixes must hash to identical block keys across DP ranks.
export PYTHONHASHSEED=42
OFFLOAD_CONFIG=$(cat <<EOF
{
"kv_connector": "SimpleCPUOffloadConnector",
"kv_role": "kv_both",
"kv_connector_extra_config": {
"cpu_bytes_to_use": ${CPU_BYTES_PER_RANK},
"enable_cross_layers_blocks": "true"
}
}
EOF
)
OFFLOAD_ARGS=(
--kv-transfer-config
"$OFFLOAD_CONFIG"
)
;;
mooncake)
require_agentic_kv_offload_backend mooncake
# Embedded mode contributes one global segment per DP rank to the
# shared store, so divide the aggregate host budget across ranks.
PER_RANK_GB=$((TOTAL_CPU_DRAM_GB / GPU_COUNT))

MOONCAKE_VERSION=0.3.11.post1
agentic_pip_install --quiet --no-cache-dir --no-deps \
--force-reinstall "mooncake-transfer-engine-cuda13==$MOONCAKE_VERSION"
Expand All @@ -139,9 +147,7 @@ EOF
export MC_ENABLE_DEST_DEVICE_AFFINITY=1
# Identical prefixes must hash to identical store keys across DP ranks.
export PYTHONHASHSEED=0
# Large agentic KV writes can exceed Mooncake Store's fixed 60-second
# transfer deadline at the default 64 KiB RDMA slice size. Reduce
# per-transfer bookkeeping and give the shared RNIC more workers.
export WITH_NVIDIA_PEERMEM=0
export MC_SLICE_SIZE=1048576
export MC_WORKERS_PER_CTX=4

Expand All @@ -164,55 +170,185 @@ EOF
exit 1
fi

unset VLLM_USE_SIMPLE_KV_OFFLOAD
OFFLOAD_CONFIG='{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_connector_extra_config":{"load_async":true}}'
OFFLOAD_ARGS=(--kv-transfer-config "$OFFLOAD_CONFIG")
;;
lmcache)
require_agentic_kv_offload_backend lmcache
# The LMCache MP server owns the host-DRAM KV pool as one shared
# tier; vLLM ranks attach via LMCacheMPConnector, so the aggregate
# host budget is passed through undivided (unlike Mooncake's
# per-rank segments). Follows the LMCache DeepSeek-V4 recipe
# (docs.lmcache.ai/recipes/deepseek_v4_flash.html); LMCache handles
# DSV4's Sparse-MLA hybrid KV geometries automatically.
LMCACHE_VERSION=0.5.1
agentic_pip_install --quiet --no-cache-dir "lmcache==$LMCACHE_VERSION"
python3 -c "import lmcache.integration.vllm.lmcache_mp_connector" >/dev/null

LMCACHE_HOST=127.0.0.1
LMCACHE_PORT=$((PORT + 12000))
LMCACHE_HTTP_PORT=$((PORT + 13000))
# LMCacheMPConnector concatenates lmcache.mp.host and port into the
# ZMQ endpoint. Bind the server to a raw host, but pass the connector
# a ZMQ-style host string.
LMCACHE_CONNECT_HOST="tcp://$LMCACHE_HOST"
# Pool target derated to 75% of the aggregate budget: pinned host
# memory is unswappable and also consumes GPU-side mapping
# resources, so leave headroom for vLLM host buffers and the OS.
# Full-budget targets OOM-killed the node (host OOM-killer or
# cudaErrorMemoryAllocation) as the cache filled past ~2 TB during
# PR #2153 bring-up.
LMCACHE_L1_SIZE_GB=$((TOTAL_CPU_DRAM_GB * 3 / 4))
# The pool grows lazily from the initial allocation, so the full
# --l1-size-gb target is not pinned at startup.
LMCACHE_L1_INIT_SIZE_GB=20
LMCACHE_MQ_TIMEOUT=300
# Identical prefixes must hash to identical cache keys across DP ranks.
export PYTHONHASHSEED=0
# Per-engine scheduler stats every 5s, to diagnose per-DP-rank KV
# cache imbalance under the session-sticky router.
export VLLM_LOG_STATS_INTERVAL=5

echo "Starting LMCache MP server on port $LMCACHE_PORT..."
# One GPU-side transfer worker avoids concurrent-GPU-transfer stalls
# under heavy async-load pressure; CPU-side workers stay at 8.
lmcache server \
--host "$LMCACHE_HOST" \
--port "$LMCACHE_PORT" \
--http-host "$LMCACHE_HOST" \
--http-port "$LMCACHE_HTTP_PORT" \
--l1-size-gb "$LMCACHE_L1_SIZE_GB" \
--l1-init-size-gb "$LMCACHE_L1_INIT_SIZE_GB" \
--max-gpu-workers 1 \
--max-cpu-workers 8 \
--chunk-size 1024 \
--l1-align-bytes 16384 \
--eviction-trigger-watermark 0.85 \
--eviction-ratio 0.10 \
--eviction-policy LRU \
--supported-transfer-mode lmcache_driven \
--no-separate-object-groups \
> "$LMCACHE_SERVER_LOG" 2>&1 &
LMCACHE_SERVER_PID=$!
LMCACHE_READY=0
for _ in $(seq 1 60); do
if ! kill -0 "$LMCACHE_SERVER_PID" 2>/dev/null; then
echo "LMCache server died during startup." >&2
cat "$LMCACHE_SERVER_LOG" >&2
exit 1
fi
if curl --output /dev/null --silent --fail \
"http://127.0.0.1:$LMCACHE_HTTP_PORT/healthcheck"; then
LMCACHE_READY=1
break
fi
sleep 2
done
if [ "$LMCACHE_READY" -ne 1 ]; then
echo "LMCache server did not become healthy in time." >&2
cat "$LMCACHE_SERVER_LOG" >&2
exit 1
fi

unset VLLM_USE_SIMPLE_KV_OFFLOAD
OFFLOAD_ARGS=(
--kv-transfer-config
'{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_connector_extra_config":{"load_async":true}}'
"{\"kv_connector\":\"LMCacheMPConnector\",\"kv_role\":\"kv_both\",\"kv_connector_extra_config\":{\"lmcache.mp.host\":\"$LMCACHE_CONNECT_HOST\",\"lmcache.mp.port\":$LMCACHE_PORT,\"lmcache.mp.mq_timeout\":$LMCACHE_MQ_TIMEOUT}}"
)
fi
;;
Comment on lines 254 to +259

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔴 The new LMCache arm's --kv-transfer-config sets kv_connector":"LMCacheMPConnector" but omits kv_connector_module_path, unlike the two sibling scripts (dsv4_fp4_mi355x_vllm.sh:344, kimik2.5_fp4_b200.sh:150) that wire up the same connector. Since LMCacheMPConnector is out-of-tree, vLLM will not be able to locate the class and vllm serve should fail at KV-connector construction, breaking every point in the new dsv4-fp4-b300-vllm-agentic-lmcache config. Fix by adding "kv_connector_module_path":"lmcache.integration.vllm.lmcache_mp_connector" to the config dict.

Extended reasoning...

The bug: In the lmcache) case of benchmarks/single_node/agentic/dsv4_fp4_b300_vllm.sh (around lines 254-259), OFFLOAD_ARGS is built as:

\n"{\"kv_connector\":\"LMCacheMPConnector\",\"kv_role\":\"kv_both\",\"kv_connector_extra_config\":{...}}"\n

This JSON blob is missing the kv_connector_module_path key.

Why it matters: LMCacheMPConnector is not a vLLM built-in connector (unlike MooncakeStoreConnector, which vLLM resolves through its internal connector-factory registry). It lives in the external lmcache package, at lmcache.integration.vllm.lmcache_mp_connector. vLLM's KVTransferConfig resolves out-of-tree connectors by dynamically importing the module named in kv_connector_module_path before looking up the connector class by name — without that field, vLLM has no way to find the class.

Confirmed by direct comparison with sibling scripts: I checked the two other scripts in this repo that wire up the identical LMCacheMPConnector:

  • benchmarks/single_node/agentic/dsv4_fp4_mi355x_vllm.sh:344: "kv_connector\":\"LMCacheMPConnector\",\"kv_connector_module_path\":\"lmcache.integration.vllm.lmcache_mp_connector\"...
  • benchmarks/single_node/agentic/kimik2.5_fp4_b200.sh:150 (also lmcache 0.5.1, matching this PR's version exactly): same pattern.

Both sibling scripts explicitly set kv_connector_module_path alongside kv_connector. This new B300 arm is the only one of the three that omits it.

Why nothing else catches this: The script does run python3 -c \"import lmcache.integration.vllm.lmcache_mp_connector\" earlier, but that only verifies the module is importable in the launcher process — it says nothing to vLLM about where to import the connector class from when it parses --kv-transfer-config. bash -n, the matrix-logic pytest suite, and generate_sweep_configs.py all validate shell syntax and config-generation logic, not the semantics of the JSON payload passed to vLLM, so none of them would catch this.

Step-by-step proof of the failure:

  1. KV_OFFLOAD_BACKEND=lmcache is selected, LMCache MP server starts and passes its healthcheck.
  2. OFFLOAD_ARGS is set to --kv-transfer-config '{\"kv_connector\":\"LMCacheMPConnector\",\"kv_role\":\"kv_both\",\"kv_connector_extra_config\":{...}}' (no kv_connector_module_path).
  3. vllm serve parses this into a KVTransferConfig. Since kv_connector_module_path is absent, vLLM falls back to its built-in connector-factory lookup by name.
  4. LMCacheMPConnector is not registered there (it is only importable via the external lmcache package), so the lookup fails and vLLM raises an error during KV-connector construction, before the server can start serving.
  5. Every concurrency point in the new dsv4-fp4-b300-vllm-agentic-lmcache search space (17 points total) uses this same code path, so the entire new config fails at server startup, not just some points.

Fix: add \"kv_connector_module_path\":\"lmcache.integration.vllm.lmcache_mp_connector\" to the JSON dict at the same spot the two sibling scripts do, e.g.:

\n"{\"kv_connector\":\"LMCacheMPConnector\",\"kv_connector_module_path\":\"lmcache.integration.vllm.lmcache_mp_connector\",\"kv_role\":\"kv_both\",\"kv_connector_extra_config\":{...}}"\n

Severity is normal, not a nit or pre-existing issue: this is new code introduced by this PR, and merging as-is causes a concrete, deterministic startup failure for the entire new lmcache arm — not a stylistic or description mismatch.

*)
echo "Error: unsupported B300 KV_OFFLOAD_BACKEND='$KV_OFFLOAD_BACKEND'" >&2
exit 1
;;
esac

PARALLEL_ARGS=(--tensor-parallel-size "$TP" --data-parallel-size 1)
if [ "$DP_ATTENTION" = "true" ]; then
PARALLEL_ARGS=(--tensor-parallel-size 1 --data-parallel-size "$TP")
fi

EP_ARGS=()
TP_ARGS=()
if [ "$DP_ATTENTION" = "true" ]; then
# LMCacheMPConnector exports the KV cache to the LMCache server through
# legacy CUDA IPC handles, and expandable-segment (cuMem/VMM) allocations
# cannot be exported that way (register_kv_caches fails with
# cudaErrorInvalidValue, same failure mode as --enable-cumem-allocator on
# the B200 lmcache arm in PR #2231), so the lmcache arm keeps the stock
# caching allocator.
if [ "$KV_OFFLOAD_BACKEND" != "lmcache" ]; then
export PYTORCH_ALLOC_CONF=expandable_segments:True
fi
else
export VLLM_ALLREDUCE_USE_FLASHINFER=1
export VLLM_FLASHINFER_ALLREDUCE_BACKEND=auto
TP_ARGS+=(--disable-custom-all-reduce)
fi

MODE_ARGS=()
if [ "$EP_SIZE" -gt 1 ]; then
EP_ARGS=(--enable-expert-parallel)
MODE_ARGS+=(
--enable-expert-parallel
--enable-ep-weight-filter
--moe-backend deep_gemm_amxf4_mega_moe
)
fi
if [ "$DP_ATTENTION" = "true" ]; then
MODE_ARGS+=(
--prefill-schedule-interval 8
--max-num-batched-tokens 8192
)
fi

# AgentX concurrency counts live session trees, not individual requests.
# Subagent fan-out can push instantaneous request concurrency above CONC, so
# leave 2x headroom rather than clipping those bursts at the scheduler.
MAX_NUM_SEQS=$((2 * CONC))
if [ "$MAX_NUM_SEQS" -eq 128 ]; then
MAX_NUM_SEQS=136
if [ "$DP_ATTENTION" = "true" ]; then
# The DEP source recipe enforces 2*CONC = DP_WORLD_SIZE*MAX_NUM_SEQS.
MAX_NUM_SEQS=$((2 * CONC / TP))
else
# Preserve the previous TP4 scheduler headroom for agentic fan-out.
MAX_NUM_SEQS=$((2 * CONC))
fi
CUDA_GRAPH_CAPTURE_SIZES=""
for ((capture_size = 1; capture_size <= MAX_NUM_SEQS; capture_size++)); do
if [ -n "$CUDA_GRAPH_CAPTURE_SIZES" ]; then
CUDA_GRAPH_CAPTURE_SIZES+=","
fi
CUDA_GRAPH_CAPTURE_SIZES+="$capture_size"
done
COMPILATION_CONFIG="{\"cudagraph_mode\":\"FULL_DECODE_ONLY\",\"cudagraph_capture_sizes\":[${CUDA_GRAPH_CAPTURE_SIZES}],\"mode\":0}"

echo "Starting vllm server..."
export TORCH_CUDA_ARCH_LIST="10.0"
export PYTHONNOUSERSITE=1
export VLLM_FLOAT32_MATMUL_PRECISION=high

vllm serve "$MODEL_PATH" --served-model-name "$MODEL" \
--host 0.0.0.0 \
--port "$VLLM_BACKEND_PORT" \
--trust-remote-code \
--kv-cache-dtype fp8 \
--block-size 256 \
"${PARALLEL_ARGS[@]}" \
"${VLLM_CP_ARGS[@]}" \
"${EP_ARGS[@]}" \
--compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE","custom_ops":["all"]}' \
--attention_config.use_fp4_indexer_cache=True \
--tokenizer-mode deepseek_v4 \
--tool-call-parser deepseek_v4 \
--enable-auto-tool-choice \
--reasoning-parser deepseek_v4 \
--enable-prefix-caching \
--no-disable-hybrid-kv-cache-manager \
--max-num-seqs "$MAX_NUM_SEQS" \
"${OFFLOAD_ARGS[@]}" > "$SERVER_LOG" 2>&1 &
{ set +x; } 2>/dev/null
VLLM_CMD=(
vllm serve "$MODEL_PATH" --served-model-name "$MODEL"
--host 0.0.0.0
--port "$VLLM_BACKEND_PORT"
--gpu-memory-utilization 0.96
--trust-remote-code
--no-enable-flashinfer-autotune
--no-disable-hybrid-kv-cache-manager
--max-num-seqs "$MAX_NUM_SEQS"
--kv-cache-dtype fp8
--block-size 256
--max-model-len 1048576
--attention-config '{"use_fp4_indexer_cache":true,"backend":"FLASHINFER_MLA_SPARSE_DSV4","use_prefill_query_quantization":true}'
--disable-uvicorn-access-log
--tokenizer-mode deepseek_v4
--tool-call-parser deepseek_v4
--enable-auto-tool-choice
--reasoning-parser deepseek_v4
--compilation-config "$COMPILATION_CONFIG"
"${PARALLEL_ARGS[@]}"
"${TP_ARGS[@]}"
"${MODE_ARGS[@]}"
"${OFFLOAD_ARGS[@]}"
)
printf '%q ' "${VLLM_CMD[@]}" | tee "$RESULT_DIR/vllm_command.txt"
printf '\n' | tee -a "$RESULT_DIR/vllm_command.txt"
"${VLLM_CMD[@]}" > "$SERVER_LOG" 2>&1 &
SERVER_PID=$!
echo "Server PID: $SERVER_PID"

Expand Down
Loading
Loading