[atom-vllm] LMCache offload for GLM-5.2

[kliuae]

Motivation

Adds support for KV cache offloading through LMCache for GLM-5.2 in ATOM-vLLM.

Technical Details

GLM-5.2's sparse MLA stores KV caches for MLA and indexer in a UniformTypeKVCacheSpecs group, but inherently they have different shapes. Overall, the group has 78 MLA and 21 indexer caches. LMCache's default connector uses the number of layers to size the caches, which is 78 for GLM-5.2, whereas there are 99 caches in total, causing conflict. This PR enables GLM-5.2 LMCache for atom-vllm by addressing the following.

• LMCache has a V3 connector that partitions caches by page shape and can handle GLM-5.2's different-shape caches case. This PR force select V3 when the KV cache layout requires it.
• However, V3 discovers its page-shape partitions lazily at the first transfer, but LMCacheEngine.store needs to allocate memory before that, and at that time metadata.get_shapes() is used for the shapes, despite the correct page group hasn't been built yet. This PR builds V3's shape map at registration, running discovery when the KV caches are registered instead of waiting for the first transfer, so that it's already populated when save path needs it.
• LMCache does not build an engine instance for the scheduler process, but when cleaning up aborted requests, it assumes there is one. This PR adds a null engine so that the cleanup can be performed.
• For MLA, LMCache by default enables save_only_first_node where rank 0 loads the KV tensors from the offloaded target, and sends them to the other ranks. The transfer occurs chunk by chunk, and in upstream LMCache, the receiving ranks retain every chunk until the whole batch is uploaded. In settings where GPU utilization is nearly full, this can cause the receiving rank to abort with out of resources, and while rank 0 is stilling waiting on them, leading to deadlock. This PR bounds the receives to 32-chunk-at-a-time on the receiving ranks so that the peak memory usage is capped at a lower level.
• For cache chunks that are pinned, LMCache only unpins them in wait_for_save, so a request that is never scheduled gets pinned forever, and they never gets evicted. This PR allows the atom-vllm path to unpin any lookup absent from the load set for two consecutive steps.

Test Plan

lm_eval gsm8k on MI355X
AgentX test

MODEL_PATH=zai-org/GLM-5.2-FP8

Server command

export LMCACHE_LOCAL_CPU=True
export LMCACHE_MAX_LOCAL_CPU_SIZE=200
export LMCACHE_CHUNK_SIZE=256

vllm serve "$MODEL_PATH" \
  --served-model-name glm52 --host 0.0.0.0 --port 8000 \
  --async-scheduling --load-format fastsafetensors --trust-remote-code \
  --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \
  --kv-cache-dtype fp8 \
  --tensor-parallel-size 4 \
  --max-num-batched-tokens 16384 \
  --gpu-memory-utilization 0.9 \
  --enable-prefix-caching \
  --kv-transfer-config '{"kv_connector":"LMCacheConnectorV1","kv_role":"kv_both"}' \
  --additional-config '{"online_quant_config": {"global_quant_config": "ptpc_fp8",
      "layer_quant_config":{"model.layers.*.mlp.experts":"per_block_fp8"},
      "exclude_layer": ["lm_head", "model.embed_tokens", "*.mlp.gate"]}}'

lm_eval

lm_eval --model local-completions --tasks gsm8k --num_fewshot 20 \
  --model_args model=glm52,base_url=http://127.0.0.1:8000/v1/completions,\
num_concurrent=65,max_retries=3,tokenized_requests=False,\
tokenizer=$MODEL_PATH

AgentX profile

export AIPERF_DATASET_WEKA_LIVE_ASSISTANT_RESPONSES=0
export AIPERF_DATASET_CONFIGURATION_TIMEOUT=1800
export AIPERF_SERVICE_PROFILE_CONFIGURE_TIMEOUT=1800
export AIPERF_HTTP_TCP_USER_TIMEOUT=900000

