Skip to content

kimik3-fp4-mi355x-vllm-agentic-mtp: add LMCache DRAM KV-offload arm - #2598

Open
sammshen wants to merge 15 commits into
mainfrom
kimik3-mi355x-lmcache
Open

kimik3-fp4-mi355x-vllm-agentic-mtp: add LMCache DRAM KV-offload arm#2598
sammshen wants to merge 15 commits into
mainfrom
kimik3-mi355x-lmcache

Conversation

@sammshen

Copy link
Copy Markdown
Collaborator

What

Adds a dedicated config key kimik3-fp4-mi355x-vllm-agentic-mtp-lmcache with an LMCache DRAM KV-offload arm at TP8 conc 4 / 8 / 10 / 12, on top of the unchanged DSpark MTP serving stack of kimik3-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

  • configs/amd-master.yaml — new key, purely additive (the base key is byte-identical to main): kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.5.4rc2" }, conc-list: [4, 8, 10, 12], spec-decoding: mtp.
  • benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh — new lmcache) case arm, modeled on the MiniMax-M3 lmcache arm:
    • installs the LMCache 0.5.4rc2 ROCm wheel (gfx950) with --no-deps plus its missing runtime deps, leaving the image's torch/ROCm stack untouched;
    • starts a single MP server per the Kimi-K3 recipe: --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;
    • wires vLLM to it via LMCacheMPConnector (lmcache.mp.port), keeping the DSpark SPEC_ARGS untouched;
    • unknown backends now fail fast instead of silently running GPU-resident.
  • perf-changelog.yaml — selects only the new key.
  • The L1 budget consumes the generated TOTAL_CPU_DRAM_GB verbatim 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-groups is the supported handling for the multi-KV-group layout that previously blocked LMCache-under-MTP arms.

Validation

process_changelog.py run exactly as CI does (base = main) emits a 4-row matrix, nothing else:

kimik3_tp8_conc4_kvdram-lmcache_spec-mtp   | lmcache 0.5.4rc2 | throughput
kimik3_tp8_conc8_kvdram-lmcache_spec-mtp   | lmcache 0.5.4rc2 | throughput
kimik3_tp8_conc16_kvdram-lmcache_spec-mtp  | lmcache 0.5.4rc2 | throughput
kimik3_tp8_conc16_kvdram-lmcache_spec-mtp  | lmcache 0.5.4rc2 | eval-only

generate_sweep_configs.py test-config passes for both the new key and the (unchanged) base key; bash -n passes on the modified script.


Reopened from #2583 on a branch in this repo rather than a fork, so run-sweep.yml runs with secrets and /stage-results can find the results. Same commits, no content changes.

Co-authored-by: ApostaC yihua98@uchicago.edu

sammshen and others added 13 commits August 12, 2026 16:54
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.
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>
@github-actions

Copy link
Copy Markdown
Contributor

@claude claude Bot left a comment

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.

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.yaml entry 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 new kimik3-fp4-mi355x-vllm-agentic-mtp-lmcache key — sets conc-list: [4, 8, 10, 12]. Running utils/matrix_logic/generate_sweep_configs.py test-config against this key confirms it emits conc4, conc8, conc10, conc12 (throughput) and conc12 (eval-only) — conc16 is never generated. Notably, this also contradicts the PR description's own "Validation" section, which independently claims the matrix output is conc4/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 shows KV_OFFLOAD_BACKEND_METADATA is only consumed by multi_node/amd_utils/server_sglang.sh and by the result-aggregation/labeling scripts (process_agentic_result.py) — not by this single-node vLLM script. So the config's version: "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.py selects which runs to include using the config-keys field, not the free-form description strings, so this inaccurate prose has zero effect on what actually runs or how results are labeled — the real sweep points come from configs/amd-master.yaml's conc-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:

    1. Read configs/amd-master.yaml diff: kimik3-fp4-mi355x-vllm-agentic-mtp-lmcachesearch-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 }].
    2. Run generate_sweep_configs.py test-config for this key → emits exp-names conc4, conc8, conc10, conc12 (throughput), conc12 (eval-only). No conc16 appears anywhere.
    3. Compare to perf-changelog.yaml:5925: "...at TP8 conc 4/8/16..." — mismatch confirmed (16 wrong, 10/12 missing).
    4. Read kimik3_fp4_mi355x_mtp.sh:153: LMCACHE_VERSION="0.5.4rc2" — a literal, not a variable read from config or env.
    5. Grep the repo for KV_OFFLOAD_BACKEND_METADATA → only hits in multi_node/amd_utils/server_sglang.sh and the aggregation/labeling scripts, never in this script.
    6. 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_METADATA claim — 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.

Comment on lines +180 to +197
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 ""
)

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 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.

Comment on lines +154 to +162
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

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 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-deps block 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 ROCm expanded_assets index, differing only in the LMCACHE_VERSION string (0.5.4rc2 vs 0.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 server argv assembly (--host/--port/--http-host/--http-port/--l1-size-gb/--l1-init-size-gb/--chunk-size/--eviction-policy LRU/--supported-transfer-mode lmcache_driven), the append_command + background-launch + wait_for_ready --endpoint .../healthcheck polling shape, and the LMCacheMPConnector kv-transfer-config JSON 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.

@github-actions

Copy link
Copy Markdown
Contributor

1 similar comment
@github-actions

Copy link
Copy Markdown
Contributor

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

agentx AgentX benchmarks, recipes, and infrastructure AMD full-sweep-enabled

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

2 participants