kimik3-fp4-mi355x-vllm-agentic-mtp: add LMCache DRAM KV-offload arm - #2598
kimik3-fp4-mi355x-vllm-agentic-mtp: add LMCache DRAM KV-offload arm#2598sammshen wants to merge 15 commits into
Conversation
Add an lmcache kv-offload-backend point at TP8 conc 10 on top of the existing DSpark MTP serving stack, mirroring the vllm-simple offload arm for a direct backend comparison. The benchmark script gains an lmcache case arm that installs the LMCache 0.5.4rc1 ROCm wheel (torch/ROCm stack untouched), starts one MP server per the Kimi-K3 recipe (chunk size 768 = K3 unified block size at 8 GPUs, --separate-object-groups for the hybrid KDA/MLA two-group KV layout, --enable-extra-logging, --max-cpu-workers 8 --max-gpu-workers 1), and wires vLLM to it via LMCacheMPConnector.
A separate kimik3-fp4-mi355x-vllm-agentic-mtp-lmcache key lets the changelog select only the LMCache points instead of re-running the resident and vllm-simple arms of the base key. The base key returns to its upstream shape.
The LMCache MP server's L1 lives in /dev/shm and the script rejects budgets above 90% of free shm. mi355x-amds nodes mount ~1.5 TB of shm (cap ~1360 GB), so 0.50's 1499 GB budget failed the check in run 31644286169. 0.40 generates ~1199 GB, which fits with margin.
vLLM sizes the K3 unified attention block to 1536 tokens on the MI355X fp8-KV TRITON_MLA path (attention page >= mamba page), and the MP connector asserts chunk %% block == 0, so the recipe's CUDA-path 768 fails engine init (run 31644990546).
The connector requires the chunk to be a multiple of every engine KV group's tokens_per_block. On this stack the hybrid layout registers attention groups at 1536 and a KDA state group at 3072 (run 31645828378), so 1536 fails registration; 3072 is the minimum valid chunk.
Auto mode loads both transfer paths; pin server-driven STORE/RETRIEVE (as the MiniMax-M3 arm does) so the benchmark measures one deterministic path. The L1 stays shm-backed either way, so the /dev/shm capacity check still applies.
The default 300s read-lock TTL expires under a single GPU worker serializing huge K3 transfers: run 31648224111 logged 57k finish-read-on-non-read-locked-key warnings starting exactly at warmup+300s, followed by a GPU illegal-access crash mid-profile. Match the MiniMax-M3 arm's 7200s read TTL.
This reverts commit 20b4fda.
Signed-off-by: ApostaC <yihua98@uchicago.edu>
Signed-off-by: ApostaC <yihua98@uchicago.edu>
Signed-off-by: ApostaC <yihua98@uchicago.edu>
Signed-off-by: ApostaC <yihua98@uchicago.edu>
# Conflicts: # perf-changelog.yaml
|
see unofficial run visualizer at https://inferencex.semianalysis.com/inference?unofficialRun=31760258443 |
There was a problem hiding this comment.
Additional findings (outside current diff — PR may have been updated during review):
-
🟡
perf-changelog.yaml:5921-5927— The new perf-changelog.yaml entry (lines 5921-5927) says the LMCache arm runs 'TP8 conc 4/8/16', but configs/amd-master.yaml actually sets conc-list: [4, 8, 10, 12] for kimik3-fp4-mi355x-vllm-agentic-mtp-lmcache — 16 is never swept and 10/12 are omitted from the description. It also claims the LMCache version is 'consumed by the script via KV_OFFLOAD_BACKEND_METADATA', but kimik3_fp4_mi355x_mtp.sh:153 hardcodes LMCACHE_VERSION="0.5.4rc2" directly and never reads that env var (only the multi-node sglang path does). Please correct both statements in the description text.Extended reasoning...
The new
perf-changelog.yamlentry added by this PR (lines 5921-5927) contains two factual inaccuracies in its free-form description text, both of which are independently verifiable against the code in the same PR.Concurrency mismatch. The description reads "at TP8 conc 4/8/16", but the config it accompanies —
configs/amd-master.yaml's newkimik3-fp4-mi355x-vllm-agentic-mtp-lmcachekey — setsconc-list: [4, 8, 10, 12]. Runningutils/matrix_logic/generate_sweep_configs.py test-configagainst this key confirms it emitsconc4,conc8,conc10,conc12(throughput) andconc12(eval-only) —conc16is never generated. Notably, this also contradicts the PR description's own "Validation" section, which independently claims the matrix output isconc4/conc8/conc16/conc16— so the stale text was seemingly carried over from an earlier iteration of the config rather than the one actually in the diff.Version-plumbing mismatch. The description says the LMCache version "is... consumed by the script via KV_OFFLOAD_BACKEND_METADATA". Looking at
benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh:153, the version is not read from any config-plumbed variable at all — it's a hardcoded literal:LMCACHE_VERSION="0.5.4rc2". Grepping the repo showsKV_OFFLOAD_BACKEND_METADATAis only consumed bymulti_node/amd_utils/server_sglang.shand by the result-aggregation/labeling scripts (process_agentic_result.py) — not by this single-node vLLM script. So the config'sversion: "0.5.4rc2"field is decorative here, exactly like every other kv-offload-backend version field in this file; bumping it would not change what the script installs.Why nothing catches this today.
process_changelog.pyselects which runs to include using theconfig-keysfield, not the free-formdescriptionstrings, so this inaccurate prose has zero effect on what actually runs or how results are labeled — the real sweep points come fromconfigs/amd-master.yaml'sconc-list. That is exactly why this is only a documentation nit and not a functional bug: a maintainer or future reader of the changelog will simply be told the wrong concurrency points and the wrong version-plumbing mechanism.Proof, step by step:
- Read
configs/amd-master.yamldiff:kimik3-fp4-mi355x-vllm-agentic-mtp-lmcache→search-space: [{ tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.5.4rc2" }, conc-list: [4, 8, 10, 12], spec-decoding: mtp }]. - Run
generate_sweep_configs.py test-configfor this key → emits exp-namesconc4,conc8,conc10,conc12(throughput),conc12(eval-only). Noconc16appears anywhere. - Compare to
perf-changelog.yaml:5925:"...at TP8 conc 4/8/16..."— mismatch confirmed (16 wrong, 10/12 missing). - Read
kimik3_fp4_mi355x_mtp.sh:153:LMCACHE_VERSION="0.5.4rc2"— a literal, not a variable read from config or env. - Grep the repo for
KV_OFFLOAD_BACKEND_METADATA→ only hits inmulti_node/amd_utils/server_sglang.shand the aggregation/labeling scripts, never in this script. - Compare to
perf-changelog.yaml:5926: claims the version is "consumed by the script via KV_OFFLOAD_BACKEND_METADATA" — mismatch confirmed.
Suggested fix: update the description to say "at TP8 conc 4/8/10/12" and drop (or rephrase) the
KV_OFFLOAD_BACKEND_METADATAclaim — e.g. note that the version field is recorded in the config for documentation/labeling purposes only, since the script pins its own LMCache version independently. - Read
| LMCACHE_CMD=( | ||
| lmcache server | ||
| --host 127.0.0.1 | ||
| --port "$LMCACHE_PORT" | ||
| --http-host 127.0.0.1 | ||
| --http-port "$LMCACHE_HTTP_PORT" | ||
| --l1-size-gb "$LMCACHE_L1_SIZE_GB" | ||
| --l1-init-size-gb 10 | ||
| --chunk-size 3072 | ||
| --separate-object-groups | ||
| --enable-extra-logging | ||
| --extra-logging-interval 30 | ||
| --max-cpu-workers 8 | ||
| --max-gpu-workers 1 | ||
| --eviction-policy LRU | ||
| --supported-transfer-mode lmcache_driven | ||
| --shm-name "" | ||
| ) |
There was a problem hiding this comment.
🔴 The new lmcache MP server command omits --l1-read-ttl-seconds, so it runs at LMCache's 300s default read-lease TTL instead of the 7200s both sibling arms explicitly set for this exact workload class (minimaxm3_fp4_mi355x_mtp.sh:134, and dsv4_fp4_mi355x_vllm_mtp.sh with a documented rationale). Under GPU-KV saturation with this arm's 100k-330k-token prefixes, the lookup-to-retrieve gap can exceed 300s, expiring the lease while the object is still in L1 -- producing spurious cache misses/retrieve failures that would silently corrupt the offload hit-rate numbers this arm exists to measure. Suggest adding --l1-read-ttl-seconds 7200 to LMCACHE_CMD (kimik3_fp4_mi355x_mtp.sh:180-197), matching both sibling arms.
Extended reasoning...
The bug: LMCACHE_CMD in the new lmcache) arm (kimik3_fp4_mi355x_mtp.sh:180-197) starts the LMCache MP server without --l1-read-ttl-seconds, so it runs at LMCache's built-in 300-second default for the L1 read lease.
Why this matters here specifically: an LMCache read lock is a lease on a chunk that lookup() has already promised vLLM it can retrieve(). If more than the TTL elapses between those two calls, the lease expires while the chunk is still physically present in L1 -- it just becomes unreadable, producing a spurious cache miss (or an outright retrieve failure) rather than a correct DRAM-offload hit. This is not a hypothetical concern for this codebase: dsv4_fp4_mi355x_vllm_mtp.sh (lines 314-322) hit exactly this failure mode on the same connector/workload shape and fixed it by raising LMCACHE_L1_READ_TTL_SECONDS to 7200 with a comment explaining that TP8/conc32 agentic queues can spend >300s between lookup and retrieve under GPU-KV saturation. minimaxm3_fp4_mi355x_mtp.sh:134 -- the arm this new kimik3 arm's own comment says it is modeled on -- independently carries the same --l1-read-ttl-seconds 7200 flag in the identical MP-server invocation shape (same connector, same --supported-transfer-mode lmcache_driven).
Why the new arm is at least as exposed, not less: the kimik3 arm's own comment (line 208) notes '100k-330k-token agentic prefixes make single retrieves large,' and the author raised lmcache.mp.mq_timeout to 6000.0s (20x the default lease) specifically to give large single retrieves enough headroom. That is an internally inconsistent combination: the code anticipates individual MQ operations taking up to 6000s, but the read lease guarding the underlying object expires after only 300s -- roughly 20x sooner than the operation it is meant to survive. Concurrency here (up to 12) is lower than dsv4's cited conc32 example, which does reduce how often GPU-KV saturation is reached, but the much larger per-request prefixes (100k-330k tokens vs. dsv4's 8k-32k-scale) push the lookup-to-retrieve gap in the opposite direction, so the net exposure at this arm's own concurrency levels is plausibly comparable to or worse than dsv4's.
Why nothing else in the diff prevents this: the new arm installs LMCache fresh via pip and constructs LMCACHE_CMD from scratch (it does not source or extend either sibling's server-start logic), so it does not inherit the fix; it must set the flag explicitly, and it does not.
Concrete walkthrough: (1) a request with a 250k-token prefix arrives under conc=12; (2) vLLM's LMCacheMPConnector calls lookup(), which finds the prefix's chunks in L1 and returns a promise that they are retrievable; (3) GPU KV is saturated (TP8, MTP, large batch), so the actual retrieve() call is queued and delayed; (4) more than 300s elapses before retrieve executes; (5) the L1 read lease on those chunks has now expired even though the chunks are still resident in L1; (6) retrieve fails or falls back to a cache miss, and the benchmark either records a failed request or silently re-computes the prefix from scratch, corrupting the DRAM-offload hit-rate metric this dedicated -lmcache config key exists specifically to measure.
Fix: add --l1-read-ttl-seconds 7200 (or an env-overridable LMCACHE_L1_READ_TTL_SECONDS, matching the dsv4 pattern) to LMCACHE_CMD in kimik3_fp4_mi355x_mtp.sh, mirroring both sibling arms.
| LMCACHE_ROCM_INDEX="https://github.com/LMCache/LMCache/releases/expanded_assets/v${LMCACHE_VERSION}-rocm" | ||
| agentic_pip_install --quiet --no-cache-dir --no-deps \ | ||
| "sortedcontainers==2.4.0" \ | ||
| "opentelemetry-exporter-prometheus==0.61b0" \ | ||
| "cupy-rocm-7-0==14.1.1" \ | ||
| "lmcache==${LMCACHE_VERSION}" --find-links "$LMCACHE_ROCM_INDEX" | ||
| python3 -c \ | ||
| "import cupy; import lmcache.integration.vllm.lmcache_mp_connector; import opentelemetry.exporter.prometheus" \ | ||
| >/dev/null |
There was a problem hiding this comment.
🟡 The LMCache MP-server bootstrap in this new arm (pip install of the 4 pinned wheels, the python3 import smoke-test, server argv assembly, and append_command + background-launch + wait_for_ready polling at lines 154-213) is copy-pasted almost verbatim from the existing minimaxm3_fp4_mi355x_mtp.sh:85-159 lmcache arm, differing only in shard count, chunk-size/eviction knobs, and the mp.port vs mp.server_urls connector key. This is a pre-existing pattern (this is now the second copy) that a shared benchmark_lib.sh helper (e.g. start_lmcache_mp_servers) could remove; not blocking.
Extended reasoning...
Comparing kimik3_fp4_mi355x_mtp.sh:148-213 against the existing minimaxm3_fp4_mi355x_mtp.sh:85-159, the new lmcache) case arm is a near-verbatim copy of the one already in the MiniMax-M3 script:
- The
agentic_pip_install --no-depsblock installs the exact same 4 pinned wheels (sortedcontainers==2.4.0,opentelemetry-exporter-prometheus==0.61b0,cupy-rocm-7-0==14.1.1,lmcache==<version>) from the same ROCmexpanded_assetsindex, differing only in theLMCACHE_VERSIONstring (0.5.4rc2vs0.5.3). - The
python3 -c "import cupy; import lmcache.integration.vllm.lmcache_mp_connector; import opentelemetry.exporter.prometheus"smoke test is byte-for-byte identical. - The
lmcache serverargv assembly (--host/--port/--http-host/--http-port/--l1-size-gb/--l1-init-size-gb/--chunk-size/--eviction-policy LRU/--supported-transfer-mode lmcache_driven), theappend_command+ background-launch +wait_for_ready --endpoint .../healthcheckpolling shape, and theLMCacheMPConnectorkv-transfer-configJSON all follow the same structure.
The real, intentional differences are exactly the ones called out in the PR description: this arm runs a single MP server for the whole node instead of one per TP rank, uses different chunk-size/eviction/logging knobs tuned for the hybrid KDA/MLA layout, and connects via lmcache.mp.port instead of lmcache.mp.server_urls.
This isn't a functional bug — the script works correctly as written, bash -n passes, and the duplication doesn't cause incorrect behavior. It's a maintainability nit: this pattern now exists in two single-node agentic scripts, and the next lmcache arm (there's already a comment in the diff about a third planned once the upstream hybrid KV recovery fix lands) will very likely copy it a third time. benchmark_lib.sh already hosts exactly this class of shared helper (agentic_pip_install, wait_for_ready, append_command, require_agentic_kv_offload_backend), so extracting at minimum the pip-install + smoke-test block, and plausibly a parameterized start_lmcache_mp_servers helper (taking shard count, chunk-size, and extra flags) into benchmark_lib.sh, would fit the established pattern and remove the duplicate.
Step-by-step: (1) open minimaxm3_fp4_mi355x_mtp.sh:85-100 — the pip-install block matches kimik3_fp4_mi355x_mtp.sh:154-162 line for line except the version string; (2) both scripts' next line is the identical three-import python3 -c smoke test; (3) both then build an lmcache server argv array with the same flag set and launch it the same way with append_command, background &, and wait_for_ready --endpoint http://.../healthcheck; (4) both close with a kv-transfer-config JSON for LMCacheMPConnector differing only in the connector-key field. No behavior differs from consolidating these into a shared helper, so this is a safe, non-blocking cleanup for a future PR.
|
see unofficial run visualizer at https://inferencex.semianalysis.com/inference?unofficialRun=31760620783 |
1 similar comment
|
see unofficial run visualizer at https://inferencex.semianalysis.com/inference?unofficialRun=31760620783 |
What
Adds a dedicated config key
kimik3-fp4-mi355x-vllm-agentic-mtp-lmcachewith an LMCache DRAM KV-offload arm at TP8 conc 4 / 8 / 10 / 12, on top of the unchanged DSpark MTP serving stack ofkimik3-fp4-mi355x-vllm-agentic-mtp(same image, script, and topology). A separate key means the changelog selects only the LMCache points — the resident and vllm-simple arms of the base key are not re-run.Changes
kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.5.4rc2" }, conc-list: [4, 8, 10, 12], spec-decoding: mtp.lmcache)case arm, modeled on the MiniMax-M3 lmcache arm:--no-depsplus its missing runtime deps, leaving the image's torch/ROCm stack untouched;--chunk-size 3072(the least common multiple of this stack's KV-group block sizes: 1536-token attention groups plus a 3072-token KDA state group; the recipe's CUDA-path 768 is not a valid multiple here),--separate-object-groups(one object group per sliding-window size for the hybrid KDA/MLA layout, which has two KV-cache groups under MTP),--enable-extra-logging,--max-cpu-workers 8 --max-gpu-workers 1,--l1-size-gb $TOTAL_CPU_DRAM_GB, LRU eviction;LMCacheMPConnector(lmcache.mp.port), keeping the DSparkSPEC_ARGSuntouched;TOTAL_CPU_DRAM_GBverbatim per the agentic README, with an early /dev/shm capacity check (same as the MiniMax-M3 arm).LMCache + DSpark MTP + the rest of this recipe's feature set has been compatibility-tested;
--separate-object-groupsis the supported handling for the multi-KV-group layout that previously blocked LMCache-under-MTP arms.Validation
process_changelog.pyrun exactly as CI does (base = main) emits a 4-row matrix, nothing else:generate_sweep_configs.py test-configpasses for both the new key and the (unchanged) base key;bash -npasses on the modified script.Reopened from #2583 on a branch in this repo rather than a fork, so
run-sweep.ymlruns with secrets and/stage-resultscan find the results. Same commits, no content changes.Co-authored-by: ApostaC yihua98@uchicago.edu