aiperf profile \
  --scenario inferencex-agentx-mvp \
  --url http://127.0.0.1:8000 --endpoint /v1/chat/completions \
  --endpoint-type chat --streaming \
  --model glm52 \
  --concurrency 16 \
  --benchmark-duration 3600 \
  --stats-interval 30 \
  --random-seed 42 \
  --failed-request-threshold 0.10 \
  --trajectory-start-min-ratio 0.25 --trajectory-start-max-ratio 0.75 \
  --warmup-requests-per-lane 10 \
  --trace-idle-gap-cap-seconds 300 \
  --warmup-grace-period 1800 \
  --use-server-token-count --no-gpu-telemetry \
  --tokenizer "$MODEL_PATH" --tokenizer-trust-remote-code \
  --max-context-length 1048576 \
  --num-dataset-entries 393 --slice-duration 1.0 \
  --public-dataset semianalysis_cc_traces_weka_062126 \
  --server-metrics http://127.0.0.1:8000/metrics \
  --output-artifact-dir "$OUT"

Test Result

lm_eval gsm8k baseline

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	20	exact_match	↑	0.9416	±	0.0065
		strict-match	20	exact_match	↑	0.9401	±	0.0065

lm_eval gsm8k with offload

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	20	exact_match	↑	0.9447	±	0.0063
		strict-match	20	exact_match	↑	0.9439	±	0.0063

AgentX baseline

Metric	avg	min	max	p99	p90	p50	std
Request Latency (ms)	214,352.42	393.17	1,663,970.55	795,741.82	352,269.77	223,567.05	172,073.12
Inter Token Latency (ms)	83.67	1.29	1,016.56	547.46	171.94	54.94	104.20
E2E Output Token Throughput (tokens/sec/user)	11.88	0.05	61.78	56.99	45.84	2.43	17.58
Output Sequence Length (tokens)	958.86	2.00	32,370.00	9,637.74	1,901.20	374.00	2,362.96
Input Sequence Length (tokens)	116,041.38	409.00	680,752.00	675,879.48	264,562.60	88,672.00	121,765.50
Output Token Throughput (tokens/sec)	79.82	N/A	N/A	N/A	N/A	N/A	N/A
Input Token Throughput (tokens/sec)	9,659.49	N/A	N/A	N/A	N/A	N/A	N/A
Request Throughput (requests/sec)	0.08	N/A	N/A	N/A	N/A	N/A	N/A
Request Count (requests)	303.00	N/A	N/A	N/A	N/A	N/A	N/A
Theoretical Prefix Cache Hit (%)	95.39	N/A	N/A	N/A	N/A	N/A	N/A

AgentX LMCache CPU offload

Metric	avg	min	max	p99	p90	p50	std
Request Latency (ms)	75,010.68	393.74	861,934.22	336,738.06	178,253.40	51,730.95	82,353.34
Inter Token Latency (ms)	29.29	0.23	131.59	80.37	37.20	27.42	11.97
E2E Output Token Throughput (tokens/sec/user)	17.44	0.01	63.43	51.11	40.46	12.73	14.92
Output Sequence Length (tokens)	1,133.77	1.00	32,370.00	10,551.30	2,674.70	395.50	2,444.85
Input Sequence Length (tokens)	132,146.58	407.00	694,151.00	678,537.05	201,050.00	108,408.50	116,739.63
Output Token Throughput (tokens/sec)	200.59	N/A	N/A	N/A	N/A	N/A	N/A
Input Token Throughput (tokens/sec)	23,379.77	N/A	N/A	N/A	N/A	N/A	N/A
Request Throughput (requests/sec)	0.18	N/A	N/A	N/A	N/A	N/A	N/A
Request Count (requests)	644.00	N/A	N/A	N/A	N/A	N/A	N/A
Theoretical Prefix Cache Hit (%)	96.61	N/A	N/A	N/A	N/A	N/A	N/A

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2142 --add-label <label>