vllm-atom: serve GLM-5.3-Flash on the plugin backend and gate it in plugin PR CI by PerryZhang01 · Pull Request #3033 · ROCm/ATOM

PerryZhang01 · 2026-11-05T07:40:35Z

Ready for review. No blocking issues. The new cell needs zai-org/GLM-5.3-Flash in the PITT runners' model cache (306 GB).

GLM-5.3-Flash has a native PR gate but no vLLM plugin support: the plugin falls through to vLLM's own model, so no plugin cell tests ATOM's kernels for it. This adds the plugin model, including the pooled (kpool) indexer, without touching atom/models or atom/model_ops, and a PR CI cell. It drops the GLM-5.2-MXFP4-MTP cell; nightly still runs that model, on Model Runner V1 only, so no PR gate runs GlmMoeDsa on V2.

Dev record

Found: Enable GLM-5.3-Flash under vllm serve on the ATOM plugin.- #2299 #2340 and [ATOM-vLLM] Add support for GLM-5.3-Flash with LMCache offload #2481 also add GLM-5.3-Flash to the plugin; both predate vLLM 0.31 (plugin/vllm: upgrade the ATOM vLLM plugin to 0.31.0 #2236) and add no CI cell.
Found: vLLM 0.31 replaced MLAAttentionSpec.compress_ratio with tokens_per_state. The pooled indexer declares it and keeps the sparse backends' 64-token kernel block.
Found: main's shared indexer hands a NoPE indexer's BF16 query to the plugin's FP8-only dense op; Glm5NextIndexerVllm turns that fusion off.
Decided: the pooled tail buffer is allocated at KDA state bind, before graph capture; a prefill chunk where no request has closed a pool still appends the tail.
Decided: the cell runs fp8 KV; bf16 misreads KV past 4 GiB in a layer (GLM-5.3-Flash with bf16 KV: aiter MLA decode misreads KV stored past 4 GiB of a layer's cache #3120).
Left undone: MTP, DCP, PCP, prefix caching and KV transfer are refused; vllm-atom: GLM-5.3-Flash prefix caching and MTP on the plugin backend #3137 adds prefix caching and MTP.

Named result

GSM8K chat 16-shot, MI355X, flexible-extract. The cell (TP4, fp8 KV): 0.9727 and 0.9682 over two runs; native ATOM (bf16 KV) 0.9666. Cell gate 0.94.

Gates

branch @ eb07d7e: tests/plugin/test_vllm_glm5_next.py 21 passed; plugin, GLM-5 and DSA indexer tests fail main's 71 (@ df61c5d), none new.
black clean. ruff: no new findings in the changed files.
@ 4c33493: the cell through atom_oot_test.sh accuracy ci, 0.9742, gate passed.

Evidence

Environment: rocm/atom-dev:vllm-latest @ sha256:dc7381266f5e (built 2026-11-04), vLLM 0.31.1.dev0+gdb9527a46, aiter wheel 0.1.1.dev1+g47b4e0869 (the CI manifest's), MI355X.

run	flexible	strict
fp8 KV, TP4, CI serve command @ `eb07d7e`, run 1	0.9727	0.9735
fp8 KV, TP4, CI serve command @ `eb07d7e`, run 2	0.9682	0.9689
fp8 KV, TP4, `atom_oot_test.sh accuracy ci` @ `4c33493`	0.9742	0.9750
bf16 KV, TP4, CI serve command, run 1	0.9697	0.9682
TP4, CI serve command, run 2	0.9666	0.9674
TP4, `atom_oot_test.sh accuracy ci`, run 3	0.9727	0.9735
TP4, `atom_oot_test.sh accuracy ci`, run 4	0.9651	0.9651
TP4, CI serve command with a geometry check, run 5	0.9682	0.9689
TP4, `atom_oot_test.sh accuracy ci` @ `a90e5a7`, run 6	0.9727	0.9719
TP4, `atom_oot_test.sh accuracy ci` @ `eb07d7e`, run 7	0.9674	0.9666
TP8, `atom_oot_test.sh accuracy ci` @ `36593ec`	0.9712	0.9704
TP4, `--max-num-batched-tokens 1027`	0.9682	0.9689
TP4, dense path (`ATOM_GLM5_KPOOL=0`), 3-shot, max-model-len 2048	0.9666	0.9659

Rows below the fp8 ones are bf16 KV, the cell's KV until 4c33493, which changes only the cell and the docs. Runs 1 to 5 and the 1027-token run used earlier commits of this branch with the same kpool computation. The eb07d7e fix only changes chunks in which no request has closed a pool, which GSM8K prompts never produce, so it does not bear on runs 1 to 6 or TP8. Run 5 checked the pooled geometry on every indexer call: 16 rows per block, state slots and block ids in range, and the cache view matched the metadata. 16-shot prompts run 1939 to 3591 tokens (mean 2728), all but one past index_topk=2048. TP8 sets the attention block to 512 tokens (TP4: 1024). With 1027-token chunks, 1684 steps resumed a request mid-pool and 3 steps carried a request with fewer than 4 new rows.

Short prompts, one request at a time, greedy, two server boots per tree, against the dense path (ATOM_GLM5_KPOOL=0, max-model-len 2048) on prompts of 1 to 7 tokens. Boots of the same tree differ by up to 0.57 in first-token logprob, so smaller gaps say nothing. On a 3-token prompt, 36593ec scores -3.69 and -3.75 against dense, with a different continuation; eb07d7e scores +0.03 and +0.02. The other 2 and 3 token prompts stay within the noise on both trees.

The 71 test failures are the same list on both trees: offload 37, DeepSeek-V4/V4.1 20, Kimi-K3 11, MHA and KV layout 3. run_unit_tests.sh skips tests/plugin, so CI does not run these suites.

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

github-actions · 2026-11-05T07:40:50Z

🏷️ 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 3033 --add-label <label>

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

PerryZhang01 · 2026-11-05T07:49:41Z

Round 2 pushed 71f82f8. Drops the GLM-5.2-MXFP4-MTP cell from the plugin PR matrix, which now has six cells; no code change.

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>

PerryZhang01 · 2026-11-05T09:11:15Z

Round 3 pushed 217d896. The indexer hooks move from atom/models/glm5_next.py into a plugin subclass, so the PR no longer changes atom/models or atom/model_ops. Same computation; the CI harness on this tree gives 0.9651, gate 0.94.

valar-cop

some changes needed..

tp8 evidence for the plugin backend? and any test for prefill chunks below index_kpool?

ci green.

valar-cop · 2026-11-05T12:49:56Z

+                    "pooled prefill metadata is missing its pooled cumulative "
+                    "lengths; the builder did not see tokens_per_state > 1"
+                )
+            if total_pools == 0:


total_pools == 0 leaves topk at -1, so the query attends slot 0 only. still call expand_pools_and_append_tail with an all -1 pool_topk?

Fixed in eb07d7e; details in the newer thread on this line.

valar-cop · 2026-11-05T12:49:56Z

+            )
+
+            budget_bytes = envs.ATOM_SPARSE_INDEXER_LOGITS_BUDGET_MB * 1024 * 1024
+            if budget_bytes > 0 and budget_bytes // (total_pools * 4) < n_rows:


