feat: Kimi-K2.5-MXFP4 LMCache MP offloading for MI355X agentic benchmarks

[andyluo7]

Summary

Adds chunked KV loading and vLLM scheduler assertion patches to the Kimi-K2.5-MXFP4 LMCache MP offloading pipeline for MI355X agentic benchmarks. These patches are required for stable, high-concurrency operation (c≥32) on ROCm.

What is LMCache MP?

LMCache MP (Multi-Process) uses an external LMCache server process that owns a large CPU DRAM pool (2.5 TB) for semantic KV caching. Unlike vLLM's built-in SimpleCPUOffloadConnector which evicts/restores KV blocks reactively, LMCache MP:

• Caches KV at the semantic level — prefix-matched KV blocks are stored once and shared across requests
• Decouples CPU memory management from vLLM's scheduler via ZMQ IPC
• Achieves 3× throughput at high concurrency where native offloading collapses

Changes in this PR

Two new ROCm-specific runtime patches added to sitecustomize.py:

1. Chunked connector patch (chunked_connector_patch.py): Caps external KV tokens loaded per scheduling step via CHUNKED_LMCACHE_MAX_TOKENS_PER_LOAD (default 32768). Without this, at c≥32, LMCache attempts to restore more KV blocks than the GPU pool can hold simultaneously, causing a deadlock where no request can make progress.

2. Scheduler assertion patch (scheduler_assertion_patch.py): Handles stale KV transfer finished_recving notifications gracefully. Under high concurrency, asynchronous H2D transfers can complete after the scheduler has already moved a request past WAITING_FOR_REMOTE_KVS → RUNNING, triggering an assertion crash.

Benchmark Results (19-config sweep on MI355X)

Environment: vLLM 0.21.0, ROCm 7.2.2, 8× MI355X (288 GB VRAM each), mia1-p02-g17
Model: amd/Kimi-K2.5-MXFP4 (MoE, 262K context)
Benchmark: AgentX v0.3 agentic trace replay, 1800s duration, 375 conversations

TP=8: No Offloading vs LMCache

Conc	none Reqs	lmc Reqs	Δ Reqs	none ITL p50	lmc ITL p50	none Out tok/s	lmc Out tok/s	Δ Throughput
16	789	721	−8.6%	44ms	48ms	214.7	204.6	−4.7%
24	563	607	+7.8%	76ms	76ms	159.3	176.2	+10.6%
32	168	257	+53%	493ms	68ms	32.8	97.0	+196%
40	174	135	−22%	511ms	67ms	33.5	41.2	+23%
48	175	75	−57%	511ms	90ms	34.9	32.5	−6.9%

Key findings:

• LMCache crossover at c=24 — below this, GPU KV cache is sufficient and LMCache adds overhead
• Sweet spot at c=32 — 3× output throughput (97 vs 33 tok/s), 7.3× lower ITL (68ms vs 493ms)
• Diminishing returns past c=40 — PCIe transfer overhead dominates at extreme saturation

TP=4: LMCache extends usable range

Conc	none Out tok/s	lmc Out tok/s	Δ
16	90.4	98.2	+8.6%
24	21.9	35.7	+63%
32	19.8	26.5	+34%
40	20.5	31.5	+54%

ROCm-Specific Patches (full stack)

Patch	Purpose	Env Var
Demand-pinned memory	Lazy-expand LMCache's 2.5 TB pinned pool on ROCm (avoids 10-min startup stall)	LMCACHE_ROCM_DEMAND_PINNED_ALLOCATOR=1
MP block transfer fallback	Python fallback for multi_layer_block_kv_transfer (no CUDA kernel on ROCm)	LMCACHE_ROCM_MP_BLOCK_FALLBACK=1
Chunked KV loading (new)	Cap tokens per scheduling step to prevent GPU block exhaustion deadlock	CHUNKED_LMCACHE_MAX_TOKENS_PER_LOAD=32768
Scheduler assertion fix (new)	Handle stale KV transfer notifications gracefully	(always active)

How to Run

export OFFLOADING=lmcache
export TP=8
export CONC=32
# Optional: tune chunk size (default 32768, set 0 to disable)
export CHUNKED_LMCACHE_MAX_TOKENS_PER_LOAD=32768

Related

• Upstream LMCache PR for chunked loading: [MP] Fix GPU block exhaustion deadlock at high concurrency with chunked KV loading LMCache/LMCache#3382
• Existing MiniMax-M2.5 LMCache PR: feat: add vLLM + LMCache CPU offloading for MiniMax-M2.5 agentic benchmark on AMD GPUs #1262

---

Note

Medium Risk
Monkey-patches LMCache connector and vLLM scheduler scheduling/KV-transfer behavior at runtime; incorrect caps or stale-notification handling could affect correctness or throughput under load, though scope is limited to the benchmark ROCm patch path.

Overview
Extends the MI355X Kimi LMCache MP benchmark bootstrap with two runtime monkey-patches dropped into the existing ROCm sitecustomize patch directory and loaded via PYTHONPATH.

Chunked KV loading (CHUNKED_LMCACHE_MAX_TOKENS_PER_LOAD, default 32768): patches LMCacheMPConnector so each scheduling step only reports/loads a capped slice of external matched tokens (with per-request chunk state and cleanup on finish). This avoids GPU KV block exhaustion deadlocks at high concurrency (e.g. c≥32). Set 0 to disable chunking.

Scheduler KV transfer fix: replaces vLLM’s _update_from_kv_xfer_finished handling so stale finished_recving notifications (request already past WAITING_FOR_REMOTE_KVS) log a warning instead of asserting, while still honoring normal receive/send completion paths.

The lmcache offload path now writes both patch modules alongside the existing ROCm patches and exports the chunk env var by default.

^{Reviewed by Cursor Bugbot for commit e5bd4d8. Bugbot is set up for automated code reviews on this repo. Configure here.}

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 39159f7. Configure here.}

[seungrokj]

LGTM