[ROCm][Perf] Use AITER single-launch KDA decode in GLM-5.3-Flash by eppaneamd · Pull Request #63612 · vllm-project/vllm

eppaneamd · 2026-10-25T19:54:57Z

Overview

On ROCm with AITER, a plain decode step of a GLM-5.3-Flash KDA layer runs the conv update, the recurrence and the gated RMSNorm as separate launches. This PR runs them as one launch of AITER's Gluon decode kernel through the existing rocm_aiter_ops.kda_fused_recurrent_packed_decode wrapper, which the Kimi-K3 layer has used since #57392. Companion prefill change for the same layer: #63568.

Claims

ROCm, VLLM_ROCM_USE_AITER=1, CDNA4 or newer (gfx950 is the only architecture run): a plain decode step of a GLM-5.3-Flash KDA layer takes one AITER launch in place of three. No new flag or environment variable; any explicit --kda-decode-backend keeps the stock chain.
GLM-5.3-Flash MXFP4, MI355X, TP2, one image plus text per request:
- FULL_DECODE_ONLY graphs: output throughput +5.63 % over concurrency 16–64 (SE 0.80 pp), +13.70 % at concurrency 1, every row positive;
- breakable FULL_AND_PIECEWISE graphs: +3.10 % over concurrency 16–64 (SE 1.13 pp), −1.75 % at concurrency 4 and −2.31 % on the long-prompt row;
- GSM8K unchanged within run-to-run variation in both.
Unchanged: CUDA, CPU and gfx942; steps that include a prefill; speculative decoding; non-fp32 recurrent states; other KDA models.

Validation

Serving.

Image vllm/vllm-openai-rocm:nightly-rocm72-c50aca14416ef29f6ed0ed256302f30646726dcd (vLLM c50aca144, AITER 0.1.26), whose kda.py and vllm/_aiter_ops.py equal this PR's base at the time of measurement (5f63ea461c). Baseline: that base's kda.py; patched: this PR's at the time. Since then the PR's kda.py inlines the eligibility check, honours --kda-decode-backend, refreshes its decode conv-weight copy after weight writes, rejects an unknown kda_decode_backend, and was merged with main's fused q|k|v conv1d parameter ([Bugfix] Fuse GLM-5.3 KDA q|k|v convs into one conv1d parameter #60991); none of these changes what a decode step runs, in the stock chain or the AITER route.
AMD Instinct MI355X (gfx950), TP2, amd/GLM-5.3-Flash-Quark-MXFP4 at revision b5688f25491202978c19c4d036eef579f61bbe07.
Two graph configurations: (i) the flags below; (ii) the same with breakable FULL_AND_PIECEWISE graphs and the conv2d custom op disabled. Without -conv2d, MIOpen solver searches for the vision encoder's convolution stalled both arms, so that run's throughput is not reported.
Per configuration, 4 server launches in the order baseline, patched, baseline, patched, each with one discarded warm-up run. 1024 input / 128 output tokens at concurrency 1–64, plus a long-prompt row (5120 / 16, concurrency 16) that weights the steps the kernel does not take.

VLLM_ROCM_USE_AITER=1 VLLM_DISABLE_COMPILE_CACHE=1 vllm serve amd/GLM-5.3-Flash-Quark-MXFP4 \
  --revision b5688f25491202978c19c4d036eef579f61bbe07 \
  --quantization quark --tensor-parallel-size 2 --max-model-len 8192 \
  --gpu-memory-utilization 0.92 --limit-mm-per-prompt '{"image": 1}' \
  --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY"}' \
  --no-enable-prefix-caching --trust-remote-code
# configuration (ii): VLLM_USE_BREAKABLE_CUDAGRAPH=1 and
#   --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE", "custom_ops": ["all", "-conv2d"]}'

