vllm-atom: GLM-5.3-Flash prefix caching and MTP on the plugin backend by PerryZhang01 · Pull Request #3137 · ROCm/ATOM

PerryZhang01 · 2026-11-09T15:13:18Z

Ready for review. No blocking issues. Contains #3033's commits.

#3033 serves GLM-5.3-Flash on the vLLM plugin backend but refuses prefix caching and MTP. This serves both, still without touching atom/models or atom/model_ops, gates them in one plugin PR CI cell, and moves the native GLM-5.3-Flash entry to fp8 KV.

Dev record

Found: the prefix-caching failure was aiter's bf16 MLA decode misreading KV past 4 GiB inside a layer (GLM-5.3-Flash with bf16 KV: aiter MLA decode misreads KV stored past 4 GiB of a layer's cache #3120). Prefix caching keeps cached blocks, so block ids climb past it. fp8 KV reads correctly.
Found: since [feature] support pp>1 for Kimi-K3 #2365 the shared Kimi-K3 KDA forward reads spec_state_indices_in_tensor on verify steps, which vLLM's KDA metadata lacks; the plugin's KDA adapter now defaults it to None. Kimi-K3 DSpark takes the same path; not run here.
Decided: each request's open pool lives in a vLLM-managed ring block (KpoolTailSpec, never prefix-cached) of native's depth: 4 rows, 16 for MTP3, gathered into contiguous scratch for native's kernels. Review this first.
Decided: the new cell runs GSM8K 30-shot twice on one server and gates the second pass, since one pass never hits; it fails if the server counted no hit.
Left undone: the MTP draft runs its own indexer every step.

Named result

GSM8K chat, fp8 KV, MI355X, flexible-extract. New plugin cell (TP4, prefix caching and MTP3, 30-shot, every gated request hitting one cached block), through atom_oot_test.sh: 0.9659, 0.9682, 0.9659, 0.9674, MTP acceptance 0.827 to 0.828. Native entry (TP8, 16-shot) through atom_test.sh: 0.9682 and 0.9697.

Gates

branch @ 6f091ea: plugin, GLM-5 and DSA indexer suites 628 passed, 41 failed, the same 41 as main @ cd5fa97.
@ 07bc99e: Pre Checkin green (ruff, black, non-GPU unit tests).

Evidence

Environment: rocm/atom-dev:vllm-latest (vLLM 0.31.1.dev0+gdb9527a46), aiter wheel 0.1.1.dev1+g47b4e0869, MI355X. Cell and entry rows ran the CI's own command through its harness; the MTP floors passed check_mtp_acceptance.sh.

run @ commit	GSM8K	cache hits	MTP overall	MTP per position
new cell @ `cb395c9` code, run 1	0.9659	2,701,312	0.827	0.951 / 0.842 / 0.689
new cell @ `cb395c9` code, run 2	0.9682	2,701,312	0.828	0.951 / 0.844 / 0.690
new cell @ `ba782c9`	0.9659	2,701,312	0.828	0.951 / 0.843 / 0.690
new cell @ `6f091ea`	0.9674	2,701,312	0.827	0.950 / 0.843 / 0.688
native entry, TP8 @ `effc516`, run 1	0.9682
native entry, TP8 @ `effc516`, run 2	0.9697
plain plugin cell, since dropped, @ `2784662`	0.9674

Dense path (ATOM_GLM5_KPOOL=0) @ 07bc99e, TP4, max-model-len 2048, chat 3-shot, max_gen_toks 512: 0.9697; with prefix caching 0.9659 then 0.9682 on one server, no hits (3-shot prompts average 573 tokens, under the 1920-token block). --block-size 32 with the cell's command @ 248d52d: 0.995 on the first 200 questions. MTP on the dense indexer failed at graph capture: the shared sparse indexer builder sizes its verify-row block-table buffer for max_model_len (32 kernel blocks), while vLLM rounds the dense indexer group's table to whole 8192-token manager blocks (128). 6f091ea copies only the columns the buffer holds. Dense + prefix caching + MTP3 @ 6f091ea: 0.9697 then 0.9674 on one server, acceptance 82.8%.