this is a copy of sparse_indexer_row_chunk and drops the < 2 GiB bound. can we just call the helper instead?

Done in eb07d7e: the chunk loop calls sparse_indexer_row_chunk with envs.ATOM_SPARSE_INDEXER_LOGITS_BUDGET_MB.

valar-cop · 2026-11-05T12:49:56Z

+    has_decode = indexer_meta.num_decodes > 0
+    has_prefill = indexer_meta.num_prefills > 0
+
+    if os.getenv("ATOM_GLM5_KPOOL_CHECK") == "1" and not (


os.getenv here? put it into envs.ATOM_xxx like the others, or drop the debug block.

Dropped in eb07d7e.

valar-cop · 2026-11-05T12:49:56Z

+def _check_unsupported(vllm_config: VllmConfig) -> None:
+    """Re-raise the guards `Config.__post_init__` applies natively.
+
+    Plugin mode builds a `PluginConfig` (atom/plugin/config.py) and never runs


the docstring premise is wrong, plugin builds a real Config. the real reason is pcp/dcp never forwarded, please say that.

Rewritten in eb07d7e: the plugin builds an ATOM Config, but vLLM's PCP and DCP sizes are never forwarded into it, so the native guard always sees 1.

valar-cop · 2026-11-05T12:49:57Z

+                get_compressed_slot_mapping,
+            )
+
+            slot_mapping = get_compressed_slot_mapping(


the pooled op never reads slot_mapping, so this per-step recompute looks dead. remove it?

Removed in eb07d7e; the builder passes common_attn_metadata.slot_mapping as main does.

valar-cop · 2026-11-05T12:49:57Z

-            lm_eval_num_fewshot: 20
-            accuracy_test_threshold: 0.91
-            runner: linux-atom-pitt-mi350x-8
          - display_name: "MiniMax-M3-MXFP4 EAGLE3 TP4"


this drops the only V2_MODEL_RUNNER=1 coverage for GlmMoeDsa. move the flag to the nightly cell or say V2 is no longer covered.

Stated in the PR body: the nightly GLM-5.2-MXFP4 MTP TP4 cell runs on V1, so the PR gate no longer runs GlmMoeDsa on V2. DeepSeek-V4-Flash keeps the gate's V2 coverage.

valar-cop · 2026-11-05T12:49:57Z

+            exc,
+        )
+
+    from vllm.transformers_utils.configs.glm5_next import Glm5NextTextConfig


unguarded import here breaks plugin load on older vllm. wrap in try/except ImportError and return?

Done in eb07d7e: an ImportError returns, and the patch also returns when the convertor has no get_head_size.

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

PerryZhang01 · 2026-11-05T15:32:36Z

Round 4 pushed a90e5a7. Review fixes: the shared sparse-MLA builders read hf_config as on main; an obsolete block-size check is dropped; the pooled tail buffer keeps its address on re-bind; a mismatched cache view now raises; the dense path (ATOM_GLM5_KPOOL=0) no longer passes a BF16 query to the plugin dense op. The harness gives 0.9727 on this tree.

valar-cop

one correctness issue still open in the pooled prefill path..

ci green.

valar-cop · 2026-11-06T01:41:20Z

+                    "pooled prefill metadata is missing its pooled cumulative "
+                    "lengths; the builder did not see tokens_per_state > 1"
+                )
+            if total_pools == 0:


short prompts (n < index_kpool) keep -1 here, then _convert_req_index_to_global_index_kernel maps tok<0 to slot 0, so they attend slot 0 instead of their own tokens. plugin backend is always sparse so this is silent wrong output. can we still expand the tail when total_pools == 0?

Fixed in eb07d7e; details in the newer thread on this line.

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

valar-cop

some changes needed..

ci green.

valar-cop · 2026-11-06T05:15:19Z

+                    "pooled prefill metadata is missing its pooled cumulative "
+                    "lengths; the builder did not see tokens_per_state > 1"
+                )
+            if total_pools == 0:


still open: total_pools == 0 skips the tail expansion too, so short prompts attend slot 0. call expand_pools_and_append_tail with an all -1 pool_topk before continue, and add a test on the op path not just the metadata?

Fixed in eb07d7e. A chunk with total_pools == 0 now builds an all -1 pool_topk and still calls expand_pools_and_append_tail, so the row at position p attends 0..p. test_prefill_rows_shorter_than_a_pool_attend_their_own_tokens runs that op path on GPU.

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>

PerryZhang01 · 2026-11-06T06:31:50Z

Round 5 pushed eb07d7e. A prefill chunk in which no request has closed a pool now appends the tail, so 1 to 3 token prompts attend their own tokens; an op-path GPU test covers it, and 36593ec adds sub-pool metadata tests. TP8 harness @ 36593ec: 0.9712. TP4 harness @ eb07d7e: 0.9674. Inline threads answered; evidence in the PR body.

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>

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

Round 6 pushed 4c33493.

The GLM-5.3-Flash cell drops its bfloat16 override and runs the harness's fp8 KV: 0.9727 and 0.9682 over two runs of its serve command, 0.9742 through atom_oot_test.sh accuracy ci. bf16 misreads KV past 4 GiB in a layer (GLM-5.3-Flash with bf16 KV: aiter MLA decode misreads KV stored past 4 GiB of a layer's cache #3120).
The prompt-length note now reads 1939 to 3591 tokens, all but one past index_topk=2048; the old lower bound (2763) was wrong.
Prefix caching and MTP for this model are in vllm-atom: GLM-5.3-Flash prefix caching and MTP on the plugin backend #3137, stacked on this branch.

PerryZhang01 · 2026-11-10T07:17:10Z

Closed in favor of #3137, which contains these commits and targets main.

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

PerryZhang01 added the ci:vllm label Nov 5, 2026

valar-cop reviewed Nov 5, 2026

View reviewed changes

zufayu requested a review from valarLip November 6, 2026 00:50

zufayu assigned valarLip Nov 6, 2026

valar-cop reviewed Nov 6, 2026

View reviewed changes

PerryZhang01 mentioned this pull request Nov 9, 2026

GLM-5.3-Flash with bf16 KV: aiter MLA decode misreads KV stored past 4 GiB of a layer's cache #3120

Open

PerryZhang01 mentioned this pull request Nov 9, 2026

vllm-atom: GLM-5.3-Flash prefix caching and MTP on the plugin backend #3137

Merged

PerryZhang01 closed this Nov 10, 2026

Conversation

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

Uh oh!

Dev record

Named result

Gates

Uh oh!

github-actions Bot commented Nov 5, 2026

🏷️ CI Guide

Uh oh!

PerryZhang01 commented Nov 5, 2026

Uh oh!

PerryZhang01 commented Nov 5, 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!

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!

PerryZhang01 commented Nov 5, 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!

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!

PerryZhang01 commented Nov 6, 2026

Uh oh!

PerryZhang01 commented Nov 9, 2026

Uh oh!

PerryZhang01 commented Nov 10, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

PerryZhang01 commented Nov 5, 2026 •

edited

Loading