# at each concurrency C in 1, 4, 16, 32, 64 with N = 64, 64, 160, 320, 640 prompts
vllm bench serve --backend openai-chat --endpoint /v1/chat/completions \
  --model amd/GLM-5.3-Flash-Quark-MXFP4 --dataset-name random-mm \
  --random-input-len 1024 --random-output-len 128 \
  --random-mm-base-items-per-request 1 --random-mm-num-mm-items-range-ratio 0 \
  --random-mm-limit-mm-per-prompt '{"image": 1, "video": 0}' \
  --random-mm-bucket-config '{(512, 512, 1): 1.0}' \
  --ignore-eos --request-rate inf --max-concurrency $C --num-prompts $N \
  --num-warmups 8 --seed 0 --percentile-metrics ttft,tpot,itl,e2el --metric-percentiles 50,99
# long-prompt row: --random-input-len 5120 --random-output-len 16 --max-concurrency 16 --num-prompts 128

Concurrency	(i) baseline tok/s	(i) patched	(i) change	(ii) baseline tok/s	(ii) patched	(ii) change
1	104.96	119.34	+13.70 % (0.21)	104.36	105.17	+0.78 % (1.36)
4	303.60	332.39	+9.48 % (0.39)	357.45	351.19	−1.75 % (0.27)
16	777.79	828.45	+6.51 % (0.19)	769.60	799.97	+3.95 % (4.42)
32	1020.67	1085.43	+6.34 % (1.58)	1085.14	1113.23	+2.59 % (0.61)
64	1298.82	1351.09	+4.02 % (0.72)	1361.57	1399.36	+2.78 % (0.78)
16, long prompt (5120 / 16)	59.58	60.20	+1.04 % (0.34)	61.36	59.94	−2.31 % (1.61)

Output throughput, mean of each arm's two launches; SE of the change in pp in brackets, from the spread between launches (rough with two launches per arm).
TPOT: −2.6 to −13.1 % in (i); −5.0 to −5.7 % at concurrency 16–64 in (ii).
TTFT: no consistent change beyond the launch-to-launch spread, except +5.1 % on the long-prompt row in (ii).
Every patched server log contains GLM-5.3-Flash KDA decode: AITER single-launch kernel. once; no baseline log does. KV cache memory is equal across arms (172.7 GiB in (i), 160.54 GiB in (ii)).
VLLM_GPU_SYNC_CHECK=error on the patched breakable-graph server (without -conv2d), all rows: no GPU<->CPU sync detected line.
GSM8K, 5-shot, 1319 questions, greedy, strict match, paired McNemar p:
- (i): 92.19 → 93.03 % (p 0.315) and 92.80 → 92.04 % (p 0.289); the two baseline launches differ by 0.61 pp (p 0.403);
- (ii): 92.04 → 92.57 % (p 0.510) and 92.19 → 92.57 % (p 0.640); baselines differ by 0.15 pp (p 0.910).

lm_eval run --model local-completions \
  --model_args "model=amd/GLM-5.3-Flash-Quark-MXFP4,base_url=http://127.0.0.1:8000/v1/completions,num_concurrent=32,max_retries=3,tokenized_requests=False,tokenizer_backend=None" \
  --tasks gsm8k --num_fewshot 5 --seed 0 --gen_kwargs temperature=0,max_gen_toks=512
# lm_eval 0.4.13

Unit tests.