Before merging main, @ 2bcfa72, the cell serve command plus the flags shown, two passes on one server:

run	pass 1	pass 2	cache hits, pass 2	MTP acceptance
`--enable-prefix-caching`, 16-shot	0.9689	0.9682	2,532,480 (1319 x 1920)
`--enable-prefix-caching` + MTP3, 16-shot	0.9682	0.9697	0	82.9%
`--enable-prefix-caching` + MTP3, 30-shot	0.9704	0.9666	2,701,312 (1319 x 2048)	82.8%

Without MTP vLLM sets the attention block to 1920 tokens, with MTP3 2048, and with MTP it drops the last matched block for the draft. 16-shot prompts run 1939 to 3591 tokens, so with MTP they never hit; 30-shot prompts run 4167 to 6039. Pass 1 never hits. Pass 2 scores sit at the low end of the 16-shot runs without hits (0.9682 to 0.9742). Per question, pass 1 to pass 2 turns 6 right answers wrong and 1 wrong answer right at 30-shot, and 8 and 7 at 16-shot; pairs of runs without hits range from 6 and 1 to 2 and 7, so these runs cannot resolve a difference this small.

Merging main first failed the new cell at graph capture with 'KimiK3KDAMetadata' object has no attribute 'spec_state_indices_in_tensor' from atom/models/kimi_k3/model.py; d6af639 fixes it. A four-row verification test checks that the pooled keys and the ring come out exactly as four single-row decodes leave them.

🤖 Generated with Claude Code

Register Glm5NextForConditionalGeneration with the vLLM plugin, so GLM-5.3-Flash runs ATOM's model and kernels under vLLM 0.31 -- including the pooled (kpool) indexer -- instead of falling through to vLLM's own Glm5Next. It is the first plugin model with three kinds of KV state at once: 11 MLA layers, 34 KDA recurrent states and 11 DSA indexer caches. * KDA: reuse Kimi-K3's vLLM adapter (Glm5NextKDAAttentionVllm), adapter first in the MRO so its overrides win over GLM's native layer. * MLA width: vLLM sizes the cache at kv_lora_rank + qk_rope_head_dim = 512, because the checkpoint keeps the rope dim at 0 for its NoPE indexer, while ATOM's MLA entry carries the 64-wide zero rope pad. Patch get_head_size to 576 for glm5_next_text only (glm5_next_mla_head_size.py). vLLM's own config classes are used as is; they carry every checkpoint field ATOM reads. * Indexer row width 132 -> 144 (aligned_index_cache_row), so the indexer page divides the MLA page and vLLM scales its block instead of padding it. * kpool: the indexer spec declares tokens_per_state = index_kpool. The sparse backends keep their 64-token kernel block, so a block holds 16 pooled rows, the preshuffled layout the decode scorer computes correctly in (checked in the indexer builder). The pooled op (glm5_kpool_indexer.py) drives the native kpool kernels from vLLM's metadata. Its per-slot tail buffer is allocated when the KDA state is bound, before warmup and graph capture: a buffer first created inside a capture has its zero-fill replayed every step. * The pure-decode sparse-MLA metadata cache keys on the pooled sparse length (pools * kpool + seq_len % kpool). min(seq_len, topk) stops changing past topk while the open tail keeps moving, so a hit would replay a stale split. * PCP, DCP, speculative decoding, prefix caching and KV transfer are refused. GSM8K chat 16-shot (every prompt past index_topk, so pooled selection decides), TP4, bf16 KV, MI355X, vLLM 0.31.1.dev0+gdb9527a46, aiter 47b4e0869, flexible-extract: 0.9697 and 0.9666 with the plugin CI serve flags, 0.9704 with vLLM's defaults; native ATOM 0.9666. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…trix TP4, chat 16-shot gsm8k so the pooled indexer is what selects, bf16 KV. Gate 0.94, the same as the native GLM-5.3-Flash-kpool-16shot entry; this exact serve command measured 0.9697 and 0.9666 flexible-extract. The evaluation takes about four minutes once the server is up. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

