fix(mla): plan persistent decode over the rows q/o have by gbyu-amd · Pull Request #2557 · ROCm/ATOM

gbyu-amd · 2026-10-13T15:45:51Z

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

prepare_decode built the persistent MLA work plan over running_bs. Under DP attention that is the group's batch, but while a peer is prefilling, running_tokens -- the height q and o are allocated at -- is this rank's own count (ForwardMode.running_tokens is not unified), so an eager decode step has fewer q/o rows than the plan has sequences. The padded sequences are published with zero KV, yet the v1_2 planner still hands each of them a work item whenever the CU it lands on has payload left (fixed_over_head_num_blocks is 16 at page size 1), and the persistent kernel loads q[qo_start] and stores o[qo_start] for every item, empty or not. On Kimi-K3 dp8ep8 with the FlyDSL PS1 decode (AITER_MLA_DECODE_PS1_FLYDSL=1), AgentX at concurrency 48 died twice, 28 and 31 min in, on "Memory access fault ... Page not present", in aiter::mla_a8w8_qh96_1tg_16mx6_64nx1_ps workgroup 253 of 256, reading past q. Where the read lands on mapped memory, the store instead writes 96 KB of zeros per padded row past o. Plan min(running_bs, running_tokens // max_seqlen_q) sequences. That is running_bs on a DP-unified step, which pads the activations as well, so graph replay and the zero output of padded rows are unchanged; while a peer is prefilling it is this rank's own batch. cprr keeps scheduled_bs. Tests: the planned batch for decoding ranks with fewer rows than the group's batch and for unified steps, at q=1 and q=3; both cases fail on main. The rest of test_h2d_attention_publication.py is unchanged by this: its 20 sparse cases fail on main as well, on the stub config lacking max_num_batched_tokens (aiter_mla.py:479). Also checked on MI455 (gfx1250): 3000 random 1-3-request batches padded to 4 or 8 gave the padded tail work items in 989 plans (workers 238..255) and in none with this change; PS1 output on the real rows stays within one bf16 ulp of the run with q/o padded to running_bs; AgentX c48 ran past 2.5 h without a fault. Co-authored-by: Cursor <cursoragent@cursor.com>

github-actions · 2026-10-13T15:46:19Z

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

gbyu-amd

No blocking issues @ bd06445. One sibling path below, two suggestions inline.

Checked: while a peer is prefilling, ForwardMode agrees running_bs across DP but keeps running_tokens per rank (since #2088), so a plan over running_bs gives padded sequences work items past the q/o rows. The new test fails on main ([4, 4, 4, 4] != [1, 2, 4, 4]) and passes here.
Accepted with reservation: running_tokens // max_seqlen_q relies on the decode rows being rectangular at max_seqlen_q (see inline).
Watch next: the is_sparse_mtp branch of prepare_decode still plans its dense ctx_mla_ps over running_bs, and per the comment in attention_mla.py the dense layers consume it, so the same mismatch applies there under DP attention with MTP. Hoist q_rows_bs above that if and use it in both branches, or note why the branch cannot hit it.

Evidence

RUN_H2D_GPU_TESTS=1 pytest tests/test_h2d_attention_publication.py on gfx950:

bd06445: 39 passed, 20 failed
main 661ef76 plus this test: 37 passed, 22 failed (the same 20, plus both new cases)

The 20 match on both sides: AttributeError: 'types.SimpleNamespace' object has no attribute 'max_num_batched_tokens'.

gbyu-amd · 2026-10-14T02:41:23Z

+            # is prefilling. The persistent kernel loads q[qo_start] and stores
+            # o[qo_start] for every work item, and the planner hands a zero-KV
+            # sequence one whenever its CU has payload left.
+            q_rows_bs = min(running_bs, running_tokens // max_seqlen_q)


This floors below scheduled_bs whenever a sequence forwards fewer than max_seqlen_q rows, e.g. a DSpark q-bucket or ragged shrink under DP, where max_seqlen_q is the group max. Real sequences then get no work item and the output is silently wrong. The rectangular positions above already break on that layout, so it is not a regression. Add assert q_rows_bs >= scheduled_bs.

gbyu-amd · 2026-10-14T02:41:23Z

+
+
+@pytest.mark.parametrize("q", [1, 3])
+def test_mla_persistent_plan_spans_only_rows_q_has(monkeypatch, q):


This file is skipped unless RUN_H2D_GPU_TESTS=1, so CI does not run this test, and 20 cases here already fail on main because the stub config has no max_num_batched_tokens. Add that field to the stub config in make_runner so the file passes as a whole again.

Review follow-up. The sparse-MTP branch of prepare_decode built its dense ctx_mla_ps -- what the dense layers read -- over running_bs, the same mismatch under DP attention while a peer is prefilling. Compute q_rows_bs once, above the branch, and plan both dense builds over it; cprr keeps scheduled_bs. Assert q_rows_bs >= scheduled_bs: like the positions above, the plan lays every sequence out at max_seqlen_q rows, and a step forwarding fewer per sequence would leave real sequences without a work item. Co-authored-by: Cursor <cursoragent@cursor.com>

…ter #6121 #6530; online FP8 rocm/fw-bringup:gfx1250-atom-k3-20261014 now points to sha256:c8260adc5e54d4bee43ff67dabae4fdd59868ea2f87f6ba7974d6697974ae9ff. That build replaces the morning one (4da3dd98...). - Commits are unchanged: ATOM 0adc645, aiter 1305f0c8 (main of 2026-10-14), mori 695d4c03. - Unmerged PRs in the image: - ATOM #2532, #2533, #2522; - ATOM #2557 (persistent MLA decode plan), #2554 (online per_1x128 FP8 with E8M0 scales), #2466 (gfx1250 MXFP4 GEMM always Triton); - aiter #6121 (gfx1250 FlyDSL gather_kv_b_proj) and #6530 (its own commit only: K3's tuned mxfp8_128 GEMMs). - ATOM #2380 is dropped. With online quantization, kv_b_proj is ptpc FP8 and #6121's FlyDSL gather serves it. The 14k/500 c32 bench ran its chunked prefill without the PHI assert, and launch_kimi_k3_gather_kv_b_proj_8wave was compiled. - Launch changes: - ATOM_USE_TRITON_GEMM and ATOM_UNFUSED_GATHER_KV_B_PROJ dropped; - ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=1 added; - --online_quant_config added (per_block_fp8, kv_b_proj ptpc_fp8; conv1d and experts excluded). - Patch files: one per PR (#2532 and #2533 share one), in apply order. - Validated on CT11 (node0) + CT10: - full GSM8K strict-match 0.9591 +-0.0055, flexible 0.9583 (morning build: 0.9659); - 14k/500 c32: 638.1 tok/s, median TTFT 4.4 s, median TPOT 32.5 ms (morning build: 477.9 tok/s, 7.7 s, 50.5 ms); - peak_torch 248 GB, KV about 101 GB/rank (was 294 GB and 56 GB). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Comment-only. The planning note keeps the one constraint the code cannot show; the assert's message stands without one. Co-authored-by: Cursor <cursoragent@cursor.com>

…; AgentX concurrency 32 - Image: ATOM and aiter main of 2026-10-16 (d82f906, dee019f6), mori 695d4c03, FlyDSL 0.4.0 (aiter main's pin), plus the four PRs not in main yet: ATOM #2557, #2554, aiter #6525, #6530. Patches replaced accordingly; merged PRs are no longer listed. Full GSM8K 0.9568 on CT7 + CT15. - Launch: ATOM main binds MegaMoE with ATOM_MEGA_STAGE2_FUSED and runs the compact plan with ATOM_MEGA_STAGE1_FUSED; the check line and the known issues follow the new log format and aiter #6525's cco sizing. - AgentX: concurrency 32 on the 20261014 image plus K3's grouped-MoE tuning, PrefillDelayer MAX_QUEUE_MS=5000, 1800 s (2026-10-16); 16 and 64 to follow. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

valarLip · 2026-10-26T09:02:11Z

-            ctx_mla_ps = self.set_mla_persistent_worker_buffers(
-                running_bs, max_seqlen_q
-            )
+            ctx_mla_ps = self.set_mla_persistent_worker_buffers(q_rows_bs, max_seqlen_q)


should always use scheduled_bs

Review follow-up (valarLip): plan both dense builds over scheduled_bs, as the cprr branch already did, instead of min(running_bs, running_tokens // max_seqlen_q). The persistent kernel is work_indptr-driven, so the padded tail gets no work item; the q_rows_bs floor and its assert go away. For the padded rows of o this is a no-op on fp8 opus and hk v32, which never wrote them. bf16 opus and the gfx1250 FlyDSL PS1 kernel wrote them as 0 (the output scale is 0 when nothing is attended); they are now left as allocated. On gfx1250 that is what the cprr branch has always done, and Kimi-K3 with DSpark + DCP8 runs over it at GSM8K 0.9545 / 0.9606. The bf16-KV models in CI (GLM-5.3-Flash and DeepSeek-V4.1-Flash with --kv_cache_dtype bf16) exercise bf16 opus. The planned-batch test now expects scheduled_bs on the unified step too: [1, 2, 2, 4]. Co-authored-by: Cursor <cursoragent@cursor.com>

Keep the two load-bearing reasons -- the unbounded q[qo_start] / o[qo_start] access and the cprr metadata kernel's div-by-0 on a 0-query tail -- and drop the restatement of what the code already shows. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

…5 notes The 20261027 image (sha256:5eefef68; ATOM d8877f1, aiter 7061cbd8) dropped the ATOM #2554 and #2557 overlays that the 20261016 and 20261022 images carried. As shipped it fails in post-load weight processing with `AssertionError: 6284 % 16 == 12`: KDA's fused in_proj (N 6284 per rank) is preshuffled without #2554's padding of N. New section on the image: what changed, the failure, and a workaround that serves /app/ATOM with both overlays. The patch files are under recipes/Kimi-K3-455-TP8-patches/; the block makes 6c960542, the same tree as the validated run. GSM8K on E15 trays 05 + 07, 2026-10-27, K3_MOE=ep, custom all-reduce: strict-match 0.9651 +-0.0051, flexible-extract 0.9636. Before you start: the data-plane NIC where it is not enp1s0f1; the checkpoint revision used (f831ab66); the file-list md5 needs LC_ALL=en_US.UTF-8 sort (C sorts to 2773d607...). Known issues: #6768 is merged, retuned, and in the new image; #2647 is in the new image; rows for the startup failure and the md5 locale. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The recipe patch recipes/Kimi-K3-455-wideEP-patches/atom-main-2557.diff, as rocm/fw-bringup:gfx1250-atom-k3-20261016 applies it to /app/ATOM. #2557 (open) Co-authored-by: Cursor <cursoragent@cursor.com> (cherry picked from commit d8ed401)

rocm/fw-bringup:gfx1250-atom-k3-20261027 was pushed again on 2026-10-27 (sha256:d4ceaa4a). The second push is the first (sha256:5eefef68) plus one layer: /app/ATOM and the launcher's ATOM_COMMIT move to 2be1506, branch zejun/k3-image-20261027, which is d8877f1 plus #2557 (8a24a91) and #2554 (e16b3dd, without its tests) as image overlays. Its tree is 161af75, the one that gave GSM8K 0.9651 on E15 05 + 07. The section now tells the two pushes apart. For the first, pull again, or point its launcher at 2be1506, which it fetches from GitHub. That replaces the patch-file workaround, so the patch files go. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

zufayu requested a review from jiayyu October 14, 2026 01:27

zufayu assigned jiayyu Oct 14, 2026

This comment was marked as duplicate.

Sign in to view

gbyu-amd commented Oct 14, 2026

View reviewed changes

gbyu-amd marked this pull request as draft October 14, 2026 06:42

gbyu-amd and others added 2 commits October 14, 2026 13:46

docs(mla): shorten the persistent-plan comments

a0f1bd0

Comment-only. The planning note keeps the one constraint the code cannot show; the assert's message stands without one. Co-authored-by: Cursor <cursoragent@cursor.com>

Merge branch 'main' into gyu/mla-persistent-plan-q-rows

3e6d653

gbyu-amd marked this pull request as ready for review October 14, 2026 13:49

gbyu-amd added the ci:atom label Oct 15, 2026

gbyu-amd requested a review from valarLip October 15, 2026 06:20

Merge branch 'main' into gyu/mla-persistent-plan-q-rows

8a24a91

gbyu-amd requested a review from zejunchen-zejun October 19, 2026 02:41

zejunchen-zejun previously approved these changes Oct 19, 2026

View reviewed changes

Merge branch 'main' into gyu/mla-persistent-plan-q-rows

9234e3b

valarLip requested changes Oct 26, 2026

View reviewed changes

gbyu-amd dismissed zejunchen-zejun’s stale review via 652422e October 26, 2026 10:07

gbyu-amd and others added 2 commits October 26, 2026 05:16

Merge branch 'main' into gyu/mla-persistent-plan-q-rows

e7d400d

gbyu-amd mentioned this pull request Oct 27, 2026

[DSpark][DCP] Allow speculative decode with DCP on gfx1250's PS1 MLA kernel #2760

Open

zejunchen-zejun mentioned this pull request Nov 4, 2026

[GFX1250][K3][LMCache] Enable LMCache MP for K3 TP8DCP8 on GFX1250 #2973

Draft

gbyu-amd mentioned this pull request Nov 5, 2026

feat(kimi-k3): add gfx1250 monokernel #2894

Draft

Navigation Menu

Repository navigation

fix(mla): plan persistent decode over the rows q/o have - #2557

fix(mla): plan persistent decode over the rows q/o have#2557
gbyu-amd wants to merge 9 commits into
mainfrom
gyu/mla-persistent-plan-q-rows

gbyu-amd commented Oct 13, 2026

Uh oh!

github-actions Bot commented Oct 13, 2026

Uh oh!

This comment was marked as duplicate.

Uh oh!

gbyu-amd left a comment

Uh oh!

gbyu-amd Oct 14, 2026

Uh oh!

gbyu-amd Oct 14, 2026

Uh oh!

valarLip Oct 26, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants



		@pytest.mark.parametrize("q", [1, 3])
		def test_mla_persistent_plan_spans_only_rows_q_has(monkeypatch, q):

Conversation

gbyu-amd commented Oct 13, 2026

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

github-actions Bot commented Oct 13, 2026

🏷️ CI Guide

Uh oh!

This comment was marked as duplicate.

Uh oh!

gbyu-amd left a comment

Choose a reason for hiding this comment

Uh oh!

gbyu-amd Oct 14, 2026

Choose a reason for hiding this comment

Uh oh!

gbyu-amd Oct 14, 2026

Choose a reason for hiding this comment

Uh oh!

valarLip Oct 26, 2026

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