test_aiter_decode_route_follows_replayed_step (tests/models/glm5next/test_kda_recurrent.py):
- Setup: two real Glm5NextLinearAttention layers sharing weights and caches, one with VLLM_ROCM_USE_AITER=0 and one with =1. It calls their process_weights_after_loading, as vLLM's loaders do.
- It runs in both conv-state layouts: the default and the dim-first one that VLLM_SSM_CONV_STATE_LAYOUT=DS selects. The fused kernel receives the conv state with different strides in each.
- It captures a breakable PIECEWISE graph of each on an all-decode batch, then replays a mixed step and a decode step.
- Mixed step: no fused launch, bit for bit equal to the stock layer.
- Decode step: exactly one fused launch of 3 rows; output within 2 bf16 ulp of its largest value; conv state exact; recurrent state within atol=2e-3, rtol=3e-2. Kimi-K3's check of this kernel uses rtol=2e-3 on up to 128 tokens; one of this test's 41.9M state elements reaches 3.8e-3, so rtol is that of Kimi-K3's spec-decode check.
- Both steps: the fused layer's conv and recurrent state keep their saved values exactly in every slot other than the batch's (1–3).
test_aiter_decode_reads_refitted_fused_conv (tests/models/glm5next/test_kda_conv_weight_refit.py):
- It loads conv weights through the weight loader, refits them through it, and reloads them through the layerwise reload.
- After each, the AITER decode must receive conv1d.weight in the [3, width, dim] layout, at an unchanged address.
- It checks that the layer is a deferred attention layer, the condition under which vLLM's loaders call its process_weights_after_loading.
- The kernel call is replaced, so the test does not need AITER.
MI355X, image vllm/vllm-openai-rocm:nightly-rocm72-3966924d566340e971a6da313e790d05492aa505 (AITER 0.1.26) with this PR's vllm/ files over the installed package, VLLM_ROCM_USE_AITER=1:
- test_kda_conv_weight_refit.py: 4 passed;
- test_kda_recurrent.py: 8 passed (the route test once per layout), 14 skipped (12 need FlashKDA, 2 are NVIDIA-only).
Each check fails when what it covers is removed. These are mutations of this PR's kda.py, run once each in the same setup:
- no refresh after a shard load: the refit test fails on the first load (test_kda_conv_weight_refit.py:181);
- no refresh in process_weights_after_loading: the refit test fails after the layerwise reload (:191), and in both layouts the route test's decode output differs from the stock layer's (test_kda_recurrent.py:480; 87 % of elements in the default layout, 83 % in the dim-first one);
- the fused route also writing conv-state slot 0: the untouched-slot check fails in both layouts (test_kda_recurrent.py:479);
- the fused route reading the conv state with the default layout's strides in either layout: the dim-first case's decode output differs (test_kda_recurrent.py:480, 78 % of elements); the default case passes.
On the serving image (above), before the review changes:
- The route test passed. The decode output's max abs error was 1.562e-02 against a bound of 5.03e-02; recurrent state 3.99e-03.
- On main's kda.py it failed the decode step: AssertionError: [1, 1, 1]: fused launches (rows each) [].
MI355X, serving image, config (i) server with --kda-decode-backend triton: starts and serves, and logs no AITER single-launch kernel line.
CI: the H200 GLM5Next step and its MI355 mirror run pytest -v -s models/glm5next. The route test needs AITER's Gluon KDA kernel and is skipped without it, as on CPU.
ruff check, ruff format --check (ruff 0.14.0), typos (1.51.1): clean. tools/pre_commit/mypy.py 3.11 (mypy 1.20.2): no issues.

Details

Scope. Glm5NextLinearAttention (vllm/models/glm5next/common/kda.py), the KDA layer of Glm5NextForCausalLM and Glm5NextForConditionalGeneration. The MTP layer is MLA and unaffected. On main only vllm/models/kimi_k3/amd/kda.py calls this kernel.
Eligibility, decided once in __init__: fp32 recurrent state, head_dim 128, width-4 conv, a lower bound, no speculative decoding, --kda-decode-backend at auto, and rocm_aiter_ops.is_kda_gluon_decode_available(). An unknown kda_decode_backend raises ValueError, as in Kimi-K3; additional_config bypasses the CLI's choices. An eligible layer logs GLM-5.3-Flash KDA decode: AITER single-launch kernel. once.
Route. _forward takes the fused launch when a step has no prefill and at least one decode. The kernel applies the gated RMSNorm, so an eligible layer passes out_gate to _forward and forward skips o_norm; other steps, and profile runs without metadata, apply the norm at the end of _forward. The route is chosen in _forward because a breakable PIECEWISE graph records forward once on a decode-only batch and replays only the eager break against each step's metadata, so a route chosen in forward would be fixed at capture. FULL graphs capture _forward inline, which is safe because GDN FULL graphs replay uniform decode batches only.
Weights and indices.
- AITER reads the q|k|v conv weight as a contiguous [3, width, dim] tensor, while conv1d.weight is [3·dim, 1, width].
- An eligible layer keeps that layout in a non-persistent buffer, decode_conv1d_weight. It is allocated once in __init__, so graphs keep its address, as in Kimi-K3's ROCm layer.
- The buffer is refreshed from conv1d.weight in two places:
  - the layer's process_weights_after_loading, which BaseModelLoader.load_model runs after loading and the layerwise reload runs at its end (reload_weights and the weight-transfer engines that use it);
  - the conv weight loader, after every shard load, for refits that call load_weights directly.