… matrix The plugin PR matrix keeps six cells, with GLM-5.3-Flash added. GLM-5.2-MXFP4 with MTP stays covered by the nightly vLLM accuracy run (oot_models_accuracy.json, "GLM-5.2-MXFP4 MTP TP4"). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Move the indexer hooks added to atom/models/glm5_next.py into a plugin subclass, Glm5NextIndexerVllm, swapped in while the model is built the same way as the KDA layer, and declare the indexer row width with native aiter_mla.aligned_index_cache_dim instead of a new geometry helper. atom/models and atom/model_ops are now unchanged by this PR. GSM8K chat 16-shot through the plugin CI harness on this tree: 0.9651 flexible-extract (gate 0.94); earlier runs of the same computation 0.9697, 0.9666, 0.9727. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* Leave the shared sparse-MLA builders reading index_topk and num_hidden_layers from hf_config, as on main: vLLM 0.31's Glm5NextConfig carries both at the top level, so GLM-5.3-Flash needs no change there. * Drop the indexer builder's ChunkK=256 block-size check. It assumed the kernel block equals the manager block, which stopped holding when the sparse backends kept their 64-token kernel block; the builders get the kernel block, so it could only reject a working configuration. * Clear the pooled tail buffer in place when the KDA state is re-bound with the same shape, so CUDA graphs captured against it keep a valid pointer. * Refuse, instead of silently re-viewing, a pooled index cache whose row count disagrees with the builder's rows_per_block. * Fix the dense indexer path (ATOM_GLM5_KPOOL=0): the shared indexer's nope-cache fusion hands the op a BF16 query, which the plugin's dense op cannot take (fp8_mqa_logits fails to compile). Glm5NextIndexerVllm turns the fusion off, so the query is quantized first, as for GLM-5.2. * List GLM-5.3-Flash in the vLLM plugin backend guide. On this tree, TP4, MI355X: the plugin CI harness gives 0.9727 GSM8K chat 16-shot (gate 0.94); the dense path gives 0.9666 3-shot at max-model-len 2048 (native ATOM 0.9689). Kpool geometry checked under ATOM_GLM5_KPOOL_CHECK over a full 16-shot run: 16 rows per block, state slots and block ids in range. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* Unit tests for the pooled indexer's prefill-chunk metadata: whole prompts shorter than one pool, a prompt of exactly one pool, chunks that resume an earlier chunk and end mid-pool, and a chunk covering a later request range. Only complete pools are counted, over each request's context so far. * Under ATOM_GLM5_KPOOL_CHECK=1, log per prefill chunk how many requests add fewer rows than one pool and how many resume mid-pool, so a run can show those cases were exercised. TP4, MI355X, --max-num-batched-tokens 1027 so chunk boundaries fall mid-pool: GSM8K chat 16-shot 0.9682. During the run, mid-pool resumes were logged in about 1,700 scheduler steps and sub-pool chunks in 2 (one of them also mid-pool), with no check failing. Prompts of 1 to 7 tokens continue coherently. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

