[MLA] Address Triton decode KV rows with int64 ids by zejunchen-zejun · Pull Request #2533 · ROCm/ATOM

zejunchen-zejun · 2026-10-12T01:21:25Z

TritonMLAMetadataBuilder hands aiter's decode_attention_fwd a dense table of token-level KV row ids (page_size 1). The kernel multiplies each id it loads by the row stride (576 for MLA) in the id's own type, so with the int32 table the offset wraps for rows >= 2**31 / 576 = 3,728,270, which is 233,016 blocks of
16. The decode then reads outside the layer's KV: garbage, NaN, or a GPU memory fault.

A pool reaches those rows once it has that many blocks and fills them. On Kimi-K3 dp8ep8 (gfx1250, 0.90, 16384 batched tokens) each rank holds about 253k-272k blocks, and with prefix caching the KDA checkpoint images (about 450 MB each) fill the pool past 90%. Decodes on every DP rank then went NaN, starting at the first token of a fresh block above 233,016, and answered !!!! up to max_tokens; all 64 corrupted blocks found were >= 233,329. Writes are unaffected: the cache kernels compute slots in int64.

The table is now allocated as torch.int64; csr_to_dense_block_table stores the int32 ids into it unchanged. Kernel repro on a 3.93M-row fp8 KV: an int32 table faults on rows past the limit, an int64 table matches a torch reference on rows below and above it.

github-actions · 2026-10-12T01:21:48Z

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

…lines) Keep what is needed to run Kimi-K3 dp8ep8 on two B0 trays today, drop the history: - A0: only the mandatory B0->A0 rocjitsu translation (detect, prefix, install, verify). - The rack: hosts, OS/BMC IPs, access; tray operations: is-the-rack-free, perf_setup after a reboot, link check, reboot (BMC Redfish, helicop AC cycle), fabric check (ubench07), checkpoint copy. - Container (rocm/atom-dev 1009) and the PRs it needs: ATOM #2380 (unfused gather_kv_b_proj), aiter #5431 (symm_mem custom AR for the allgather LM head), ATOM #2522 (flydsl dispatch on the fp4 wire), ATOM #2532 and #2533 (prefix caching correctness), with apply/verify commands. - The exact dp8ep8 launch (env + CLI, accuracy and AgentX modes), accuracy gate, GSM8K command and results (0.9621 with and, with #2532/#2533, without prefix caching). - AgentX: ATOM CI method and command, the four-point dp8ep8 results. - Known issues and fixes. Removed: the 0918-ep8 / dp16ep16 configuration and results, old throughput tables, profiling, the optimized source stack, Appendix A, the dispatch-backend history and other superseded notes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Copilot

🟡 Changes recommended

The regression test should fail fast on an incorrect dtype rather than risk an out-of-bounds GPU access.

1 open finding

Assert index dtype before launching kernel · New

What changed in this PR

Prevents Triton MLA decode address overflow for large KV pools.

Changes:

Uses int64 token-row IDs for Triton MLA decode tables.
Adds a GPU numerical regression test across the former int32 limit.

File	Description
`atom/model_ops/attentions/triton_mla.py`	Allocates decode tables with 64-bit IDs.
`tests/test_triton_mla_block_table_dtype.py`	Tests decoding above the 32-bit offset boundary.

🧠 Review effort: Balanced

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

+
+    q = torch.randn(2, H, D, dtype=torch.bfloat16, device=dev)
+    scale = D**-0.5
+    table = torch.stack([low, high]).to(TRITON_BLOCK_TABLE_DTYPE)


XiaobingSuper

LGTM