- The IPC weight cache maps the buffer the daemon filled, so the hook skips it there.
- State indices go in as an [N, 1] view because the block-table column can be strided.
Numerics. AITER keeps the conv and attention outputs in fp32 through the norm; the stock chain rounds to bf16 before o_norm. Decode outputs therefore agree to 2 bf16 ulp, not bit for bit.
Trade-offs. A step that does not take the kernel pays one extra copy_ of the layer output, and under breakable PIECEWISE graphs runs o_norm eagerly, the likely cause of the (ii) long-prompt result. Fusing only the conv update and the recurrence, keeping o_norm in forward, would avoid it, at the cost of part of the decode gain and a new _aiter_ops wrapper.
Limits. gfx1250 passes the probe and has not been run. Without a lower bound the stock kernel falls back to −5.0 while AITER's gate would be unbounded, so that configuration keeps the stock chain. AITER asserts a recurrent slot stride of whole 128-float rows. Not measured: TP 1, 4 and 8; inputs over 5120 tokens; concurrency above 64.
Weight writes the decode copy does not follow (from reading the code, not run):
- nccl_m2n's direct resharding writes parameters without load_weights or a reload, so the copy stays stale.
- Sharded RDT counts the buffer as part of the layer, cannot bake the layer, and stops while building its plan.
- --kda-decode-backend triton keeps the stock chain for both.
- Not checked: the tensorizer loader, which overrides load_model.
- An AITER kernel that reads conv1d.weight in place would remove the copy.
Related PRs.
- [Perf][GLM-5.3-Flash] KDA decode and prefill through FlashInfer #63544 (draft): FlashInfer fused KDA decode for this layer on CUDA, same structure; whichever lands second adapts.
- [Perf][ROCm][GLM-5.3-Flash] Default to breakable CUDA graphs on ROCm #63070 (draft): breakable CUDA graphs as the ROCm default for GLM-5.3-Flash; configuration (ii) then applies.
- [Bugfix][GLM] Invalidate merged KDA convolution weights after refit #56298 (open): targets the separate q/k/v convs, which [Bugfix] Fuse GLM-5.3 KDA q|k|v convs into one conv1d parameter #60991 merged into one conv1d parameter; this PR no longer interacts with it.
- [Perf][ROCm][GLM-5.3-Flash] Defer KDA spec decode state writes #63071, [Perf][ROCm][GLM-5.3-Flash] Gluon KDA spec decode kernel for gfx950 #63073 (drafts): speculative decoding on this layer, which keeps the stock chain here.

This change was developed with AI assistance (Claude); I reviewed every line and ran the tests listed.

Pull Request Checklist

I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).
AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

^{Need help on this PR? Tag @codesmith-bot with what you need.}

With VLLM_ROCM_USE_AITER=1 on CDNA4 or newer (used so far on gfx950 only), plain decode steps of a GLM-5.3-Flash KDA layer run AITER's Gluon decode kernel through rocm_aiter_ops, as the Kimi-K3 AMD layer does: one launch for the causal conv update, the delta-rule recurrence and the gated RMSNorm. A layer is eligible with an fp32 recurrent state, head_dim 128, a width-4 conv, a lower bound, no speculative decoding and the default --kda-decode-backend; any explicit value keeps the Triton chain. Other layers and steps keep the Triton chain as well. An eligible layer hands the gated RMSNorm to _forward through out_gate, so the route is chosen per step in _forward from the step's metadata. Breakable PIECEWISE graphs replay _forward against the replayed step's metadata, so a prefill or mixed step replayed from a graph captured with a decode batch does not take the decode kernel. Those steps, and profile runs without metadata, apply the norm at the end of _forward. The decode copy of the conv weight is built inside the existing lazy merged-weight block, from the merged weight, so invalidating the merged weight rebuilds both. Test: a ROCm test captures a breakable PIECEWISE graph of a real layer with and without AITER on an all-decode batch, then replays a mixed step (stock chain, bit for bit) and a decode step (one fused launch, within tolerance of the stock chain). Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Rony Leppänen <rony.leppanen@amd.com>