A prefill chunk whose requests have not closed a single pool yet (total_pools == 0, e.g. a 1 to 3 token prompt with index_kpool=4) skipped expand_pools_and_append_tail along with the scoring, so every row kept topk = -1 and attended slot 0 instead of its own tokens. The chunk now builds an all -1 pool_topk and still appends the tail, so each row attends positions 0..pos. Review follow-ups on the same op and its metadata: - Size scoring row chunks with the shared sparse_indexer_row_chunk helper instead of a copy that dropped its 2 GiB bound. - Drop the ATOM_GLM5_KPOOL_CHECK debug block. - Drop the pooled slot_mapping recompute in the indexer metadata builder; the pooled op computes its own write slots. - Return instead of raising when vLLM has no Glm5NextTextConfig or no get_head_size hook; the head-size patch runs for every model. - Restate why _check_unsupported exists: plugin mode builds an ATOM Config, but vLLM's PCP and DCP sizes are never forwarded into it. Adds a GPU test that runs the prefill chunk op with total_pools == 0 and checks each row attends exactly its own positions. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Drop the cell's bfloat16 override so it runs the harness's fp8 KV cache: GSM8K chat 16-shot scores 0.9727 / 0.9682 over two runs of the cell's serve command. Correct the prompt-length note: the 16-shot prompts run 1939-3591 tokens, all but one past index_topk=2048. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The pooled indexer kept each request's open pool (the tail) in a side buffer keyed by the KDA state slot. That ties the tail to a slot an MTP draft layer does not have (the draft layer is MLA plus an indexer, no KDA), and that vLLM moves under prefix caching. Declare the tail as vLLM's own GLM-5.3-Flash tail cache instead (KpoolTailSpec: one block per request for its lifetime, never prefix-cached), with native ATOM's ring depth so native's verification kernels stay correct: one pool without speculation, speculative_kpool_history_size with it (16 rows for MTP3, where vLLM sizes 8). vLLM aliases each ring into the indexer page, so the block stride is padded; the op gathers the batch's rings into a contiguous buffer for the native kernels and scatters them back. The heterogeneous MLA/MHA EAGLE3 KV pool check treated the ring spec as an MHA draft layer; skip it there. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Register the checkpoint's MTP layer (Glm5NextMTPModel) as ATOM's Glm5NextMTP with the plugin indexer swapped in during construction, under the DeepSeek-MTP draft contract (position-0 embedding mask, spec_step_idx, (sample, recycle) hidden). vLLM passes the draft the target's positions, which is what native's MTP indexer derives with positions - 1, so the plugin indexer takes them as-is. The pooled op verifies drafts with native ATOM's kernels: build_speculative_pool_candidates forms the pool ending at each row from the step's rows and the ring, update_speculative_kpool_history writes the rows into the ring, and scoring reuses the per-row decode path the builder already flattens a verification into. A test checks that a four-row verification leaves the pooled keys and the ring exactly as four single-row decodes do. The KDA state shape takes vLLM's num_speculative_tokens, as Kimi-K3's does. Speculative methods other than mtp are refused. Prefix caching stays refused on the pooled path here; the next commit lifts it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The refusal blamed align-mode KDA state migration. The real cause is aiter's bf16 MLA decode misreading KV past 4 GiB inside a layer (#3120): prefix caching only made vLLM hand out higher block ids. With fp8 KV, prefix caching scores 0.9727 then 0.9682 (cache hits on the second pass) on GSM8K 16-shot, so drop the refusal. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

ATOM's shared Kimi-K3 KDA forward now also reads spec_state_indices_in_tensor on a verify step (#2365, pp>1 checkpoint forks). vLLM's KimiK3KDAMetadata has no such field, so any plugin model on that forward with speculative decoding (GLM-5.3-Flash MTP, Kimi-K3 DSpark) raised AttributeError during CUDA graph capture. Default it to None, as the non-speculative index already is: vLLM has no state fork, and None keeps the read on the write slot. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Add a TP4 cell serving GLM-5.3-Flash on fp8 KV with prefix caching and MTP3. lm_eval draws each question's exemplars separately, so nothing hits within one pass, and with MTP3 a prompt hits only past two 2048-token blocks: the client runs GSM8K 30-shot (4167-6039 token prompts) twice on one server, gates the second pass, and fails if the server counted no prefix-cache hit. Two runs through atom_oot_test.sh: 0.9659 / 0.9682, MTP acceptance 0.827 / 0.828, every second-pass request hitting one cached block. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

PerryZhang01 · 2026-11-10T06:54:26Z

Round 1 pushed cb395c9; the base is now main.

Merged main. Since [feature] support pp>1 for Kimi-K3 #2365, speculative decoding on the shared Kimi-K3 KDA forward failed at graph capture on the plugin path; d6af639 defaults the missing read index.
New PR CI cell "GLM-5.3-Flash kpool MTP prefix-cache TP4": 0.9659 and 0.9682 through atom_oot_test.sh, MTP floors 0.58 overall and 0.66/0.59/0.48 per position.

…R CI The plain GLM-5.3-Flash cell runs a subset of what the MTP prefix-cache cell runs, so drop it and gate the model with the one cell. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

aiter's bf16 MLA decode misreads KV past 4 GiB inside a layer (#3120), which a long-running server reaches; fp8 reads correctly. Switch the entry's KV cache to fp8. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

valar-cop

some changes needed..

gsm8k for the dense path (ATOM_GLM5_KPOOL=0)?

ci Check Code Style with Ruff failed — looks like the pr itself (high): reviewdog's diff-context lint reports a pre-existing BLE001 on spec_decode_patch.py:279 because this PR added lines to that hunk. The captured log is 0 bytes, so the root cause comes from the reviewdog annotation, not the log.

⚠️ [ruff] BLE001 Do not catch blind exception: Exception — atom/plugin/vllm/spec_decode_patch.py:279
next: Narrow the except to except ImportError: on line 279. https://github.com/ROCm/ATOM/actions/runs/42775509059/job/129692968984

ci Run non-GPU unit tests failed — not sure yet (low): The captured log is empty (0 bytes), so the failure cannot be attributed. The job ignores tests/plugin and the new diff touches only plugin modules, but that is inference, not evidence.

job_129692969370.log is 0 bytes
next: Re-run the job or attach the log before attributing it. https://github.com/ROCm/ATOM/actions/runs/42775509059/job/129692969370

ci Atomesh Accuracy Validation failed — not sure yet (low): The run is marked failed but no failed job or log was captured, so nothing can be attributed.

Atomesh Accuracy Validation: failure (no failed job listed, no log)
next: Open the run and identify the failing job before treating it as PR-caused. https://github.com/ROCm/ATOM/actions/runs/42775508658

valar-cop · 2026-11-10T07:53:26Z

@@ -278,12 +278,20 @@ def _spec_has_heterogeneous_mla_mha_backend(kv_cache_spec) -> bool:
        from vllm.v1.kv_cache_interface import AttentionSpec, MLAAttentionSpec
    except Exception:


narrow this? except ImportError:

Narrowed to except ImportError in ba782c9.

valar-cop · 2026-11-10T07:53:26Z

+    index_kpool = _pool_size(vllm_config)
+    if pooled_path_enabled(index_kpool):
+        block_size = vllm_config.cache_config.block_size
+        if block_size % index_kpool != 0:


block_size 16/32 passes here then blows up later, use glm5_kpool_block_size(index_kpool) so we fail at config time?

Not changed: the default launch reaches this check with 16. vLLM raises the attention block to a multiple of the 64-token kernel block (2048 in the CI cell) only after the model loads, so a glm5_kpool_block_size check here would refuse the default launch. --block-size 32 serves: 0.995 on 200 questions, CI cell command @ 248d52d.

valar-cop · 2026-11-10T07:53:26Z

+        triton_convert_req_index_to_global_index,
+    )
+
+    result = weights.to(dtype=torch.float32, copy=True)


