-
Notifications
You must be signed in to change notification settings - Fork 224
[LMCache] Add LMCache configs for dsv4 vllm b200/b300 agentix setups #2153
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
10e6639
73137f1
951f2c8
4244f91
380ee5a
6623385
d9209dd
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -21,7 +21,7 @@ set -x | |
| # Required env vars: | ||
| # MODEL, TP, CONC, KV_OFFLOADING, TOTAL_CPU_DRAM_GB, RESULT_DIR | ||
| # | ||
| # KV_OFFLOADING=dram requires KV_OFFLOAD_BACKEND=mooncake. | ||
| # KV_OFFLOADING=dram requires KV_OFFLOAD_BACKEND=mooncake or lmcache. | ||
|
|
||
| source "$(dirname "$0")/../../benchmark_lib.sh" | ||
|
|
||
|
|
@@ -101,15 +101,20 @@ export VLLM_PREFIX_CACHE_RETENTION_INTERVAL=32768 | |
| 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="" | ||
|
|
||
| OFFLOAD_ARGS=() | ||
|
|
||
| if require_agentic_kv_offload_backend mooncake; then | ||
| if agentic_kv_offload_enabled; then | ||
| case "$KV_OFFLOAD_BACKEND" in | ||
| mooncake) | ||
| require_agentic_kv_offload_backend mooncake | ||
| # Embedded mode contributes one segment per GPU rank to a shared | ||
| # distributed store, so pre-divide the aggregate host-memory budget. | ||
| PER_RANK_GB=$((TOTAL_CPU_DRAM_GB / GPU_COUNT)) | ||
|
|
@@ -118,6 +123,7 @@ if require_agentic_kv_offload_backend mooncake; then | |
| agentic_pip_install --quiet --no-cache-dir --no-deps \ | ||
| --force-reinstall "mooncake-transfer-engine-cuda13==$MOONCAKE_VERSION" | ||
| python3 -c "from mooncake.store import MooncakeDistributedStore" >/dev/null | ||
| python3 "$(dirname "$0")/patch_vllm_pr45406.py" | ||
|
|
||
| MOONCAKE_MASTER_PORT=$((PORT + 12000)) | ||
| MOONCAKE_CONFIG_PATH="$RESULT_DIR/mooncake_config.json" | ||
|
|
@@ -169,6 +175,96 @@ EOF | |
| --kv-transfer-config | ||
| '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_connector_extra_config":{"load_async":true}}' | ||
| ) | ||
| ;; | ||
| 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 | ||
| # Async KV loads park requests in WAITING_FOR_REMOTE_KVS holding | ||
| # blocks; without the PR #45406 scheduler fix the waiting-queue scan | ||
| # can freeze permanently once nothing is running (hung warmup | ||
| # stragglers in PR #2153 bring-up). | ||
| python3 "$(dirname "$0")/patch_vllm_pr45406.py" | ||
|
|
||
| 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 | ||
|
|
||
| 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 \ | ||
| > "$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\":\"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}}" | ||
| ) | ||
| ;; | ||
| *) | ||
| echo "Error: unsupported KV_OFFLOAD_BACKEND '$KV_OFFLOAD_BACKEND' (expected one of: mooncake, lmcache)" >&2 | ||
| exit 1 | ||
| ;; | ||
| esac | ||
|
Comment on lines
+263
to
+267
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🔴 The new Extended reasoning...What the bug isIn the new {
"kv_connector": "LMCacheMPConnector",
"kv_role": "kv_both",
"kv_connector_extra_config": {
"lmcache.mp.host": "...",
"lmcache.mp.port": ...,
"lmcache.mp.mq_timeout": ...
}
}It is missing the top-level Why the existing in-tree pattern says this is requiredEvery other in-tree call-site of
Those two sites were authored by independent PRs and both pass the module path. The consistency of that pattern is strong evidence the key is load-bearing, not decorative. By contrast, Why the pre-existing import check does not save usThe script does run Impact and step-by-step failure
Because both FixAdd the module path in the JSON built at OFFLOAD_ARGS=(
--kv-transfer-config
"{\"kv_connector\":\"LMCacheMPConnector\",\"kv_connector_module_path\":\"lmcache.integration.vllm.lmcache_mp_connector\",\"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}}"
)This mirrors the pattern already established at
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. No we don't need this in newer versions of vLLM |
||
| fi | ||
|
|
||
| PARALLEL_ARGS=(--tensor-parallel-size "$TP" --data-parallel-size 1) | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,63 @@ | ||
| #!/usr/bin/env python3 | ||
| """Backport the scheduler fix from vllm-project/vllm PR #45406.""" | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import importlib.util | ||
| from pathlib import Path | ||
|
|
||
|
|
||
| PATCH_MARKER = "See https://github.com/vllm-project/vllm/issues/45388" | ||
|
|
||
|
|
||
| def main() -> None: | ||
| spec = importlib.util.find_spec("vllm") | ||
| if spec is None or not spec.submodule_search_locations: | ||
| raise RuntimeError("Could not locate the installed vllm package") | ||
|
|
||
| package_root = Path(next(iter(spec.submodule_search_locations))) | ||
| scheduler = package_root / "v1" / "core" / "sched" / "scheduler.py" | ||
| text = scheduler.read_text() | ||
| if PATCH_MARKER in text: | ||
| print(f"vLLM PR #45406 backport already present in {scheduler}") | ||
| return | ||
|
|
||
| old = """ if request.has_encoder_inputs: | ||
| self.encoder_cache_manager.free(request) | ||
| break | ||
|
|
||
| # KVTransfer: the connector uses this info to determine | ||
| """ | ||
| new = """ if request.has_encoder_inputs: | ||
| self.encoder_cache_manager.free(request) | ||
| if self.running: | ||
| # Running requests will free blocks when they | ||
| # complete; stop here to preserve queue-order | ||
| # admission. | ||
| break | ||
| # Nothing is running, so no future event frees blocks and | ||
| # stopping at this request would freeze this state | ||
| # permanently. Requests behind this one may hold blocks | ||
| # while parked (async KV loads in WAITING_FOR_REMOTE_KVS) | ||
| # and are only promoted when this traversal reaches them. | ||
| # Keep scanning so they can be promoted, scheduled, and | ||
| # eventually free the blocks this request needs. | ||
| # See https://github.com/vllm-project/vllm/issues/45388 | ||
| request_queue.pop_request() | ||
| step_skipped_waiting.prepend_request(request) | ||
| continue | ||
|
|
||
| # KVTransfer: the connector uses this info to determine | ||
| """ | ||
| if text.count(old) != 1: | ||
| raise RuntimeError( | ||
| f"Expected exactly one vLLM PR #45406 patch target in {scheduler}" | ||
| ) | ||
|
|
||
| scheduler.write_text(text.replace(old, new, 1)) | ||
| compile(scheduler.read_text(), str(scheduler), "exec") | ||
| print(f"Applied vLLM PR #45406 backport to {scheduler}") | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| main() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
generally we wanna avoid patches too
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah, this is what I picked from #2138. Once that is in, I will remove this.