github-actions · 2026-10-25T19:55:05Z

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

amd-sriram · 2026-10-26T19:36:42Z

TP4 MI355X numbers for this PR together with #63568 (+6.1% throughput, TTFT -7.3%, TPOT -5.6%, GPQA/NIAH unchanged): #63568 (comment)

mergify · 2026-10-27T12:51:02Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @eppaneamd.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Signed-off-by: Rony Leppänen <rony.leppanen@amd.com> Co-authored-by: Claude <noreply@anthropic.com>

akii96

The overall routing approach makes sense to me. My questions are mainly about making the weight-update lifecycle and fallback behavior explicit, so reviewers can more easily confirm that the optimization stays safe across refits and non-default configurations. Thanks for the good work!

akii96 · 2026-10-28T07:51:38Z

+        if out_gate is not None and self._decode_conv_weight is None:
+            self._decode_conv_weight = (
+                conv_weights.view(3, self.local_projection_size, self.conv_size)
+                .transpose(1, 2)
+                .contiguous()
+            )


Am I reading the _decode_conv_weight lifecycle correctly? It looks like the transposed copy is created on the first eligible forward and then reused. If conv1d.weight is updated later by a refit or reload, what refreshes this copy? Would extending tests/models/glm5next/test_kda_conv_weight_refit.py:87-133 to exercise the AITER path make this clearer to reviewers?

Thanks for noticing this! I have changed as follows:

The [3, width, dim] copy is now a non-persistent buffer, decode_conv1d_weight, allocated once in __init__, as in Kimi-K3's ROCm KDA layer.

It is refreshed from conv1d.weight in two places:

the layer's process_weights_after_loading, which BaseModelLoader.load_model runs after loading and the layerwise reload runs at its end;

the conv weight loader, after each shard load, which covers refits that call load_weights directly.

test_aiter_decode_reads_refitted_fused_conv adds the AITER path to the refit test. It checks the weight the AITER decode receives after a load, after a refit through the loader, and after a layerwise reload, each time at the
same address.

akii96 · 2026-10-28T07:56:35Z

+    # The fused kernel keeps outputs in fp32 where the stock chain rounds to bf16:
+    # 2 bf16 ulp of the largest value. State rtol is Kimi-K3's spec-decode one, as
+    # one of these 41.9M elements reaches 3.8e-3, past its packed-decode 2e-3.
+    fused_rtol = 2 * torch.finfo(torch.bfloat16).eps
+    fused_state_tol = {"atol": 2e-3, "rtol": 3e-2}
+    # (query lengths, expected fused launches (rows each), output rtol, state tol)
+    steps: list[tuple[list[int], list[int], float, dict[str, float]]] = [
+        ([1, 1, 50], [], 0, {"atol": 0, "rtol": 0}),  # Mixed: stock, bit for bit.
+        ([1, 1, 1], [3], fused_rtol, fused_state_tol),  # Decode: one fused launch.
+    ]


The underlying packed-decode test already checks untouched slots exactly. Would it also be useful to check them here, so this GLM integration test confirms that its state-index mapping only updates slots 1–3?

Thanks. Now every slot other than 1–3 of the fused layer's conv and recurrent state must equal the saved caches.

akii96 · 2026-10-28T08:07:33Z