why this copy? return None like sparse_attn_indexer_plugin_mode

Returns nothing now, with the dense op's fake impl; the caller never read the copy (ba782c9). The CI cell re-ran at 0.9659 after the change.

valar-cop · 2026-11-10T07:53:26Z

            # runs; these floors sit near 70% of that.
            mtp_per_pos_threshold: "0.55,0.40,0.25,0.18,0.10"
            runner: linux-atom-pitt-mi350x-8
+          - display_name: "GLM-5.3-Flash kpool MTP prefix-cache TP4"


why drop the V2 model runner MTP cell instead of keeping it? that was our only gate

Restored in ba782c9.

- Register Glm5NextMTPModel's predictor (DeepSeekMultiTokenPredictor, which Glm5NextMTP inherits) in the MTP input-mask contract test. - Narrow the AttentionSpec import guard to ImportError. - The kpool indexer op returns nothing, like the dense one: its caller never read the float32 copy of the weights it returned. - Restore the GLM-5.2-MXFP4-MTP cell: it is the plugin CI's only Model Runner V2 MTP gate. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Two runs of the entry's exact command on fp8 KV (TP8): flexible-extract 0.9682 and 0.9697. Baseline 0.9682, threshold unchanged at 0.94. Correct the note's prompt lengths: the 16-shot prompts run 1939-3591 tokens, all but one past index_topk=2048. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

