[kliuae]

Motivation

This PR adds support for GLM-5.3-Flash in ATOM-vLLM mode, with prefix caching, LMCache KV offloading, and MTP speculative decoding. Currently only tested on single-node MI355X x8, TP8, fp8 KV cache.

Technical Details

Plugin

• Glm5NextForConditionalGeneration runs ATOM's native glm5_next model. Its KDA layers are swapped for Kimi-K3's vLLM-backed KDA layer, so vLLM manages their recurrent state, and the indexer runs its top-k through an op that reads vLLM's attention metadata.
• The pooled index keys are stored inside the MLA cache rather than in a cache of their own. vLLM gives one shared set of block ids only to layers with the same cache spec, so the keys live in one extra proxy layer that reports MLA's spec. A prefix-cache hit then brings back a block's index keys together with its MLA entries.
• A pool that is still filling up with fewer than 4 tokens has no pooled key yet, so its raw keys are kept per request in a small buffer owned by the plugin. Prefix-cache hits always end on a block boundary, which is also a pool boundary, so this buffer never needs caching.
• vLLM does not recognize glm5_next as an MLA model. It's achieved through arch-config convertor so that the cache pages are sized correctly. The sparse MLA metadata counts a pooled selection as whole pools plus the tokens of the pool still filling up.

LMCache offload

• AtomLMCacheOffloadConnector needs no GLM-specific code since it already moves the MLA cache, which now includes the index keys, and the KDA recurrent state. Restores end on block boundaries, so the per-request buffer above is never offloaded either.
• ATOMPlatform now stays active when vLLM resolves the platform during its own import. Its config hook selects a scheduler that can recover from a failed load on a model with several KV cache groups; vLLM's default scheduler stops the engine on the first one.
• Two connector fixes that apply to every plugin model: a request that finishes on the step that ran its last prefill chunk now saves that chunk, and an idle server keeps stepping until its last pending save is sent, so the blocks that save holds are released.

MTP

• Checkpoint layer 45 serves as vLLM's MTP draft (Glm5NextMTPModel) through ATOM's native NextN module, with the same plugin layers as the target. vLLM maps glm5_next to the glm5_next_mtp draft.
• With MTP, draft tokens are written to the per-request buffer before verification, and rejected ones must not overwrite keys that are still needed. The buffer becomes a ring sized to the speculative history (e.g. 8 rows for 1 draft token, 16 for 3, instead of 4), addressed by token position.
• In ATOM's native prefill and decode kernels, address the k-pool tail as a ring.

Test Plan

Model: zai-org/GLM-5.3-Flash, single node MI355X x8, TP8

Accuracy

Server command

vllm serve zai-org/GLM-5.3-Flash \
    --host localhost \
    --port 8000 \
    --tensor-parallel-size 8 \
    --trust-remote-code \
    --kv-cache-dtype fp8 \
    --max-model-len 8192 \
    --max-num-seqs 32 \
    --max-num-batched-tokens 8192 \
    --gpu-memory-utilization 0.85 \
    --enable-prefix-caching \
    --speculative-config '{"method": "mtp", "num_speculative_tokens": 3}'  # For MTP

LMCache server command

export PYTHONHASHSEED=0
export VLLM_SERVER_DEV_MODE=1
export LMCACHE_LOCAL_CPU=True
export LMCACHE_MAX_LOCAL_CPU_SIZE=96
export LMCACHE_CHUNK_SIZE=1024
export LMCACHE_CACHE_POLICY=ATOM_SLRU
export LMCACHE_TRACK_USAGE=false
export OFFLOAD_MIN_LOAD_TOKENS=256
vllm serve zai-org/GLM-5.3-Flash \
    --host localhost \
    --port 8000 \
    --tensor-parallel-size 8 \
    --trust-remote-code \
    --kv-cache-dtype fp8 \
    --max-model-len 8192 \
    --max-num-seqs 32 \
    --max-num-batched-tokens 8192 \
    --gpu-memory-utilization 0.85 \
    --enable-prefix-caching \
    --kv-transfer-config '{"kv_connector":"AtomLMCacheOffloadConnector","kv_connector_module_path":"atom.plugin.vllm.kv_transfer.connector","kv_role":"kv_both","kv_load_failure_policy":"recompute"}'