+        additional_config = vllm_config.additional_config
+        decode_backend = (
+            additional_config.get("kda_decode_backend", "auto")
+            if isinstance(additional_config, dict)
+            else "auto"
+        )
+        # AITER's single-launch decode needs an fp32 state, head_dim 128 and a
+        # width-4 conv. Without a lower bound the stock kernel falls back to -5.0
+        # while AITER switches to the unbounded gate, so that config keeps the
+        # stock chain, as do spec decode and any explicit --kda-decode-backend.
+        self._use_aiter_decode = (
+            decode_backend == "auto"


Would it be worth validating kda_decode_backend here, similar to _resolve_kda_prefill_backend? The CLI already checks its choices, but an unknown value supplied directly through additional_config appears to silently disable the AITER path.

An unknown kda_decode_backend now raises ValueError. The accepted values and the message are the
same as in Kimi-K3's resolve_kda_decode_backend.

I like @akii96 's suggestion of adding resolve_kda_decode_backend, this is similar to NV

…code in GLM-5.3-Flash PR head f7b184b, applied as its diff against merge base 5f63ea4.

AITER's single-launch decode reads a [3, width, dim] copy of the q|k|v conv weight. That copy was built on the first eligible decode and never refreshed, so a refit or reload of conv1d.weight left the AITER decode on the old weights. The copy is now a non-persistent buffer, allocated once in __init__ so cudagraphs keep its address, as in Kimi-K3's ROCm KDA layer. It is refreshed from conv1d.weight by process_weights_after_loading, which BaseModelLoader runs after loading and a layerwise reload runs at its end, and by the conv weight loader after every shard load, for refits that call load_weights outside the reload lifecycle. Defining process_weights_after_loading also makes the layer a deferred attention layer: a layerwise reload finalizes it after the other layers, as it does Kimi-K3's ROCm layer. An unknown kda_decode_backend in additional_config now raises ValueError, as in Kimi-K3, instead of silently keeping the stock chain. Tests: test_aiter_decode_reads_refitted_fused_conv checks the weight the AITER decode receives after a load, a refit through the loader and a layerwise reload. The GLM route test fills the copy as vLLM's loaders do, and checks that cache slots outside the batch keep their values. Signed-off-by: Rony Leppänen <rony.leppanen@amd.com> Co-authored-by: Claude <noreply@anthropic.com>

akii96

Thanks for addressing everything! I am overall happy with how this looks now just one optional nit but overall LGTM 🚀

IMO this is ready to come out of draft so maintainers can take a look!

akii96 · 2026-10-30T13:08:14Z

+def test_aiter_decode_route_follows_replayed_step(dist_init, monkeypatch):
+    """AITER's single-launch decode kernel runs on plain decode steps only.
+
+    Breakable PIECEWISE graphs replay ``_forward`` against each step's metadata,
+    so the route is chosen there. Capture on an all-decode batch, then replay a
+    mixed step (stock chain) and a decode step (one fused launch); a route fixed
+    in ``forward`` fails one of the two. The reference is a layer built without
+    AITER, with the same weights and caches.
+    """


Super Optional nit -> this test only covers the default conv-state layout. Could you also run it with the other layout (_conv_state_dim_first=True), the same way test_prefill_checkpoint_resumes_suffix does? Then we'd notice if the AITER path ever stops working for it.

Done in ba1ed2b: the route test now runs in both conv-state layouts, and both pass on MI355X.

test_aiter_decode_route_follows_replayed_step now also runs with the dim-first (DS) conv-state layout that VLLM_SSM_CONV_STATE_LAYOUT=DS selects. The AITER route runs in either layout and receives the conv state with different strides; the test covered only the default (SD). Signed-off-by: Rony Leppänen <rony.leppanen@amd.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Main now type-checks tests/models/glm5next strictly (tools/pre_commit/mypy.py), which rejects GDNAttentionMetadata as a value of ForwardContext's AttentionMetadata dict in the route test. Signed-off-by: Rony Leppänen <rony.leppanen@amd.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Rony Leppänen <rony.leppanen@amd.com>

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.

okorzh-amd · 2026-10-30T23:52:56Z

+                self._copy_decode_conv1d_weight()
+
+            delattr(weight, "weight_loader")
+            set_weight_attrs(weight, {"weight_loader": weight_loader})


decode_conv1d_weight is non-persistent, so it is not in state_dict(). Elastic EP scale-up builds each new rank with load_dummy_weights=True, so this hook copies the dummy conv1d.weight into the buffer. The rank then receives only model.state_dict() tensors from a peer and runs no post-load pass. conv1d.weight ends up correct but this copy keeps the dummy values, so on MI355X with AITER and --enable-eplb --enable-elastic-ep plain decode on scaled-up ranks runs the conv with dummy weights.

simondanielsson

Thanks for adding this!

Suggestions:

can we run a vllm bench serve sweep also at 64k/1k and 128k/1k workloads? Good to validate long context decode is also sped up
Can we run niah @ 128k as well?

simondanielsson · 2026-11-02T17:35:55Z

+        additional_config = vllm_config.additional_config
+        decode_backend = (
+            additional_config.get("kda_decode_backend", "auto")
+            if isinstance(additional_config, dict)
+            else "auto"
+        )
+        # AITER's single-launch decode needs an fp32 state, head_dim 128 and a
+        # width-4 conv. Without a lower bound the stock kernel falls back to -5.0
+        # while AITER switches to the unbounded gate, so that config keeps the
+        # stock chain, as do spec decode and any explicit --kda-decode-backend.
+        self._use_aiter_decode = (
+            decode_backend == "auto"


I like @akii96 's suggestion of adding resolve_kda_decode_backend, this is similar to NV

simondanielsson · 2026-11-02T17:37:39Z

+        )
+        if decode_backend not in ("auto", "native", "flashinfer", "triton"):
+            raise ValueError(f"Unsupported KDA decode backend: {decode_backend}")
+        # AITER's single-launch decode needs an fp32 state, head_dim 128 and a


Nit: we can remove this comment as the code says the same thing 👍

simondanielsson · 2026-11-02T17:43:16Z

+        )
+
+    @torch.no_grad()
+    def _copy_decode_conv1d_weight(self) -> None:


Question: can we check if this helper is strictly needed?

simondanielsson · 2026-11-02T17:49:59Z

+        # AITER's single-launch decode applies the gated RMSNorm in-kernel, so an
+        # eligible layer hands the norm to _forward. The route is chosen there
+        # per step: breakable PIECEWISE graphs replay _forward, not this function.
+        out_gate = g2 if self._use_aiter_decode else None


Nit: Is fuse_gated_norm a more clear name? Then

self._forward( qkv_proj_states=qkv, g1=g1, beta=beta, core_attn_out=core_attn_out, g2=g2 if fuse_gated_norm else None, ) if not fuse_gated_norm: core_attn_out = self.o_norm(core_attn_out, g2)

and we wouldn't need to extra docstring comment
"""With ``out_gate``, also apply the gated RMSNorm that ``forward`` skips."""

simondanielsson · 2026-11-02T17:52:39Z

+        ):
+            # One launch for the conv update, recurrence and gated RMSNorm
+            # (out_gate implies no spec decode). The [N, 1] indices: the
+            # block-table column can be strided, which the flat form rejects.


Nit: should we also assert we're using aiter decode here just to be on the safe side?

simondanielsson · 2026-11-02T17:53:30Z

                    0, :num_actual_tokens
                ]
+
+        if out_gate is not None:


Question: Can this happen?

simondanielsson · 2026-11-02T17:56:34Z

+    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.platforms.rocm import get_cdna_version
+
+    return get_cdna_version() >= 4 and rocm_aiter_ops._kda_gluon_decode_importable()


Suggestion: Should we re-use the util that the code is using?rocm_aiter_ops.is_kda_gluon_decode_available

simondanielsson · 2026-11-02T17:57:21Z

+    from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
+    from vllm.v1.attention.backends.utils import compute_causal_conv1d_metadata
+
+    torch.manual_seed(0)


Nit: We can use set_random_seed here 👍

simondanielsson · 2026-11-02T17:59:21Z