valar-cop

LGTM

ci Atomesh Accuracy Validation failed — not sure yet (low): The run is marked failed but no failed job or log was captured, so nothing can be attributed. The workflow triggers on this PR's models_accuracy.json edit, and its own comment says its PR whitelist selects no model.

Atomesh Accuracy Validation: failure (no failed job listed, no log)
next: Open the run (or attach the job log) and identify the failing job before treating it as PR-caused. https://github.com/ROCm/ATOM/actions/runs/42790533050

With ATOM_GLM5_KPOOL=0 and MTP, the dense indexer's metadata builder flattens the verify rows into a block-table buffer sized from the indexer's 64-token kernel block (32 columns at max-model-len 2048), while the hybrid model's block table has 128, and the engine fails at graph capture. Refuse the combination at startup. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

valar-cop

LGTM

ci Atomesh Accuracy Validation failed — not sure yet (low): Run is marked failed but no failed job or log was captured, so nothing can be attributed. In PR mode this workflow's whitelist (DeepSeek-R1-0528=main, gpt-oss-120b=nightly) selects no model, so its test matrix is empty; this PR only edits models_accuracy.json, which is one of the workflow's path triggers.

Atomesh Accuracy Validation: failure (no failed job listed, no log)
.github/workflows/atomesh-accuracy-validation.yaml: # NOTE: ... atomesh standalone currently validates nothing on a pull request
next: Open the run and attach the failing job's log before treating it as PR-caused; if it is check-signal/download_aiter_wheel or the empty matrix, it is not this PR's code. Same non-attributable run as in the previous two rounds. https://github.com/ROCm/ATOM/actions/runs/42794752659

PerryZhang01 · 2026-11-10T10:28:40Z

Round 2 pushed 07bc99e; Pre Checkin is green.

Review fixes in ba782c9, replies inline; the GLM-5.2-MXFP4-MTP cell is back.
Dense path (ATOM_GLM5_KPOOL=0), fp8, TP4, 2048 context, chat 3-shot: 0.9697; with prefix caching 0.9659 then 0.9682, no hits (prompts under one 1920-token block).
MTP on the dense indexer failed at graph capture (block-table buffer 32 columns, table 128); 07bc99e refuses it at startup.
Native entry on fp8, TP8: 0.9682 and 0.9697; baseline 0.9682.

MTP verify rows on the sparse MLA indexer are flattened into a block-table buffer sized for max_model_len in kernel blocks. vLLM rounds each request's table up to whole manager blocks, and on a hybrid model the indexer group's manager block can be several times the attention one: GLM-5.3-Flash's dense indexer (ATOM_GLM5_KPOOL=0) gets 8192-token manager blocks, 128 kernel blocks per row at max_model_len 2048 against a 32-column buffer, and the engine failed at graph capture. Copy only the columns the buffer holds; the rest address no position below max_model_len. This replaces the startup refusal from 07bc99e. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

PerryZhang01 · 2026-11-10T12:29:24Z