…KERNEL (ATOM #2522) ATOM #2522 now names the MegaMoE dispatch kernel with ATOM_MEGA_DISPATCH_KERNEL (mori | flydsl_token_major | flydsl_compact), which replaces aiter's MEGA_DISPATCH. The launch sets flydsl_compact, the kernel every result on this page was measured with, so #2522 is now needed always. - Patches: a container patched before the 10-13 revision must drop the old diff; every rank's "Created MegaMoE" line must read kernel=flydsl_compact. - Launch notes: A/B with ATOM_MEGA_DISPATCH_KERNEL=mori; keep flydsl_token_major off the fp4 wire; leave MEGA_DISPATCH unset. - Known issues: the MEGA_DISPATCH conflict error, and a missing or stale #2522. - Result labels say flydsl_compact. Checked: the 10-13 diffs of ATOM #2380, #2522, #2532 and #2533 apply in the recipe's order to the image's ATOM (16cc652) and reverse-apply cleanly.

…20261012 Kimi-K3 dp8ep8 now runs on rocm/fw-bringup-mlperf:gfx1250-atom-20261012 (ROCm 10.1 TheRock wheels, torch 2.11, editable ATOM/aiter/mori in /app), validated on CT15 (node0) + CT14 on 2026-10-13: full GSM8K strict-match 0.9659 +-0.0050, flexible-extract 0.9651, gate pass, ~8 min cold start. - Container and patches: the new image, its private repository, stack and source branches; the DeepSeek-R1 ENV it carries. - Two ATOM patches next to the page (Kimi-K3-455-wideEP-patches/): atom-r1_1012-k3.diff = #2380 (one hunk moved) + #2532 + #2533; atom-r1_1012-mega-vmm.diff = #2522's compact-plan cco arena sizing, needed because this aiter has no MegaMoEGfx1250.required_vmm_bytes() (otherwise ccoMemAlloc failed: -1 at MoE init). #2522 itself is superseded by upstream #2384 (ATOM_MEGA_STAGE1_FUSED); aiter #5431 and #5844 are in the image. - Launch: no venv; MORI_AMD_SMI_PATH (mori dlopens libamd_smi.so, TheRock ships only .so.27, ranks exit 255); ATOM_MLA_QH128_ASM_PAGE_SIZE=0; ATOM_MEGA_STAGE1_FUSED=1; DSR1 decode fast path and pad-row mask off. - Results and known issues updated for this image; AgentX not re-measured. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

TritonMLAMetadataBuilder hands aiter's decode_attention_fwd a dense table of token-level KV row ids (page_size 1). The kernel multiplies each id it loads by the row stride (576 for MLA) in the id's own type, so with the int32 table the offset wraps for rows >= 2**31 / 576 = 3,728,270, which is 233,016 blocks of 16. The decode then reads outside the layer's KV: garbage, NaN, or a GPU memory fault. A pool reaches those rows once it has that many blocks and fills them. On Kimi-K3 dp8ep8 (gfx1250, 0.90, 16384 batched tokens) each rank holds about 253k-272k blocks, and with prefix caching the KDA checkpoint images (about 450 MB each) fill the pool past 90%. Decodes on every DP rank then went NaN, starting at the first token of a fresh block above 233,016, and answered `!!!!` up to max_tokens; all 64 corrupted blocks found were >= 233,329. Writes are unaffected: the cache kernels compute slots in int64. The table is now int64; csr_to_dense_block_table stores the int32 ids into it unchanged. Kernel repro on a 3.93M-row fp8 KV: an int32 table faults on rows past the limit, an int64 table matches a torch reference on rows below and above it.

Copilot

🟡 Changes recommended

The correctness fix lacks regression coverage for its dtype and overflow-boundary contract.

2 open findings

🧠 Review effort: Balanced

zejunchen-zejun · 2026-10-14T03:52:07Z

                self.max_bs,
                self.max_num_blocks_per_seq,
-                dtype=torch.int32,
+                dtype=torch.int64,


int64 covers the range that int32 express.

…, ATOM/aiter/mori on main The 1013 image ran ATOM, aiter and mori from the bring-up team's private MLPerf branches. The 1014 image switches all three to public main commits: ATOM 0adc645 and aiter 1305f0c8 (main of 2026-10-14), and mori 695d4c03, the last main commit before mori #745, whose interface aiter main does not use yet. The ROCm/torch/Triton/FlyDSL stack still comes from rocm/fw-bringup-mlperf:gfx1250-atom-20261012. - Patches, remade for ATOM main 0adc645: - atom-main-k3.diff: #2380, #2532, #2533 (4229c75). - atom-main-2522.diff: #2522 r3 (9cd096a). aiter #5431 and #5844 are merged in main. The r1_1012 patches are gone. - Launch: ATOM main's names. ATOM_MORI_V2_FUSED=1 binds MegaMoE (was ATOM_MEGA_STAGE2_FUSED), and MEGA_DISPATCH=flydsl picks the compact plan through #2522 (was ATOM_MEGA_STAGE1_FUSED). ATOM main does not read ATOM_MLA_QH128_ASM_PAGE_SIZE. ATOM_MEGA_MASK_PAD_ROWS stays 0. - Validated on CT11 (node0) and CT10, 2026-10-14: - full GSM8K strict-match 0.9659 +-0.0050, flexible 0.9651; - gate pass; prefix-cache probe pass; - 14k/500 c32: 477.9 output tok/s, median TPOT 50.5 ms, median TTFT 7.7 s. - The image digest is sha256:4da3dd98...8d8d. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.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>

Copilot

🟢 Approval recommended

The focused dtype change addresses the documented overflow without introducing identified blocking issues.

2 open findings

🧠 Review effort: Balanced

…355X reference points - Five of the image's eight PRs have merged since (#2532, #2533, #2466, aiter #6121) and #2522 was closed in favour of ATOM #2384; the table gains an upstream-state column. aiter #6438 (K3 grouped-MoE tuning, merged 10-15) is noted as not in the image. - Results: ATOM on MI355X from the InferenceX Kimi-K3 dashboard (2026-10-15 run), concurrency 16, 48 and 56: P90 interactivity (1 / p90 ITL) and throughput per GPU, with the setup differences (TP 8, MTP on). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

zejunchen-zejun marked this pull request as ready for review October 12, 2026 06:46

zejunchen-zejun requested review from XiaobingSuper and ganyi1996ppo and a balanced review from Copilot October 12, 2026 06:46

Copilot started reviewing on behalf of zejunchen-zejun October 12, 2026 06:47

View session

Copilot AI reviewed Oct 12, 2026

View reviewed changes

Comment thread tests/test_triton_mla_block_table_dtype.py Outdated

q = torch.randn(2, H, D, dtype=torch.bfloat16, device=dev)

scale = D**-0.5

table = torch.stack([low, high]).to(TRITON_BLOCK_TABLE_DTYPE)

XiaobingSuper previously approved these changes Oct 12, 2026

View reviewed changes

zejunchen-zejun requested a review from zufayu October 13, 2026 02:58

zejunchen-zejun mentioned this pull request Oct 13, 2026

[MegaMoE] FlyDSL dispatch on fp8/fp4 wires; Kimi-K3 defaults to the compact plan #2522

Closed

zejunchen-zejun dismissed XiaobingSuper’s stale review via 4229c75 October 14, 2026 03:35

Copilot AI balanced review requested due to automatic review settings October 14, 2026 03:35

zejunchen-zejun force-pushed the zejun/triton-mla-int64-block-table branch from 7351ef0 to 4229c75 Compare October 14, 2026 03:35

Copilot started reviewing on behalf of zejunchen-zejun October 14, 2026 03:36

View session

Copilot AI reviewed Oct 14, 2026

View reviewed changes

zejunchen-zejun added the ci:atom label Oct 14, 2026

zejunchen-zejun requested a review from valarLip October 15, 2026 05:36

Merge branch 'main' into zejun/triton-mla-int64-block-table

ef06615

Copilot AI balanced review requested due to automatic review settings October 16, 2026 01:31

Copilot started reviewing on behalf of zejunchen-zejun October 16, 2026 01:32

View session

Copilot AI reviewed Oct 16, 2026

View reviewed changes

valarLip approved these changes Oct 16, 2026

View reviewed changes

zejunchen-zejun merged commit cd3bf5f into main Oct 16, 2026

zejunchen-zejun deleted the zejun/triton-mla-int64-block-table branch October 16, 2026 03:51

Navigation Menu

Repository navigation

[MLA] Address Triton decode KV rows with int64 ids - #2533

[MLA] Address Triton decode KV rows with int64 ids#2533
zejunchen-zejun merged 2 commits into
mainfrom
zejun/triton-mla-int64-block-table

zejunchen-zejun commented Oct 12, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Oct 12, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

XiaobingSuper left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

zejunchen-zejun Oct 14, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

zejunchen-zejun commented Oct 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Oct 12, 2026

🏷️ CI Guide

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

XiaobingSuper left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

zejunchen-zejun Oct 14, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟢 Approval recommended

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

zejunchen-zejun commented Oct 12, 2026 •

edited

Loading