+    fused_rtol = 2 * torch.finfo(torch.bfloat16).eps
+    fused_state_tol = {"atol": 2e-3, "rtol": 3e-2}
+    # (query lengths, expected fused launches (rows each), output rtol, state tol)
+    steps: list[tuple[list[int], list[int], float, dict[str, float]]] = [


Suggestion: Could this be a pytest parametrization?

Net diff of vllm-project#63612 at 749147c, which has these commits on main: - [ROCm][Perf] Use AITER single-launch KDA decode in GLM-5.3-Flash - [ROCm][GLM] Refresh the AITER KDA decode conv weight after weight writes - [ROCm][GLM] Test the AITER KDA decode route in both conv-state layouts - [ROCm][GLM] Satisfy mypy in the AITER KDA decode route test Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Rony Leppänen <rony.leppanen@amd.com>

mergify Bot added glm rocm Related to AMD ROCm labels Oct 25, 2026

github-project-automation Bot added this to AMD Oct 25, 2026

github-project-automation Bot moved this to Todo in AMD Oct 25, 2026

amd-sriram mentioned this pull request Oct 26, 2026

[ROCm][Perf] Add AITER FlashKDA prefill backend for GLM KDA #63568

Draft

5 of 6 tasks

mergify Bot added the needs-rebase label Oct 27, 2026

Merge branch 'main' into rocm-aiter-kda-decode-glm-single-launch

9997cfa

Signed-off-by: Rony Leppänen <rony.leppanen@amd.com> Co-authored-by: Claude <noreply@anthropic.com>

mergify Bot removed the needs-rebase label Oct 27, 2026

akii96 reviewed Oct 28, 2026

View reviewed changes

waqahmed-amd-fi added a commit to waqahmed-amd-fi/vllm that referenced this pull request Oct 28, 2026

Apply vllm-project#63612: [ROCm][Perf] Use AITER single-launch KDA de…

42399c6

…code in GLM-5.3-Flash PR head f7b184b, applied as its diff against merge base 5f63ea4.

waqahmed-amd-fi mentioned this pull request Oct 29, 2026

[ROCm][Perf][GLM-5.3-Flash] Split decodes off the AITER FlashKDA prefill on gfx950 (do not merge) #64243

Draft

6 tasks done

eppaneamd and others added 2 commits October 29, 2026 22:51

Merge branch 'main' into rocm-aiter-kda-decode-glm-single-launch

9486b31

akii96 approved these changes Oct 30, 2026

View reviewed changes

eppaneamd and others added 2 commits October 30, 2026 13:23

eppaneamd marked this pull request as ready for review October 30, 2026 13:44

eppaneamd requested review from AndreasKaratzas, DarkLight1337 and ywang96 as code owners October 30, 2026 13:44

claude Bot reviewed Oct 30, 2026

View reviewed changes

Merge branch 'main' into rocm-aiter-kda-decode-glm-single-launch

749147c

eppaneamd mentioned this pull request Oct 30, 2026

[Triton/Gluon] KDA packed decode: accept the fused conv1d weight ROCm/aiter#7225

Open

4 of 10 tasks

okorzh-amd reviewed Oct 30, 2026

View reviewed changes

okorzh-amd mentioned this pull request Oct 31, 2026

[Bugfix][ROCm][Kimi-K3] Refresh derived weight copies after in-place weight writes #64886

Merged

simondanielsson reviewed Nov 2, 2026

View reviewed changes

amd-sriram mentioned this pull request Nov 4, 2026

[Test only] vllm#63568 + #63612 rebased on nightly dd4e4014 (week-45 control) amd-sriram/vllm#65

Draft

waqahmed-amd-fi mentioned this pull request Nov 4, 2026

[ROCm][Perf][GLM-5.3-Flash] Split decodes off the AITER FlashKDA prefill on gfx950 #65224

Draft

6 tasks done

Conversation

eppaneamd commented Oct 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

Claims

Validation

Details

Uh oh!

github-actions Bot commented Oct 25, 2026

Uh oh!

amd-sriram commented Oct 26, 2026

Uh oh!

mergify Bot commented Oct 27, 2026

Uh oh!

akii96 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!

akii96 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!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

Choose a reason for hiding this comment

Uh oh!

simondanielsson 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!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

eppaneamd commented Oct 25, 2026 •

edited

Loading