Round 3 pushed 6f091ea: a fix replaces 07bc99e's refusal of MTP on the dense indexer.

Cause: the shared sparse indexer builder sizes the verify-row block-table buffer for max_model_len (32 kernel blocks at 2048); vLLM rounds the dense indexer group's table to whole 8192-token manager blocks (128). The expansion now copies only the columns the buffer holds; the rest address nothing below max_model_len.
Dense + prefix caching + MTP3, 3-shot, 2048: 0.9697 then 0.9674, acceptance 82.8%. New cell 0.9674.

#2365 made the shared K3 KDA forward read `kda_metadata.spec_state_indices_in_tensor` on its speculative-decode branch. vLLM's KimiK3KDAMetadata has no such field, and the plugin only defaulted its non-spec twin, so K3 on the plugin dies in profiling under DSpark (AttributeError on the first verify-shaped forward). None keeps the forward on the write slot, as the non-spec default does. The same hunk as #3137's, so the two apply cleanly in either order. The K3 TP8 LMCache MP runs on this PR ran with it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

… object (#3117) * Kimi-K3 plugin: ride the KDA state on the PAGE object through the LMCache MP server Kimi-K3 on the MP backend was refused ("recurrent (mamba) state offloads in-process only"). It now takes the path DeepSeek-V4.1's checkpoints take (#2598): the KDA state at a boundary rides the PAGE object as a native-state image, so KV and state commit and restore as one LMCache object. No copy is needed. vLLM lays every KV cache group out layer-outer from byte 0 of one allocation, so mamba group layer j aliases MLA plane j. On K3 the three mamba groups of 23 KDA layers sit on the first 23 of the 29 MLA planes, and the state a group holds in block b is PAGE unit b of those planes. A boundary's image is one unit per mamba group (3 x 23 x 884,736 B). - kda_ride.py: KdaRidePlanner (caps hits to the chunk grid, names the load destination row (n - 1) // block in every mamba group or none), kda_ride_geometry (the checkpoint spec both halves derive from vLLM's KV cache config), publish_kda_ride_state (checks on the registered tensors that every KDA layer starts at, and strides with, one of the leading PAGE planes before declaring the image), and the startup refusals (align mode, prefix caching, kv_load_failure_policy=recompute). - V41StateMPScheduler had nothing V4.1-specific; renamed StateRideMPScheduler (ride_mp_scheduler.py) and shared by both models. _v41_request_of is now _request_of. - The KDA state is per rank (TP-sharded heads); the in-process K3 leg is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Kimi-K3 plugin: size the KDA image by attention layer, and ride only K3 Review of #3117: - The image covered the leading N PAGE planes with N = mamba layers per group, which counts layers, not planes. An attention layer with a DSA index cache publishes two planes, so the count fell short. The spec now sums each leading attention layer's movable bytes from vLLM's specs (an index cache folded into the layer that owns it, as build_kv_cache_tensors publishes it), and the worker counts the region planes off the published regions' layer names before it checks the aliasing. - Only Kimi-K3 rides on MP. Its KDA layers report GDN_ATTN, as Qwen's GDN layers do, so the gate is the kimi_k3 offload layout; any other mamba model on MP is refused as before. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Kimi-K3 plugin: put the layers the KDA state aliases first in the PAGE order The first K3 TP8 boot over mp refused at registration: mamba group 0 layer 18 started on PAGE plane 23, outside the image's leading 23 planes. ATOM publishes PAGE planes by layer name, while vLLM aliases mamba layer j onto the j-th attention layer it allocates, and with DSpark the two orders differ (29 attention layers, five sort ahead of the first aliased one). kda_ride_page_order moves the attention layers the mamba layers alias to the front, in mamba-layer order, matched on addresses, and the scheduler sizes the image from vLLM's allocation order. The end-to-end test gains a layer that sorts first without being aliased; without the reorder it raises the same refusal the boot did. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Kimi-K3 plugin: default spec_state_indices_in_tensor so DSpark starts #2365 made the shared K3 KDA forward read `kda_metadata.spec_state_indices_in_tensor` on its speculative-decode branch. vLLM's KimiK3KDAMetadata has no such field, and the plugin only defaulted its non-spec twin, so K3 on the plugin dies in profiling under DSpark (AttributeError on the first verify-shaped forward). None keeps the forward on the write slot, as the non-spec default does. The same hunk as #3137's, so the two apply cleanly in either order. The K3 TP8 LMCache MP runs on this PR ran with it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: perzhang <perzhang@amd.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

zgplvyou and others added 11 commits November 5, 2026 01:27

PerryZhang01 mentioned this pull request Nov 9, 2026

vllm-atom: serve GLM-5.3-Flash on the plugin backend and gate it in plugin PR CI #3033

Closed

zgplvyou and others added 3 commits November 10, 2026 00:01

Merge branch 'main' into perzhang/glm53-flash-plugin-prefix-mtp

2784662

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

PerryZhang01 changed the base branch from perzhang/glm53-flash-plugin-ci to main November 10, 2026 06:53

zgplvyou and others added 2 commits November 10, 2026 01:28

ci: keep only the GLM-5.3-Flash MTP prefix-cache cell in the plugin P…

b0cc701

…R CI The plain GLM-5.3-Flash cell runs a subset of what the MTP prefix-cache cell runs, so drop it and gate the model with the one cell. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

ci: run the native GLM-5.3-Flash kpool entry on fp8 KV

effc516

aiter's bf16 MLA decode misreads KV past 4 GiB inside a layer (#3120), which a long-running server reaches; fp8 reads correctly. Switch the entry's KV cache to fp8. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

valar-cop reviewed Nov 10, 2026

View reviewed changes

zgplvyou and others added 2 commits November 10, 2026 03:12

valar-cop reviewed Nov 10, 2026

View reviewed changes

PerryZhang01 added the ci:vllm label Nov 10, 2026

PerryZhang01 mentioned this pull request Nov 10, 2026

Kimi-K3 plugin: LMCache MP offload with the KDA state riding the PAGE object #3117

Merged

zufayu requested a review from valarLip November 11, 2026 00:51

zufayu assigned valarLip Nov 11, 2026

PerryZhang01 added the ci:atom label Nov 11, 2026

valarLip approved these changes Nov 11, 2026

View reviewed changes

valarLip merged commit e91df95 into main Nov 11, 2026
67 of 80 checks passed

valarLip deleted the perzhang/glm53-flash-plugin-prefix-mtp branch November 11, 2026 03:16

ZLkanyo009 mentioned this pull request Nov 11, 2026

fix(offload): key multimodal prompts by media, one LMCache engine id per server; a native MP restore has its own deadline and waits for a free slot before dispatch #3027

Open

MengqingCao mentioned this pull request Nov 11, 2026

fix(runner): reconcile KV block counts across TP and PCP-group workers #3171

Merged

3 of 5 tasks

wanzhenchn mentioned this pull request Nov 12, 2026

ci(mesh): skip the mocker benchmark and empty accuracy runs on PRs, fix the mocker cargo cache #3232

Merged

		@@ -278,12 +278,20 @@ def _spec_has_heterogeneous_mla_mha_backend(kv_cache_spec) -> bool:
		from vllm.v1.kv_cache_interface import AttentionSpec, MLAAttentionSpec
		except Exception:

Conversation

PerryZhang01 commented Nov 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Dev record

Named result

Gates

Uh oh!

PerryZhang01 commented Nov 10, 2026

Uh oh!

valar-cop left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

valar-cop left a comment

Choose a reason for hiding this comment

Uh oh!

valar-cop left a comment

Choose a reason for hiding this comment

Uh oh!

PerryZhang01 commented Nov 10, 2026

Uh oh!

PerryZhang01 commented Nov 10, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

PerryZhang01 commented Nov 9, 2026 •

edited

Loading