lm_eval was run twice on each LMCache configured servers. In between the two runs, the HBM prefix cache was reset (curl -X POST http://localhost:8000/reset_prefix_cache), so that the second run can only reuse prefixes by loading them from LMCache.

lm_eval command

PYTHONHASHSEED=0 \
lm_eval --model local-chat-completions \
        --model_args model=zai-org/GLM-5.3-Flash,base_url=http://localhost:8000/v1/chat/completions,num_concurrent=32,max_retries=3,tokenized_requests=False,trust_remote_code=True \
        --tasks gsm8k \
        --num_fewshot 20 \
        --gen_kwargs max_gen_toks=512 \
        --apply_chat_template \
        --fewshot_as_multiturn

Agentic workload with shared prefix was used to evaluate the effectiveness of KV cache offloading. The workload used has 128 shared prefixes of 27648 tokens, 4608 unique tokens and 512 output tokens per request, at concurrency 16, lasting for 900s.

Server command: the accuracy server commands, appended with --max-model-len 65536 --max-num-seqs 64 --max-num-batched-tokens 16384 --enable-prompt-tokens-details, in four configurations

• --no-enable-prefix-caching
• --enable-prefix-caching
• --enable-prefix-caching --num-gpu-blocks-override 2048
• LMCache server command with --num-gpu-blocks-override 2048
  (the last two with MTP 3)

Benchmark command

PYTHONHASHSEED=0 \
aiperf profile --model zai-org/GLM-5.3-Flash --url http://127.0.0.1:8000 \
    --endpoint-type chat --streaming --tokenizer zai-org/GLM-5.3-Flash \
    --tokenizer-trust-remote-code \
    --isl 4608 --isl-stddev 0 --osl 512 --osl-stddev 0 \
    --prompt-prefix-length 27648 --prompt-prefix-pool-size 128 \
    --num-dataset-entries 800 --concurrency 16 --warmup-request-count 16 \
    --benchmark-duration 900 --benchmark-grace-period 60 \
    --extra-inputs ignore_eos:true --use-server-token-count \
    --cache-bust first-turn-suffix --random-seed 530419

Test Result

lm_eval with gsm8k, 20-shot

configuration	pass	flexible-extract	strict-match	prompt tokens from HBM / LMCache
prefix caching	1	0.9682	0.9689	-
prefix caching	2	0.9689	0.9697	87.4% / -
prefix caching, MTP 3	1	0.9712	0.9697	-
prefix caching, MTP 3	2	0.9704	0.9712	57.2% / -
LMCache	1	0.9697	0.9682	-
LMCache, HBM reset	2	0.9682	0.9674	0 / 87.4%
LMCache, MTP 3	1	0.9682	0.9674	-
LMCache, MTP 3, HBM reset	2	0.9704	0.9682	0 / 57.3%

MTP 3 mean acceptance length: 3.47 - 3.48

Shared-prefix performance, MI355X TP8

configuration	Total tok/s	median TTFT (ms)	prompt tokens from HBM / LMCache
no prefix caching	37,028	2,606	-
prefix caching	76,220	448	86.3% / -
prefix caching, 2048 GPU blocks	42,470	2,532	22.7% / -
prefix caching + LMCache, 2048 GPU blocks	68,954	951	11.5% / 70.9%
prefix caching, MTP 3, 2048 GPU blocks	42,068	859	18.5% / -
prefix caching + LMCache, MTP 3, 2048 GPU blocks	72,752	465	9.1% / 69.1%

MTP 3 mean acceptance length: 2.57, with and without LMCache

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.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 2481 --add-label <label>