[gfx1250] feat(mla): unfused fallback for gather_kv_b_proj by zejunchen-zejun · Pull Request #2380 · ROCm/ATOM

zejunchen-zejun · 2026-09-24T02:22:58Z

The fused Triton gather_kv_b_proj does the whole MLA cached-prefix expansion in one launch. On some targets Triton has no working backend for the shapes a chunked prefill produces and codegen asserts inside LLVM before any kernel runs:

llvm/IR/Instructions.h: PHINode::setIncomingValue:
Assertion `getType() == V->getType()' failed.

That is a compile-time abort with no fallback -- a long prompt at concurrency takes the engine down -- and the FlyDSL gather is not an alternative there, being gfx950-only and aborting in its own compiler if forced. Every step of the chain (row gather, cache dequant, kv_b_proj, the k_nope/v split, the k_pe concat) is an ordinary torch op, so an unfused version needs no codegen and cannot miscompile.

Gated by ATOM_UNFUSED_GATHER_KV_B_PROJ, default off, and checked after FlyDSL: wherever a kernel backend compiles, it wins. This path trades throughput for existing.

It refuses rather than guesses on layouts private to the fused kernel -- the shuffled-KV cache, a preshuffled weight, MXFP4 kv_b_proj weights. Silently wrong K/V is worse than a clear failure, since nothing downstream inspects it. The one combination that can be known up front (unfused + ATOM_USE_TRITON_MLA_SHUFFLE_KV) raises at construction rather than from inside the first chunked prefill.

Tests cover the dispatch order and numerics against a per-token reference: page_size 1 and paged addressing, ragged tails, multiple and single-token sequences, chunk-size invariance, per-row and 128x128 block weight scales, and that a bf16 cache is not scaled by k_scale (the fused kernel's rule, and the failure no assertion downstream would catch).

The fused Triton gather_kv_b_proj does the whole MLA cached-prefix expansion in one launch. On some targets Triton has no working backend for the shapes a chunked prefill produces and codegen asserts inside LLVM before any kernel runs: llvm/IR/Instructions.h: PHINode::setIncomingValue: Assertion `getType() == V->getType()' failed. That is a compile-time abort with no fallback -- a long prompt at concurrency takes the engine down -- and the FlyDSL gather is not an alternative there, being gfx950-only and aborting in its own compiler if forced. Every step of the chain (row gather, cache dequant, kv_b_proj, the k_nope/v split, the k_pe concat) is an ordinary torch op, so an unfused version needs no codegen and cannot miscompile. Gated by ATOM_UNFUSED_GATHER_KV_B_PROJ, default off, and checked after FlyDSL: wherever a kernel backend compiles, it wins. This path trades throughput for existing. It refuses rather than guesses on layouts private to the fused kernel -- the shuffled-KV cache, a preshuffled weight, MXFP4 kv_b_proj weights. Silently wrong K/V is worse than a clear failure, since nothing downstream inspects it. The one combination that can be known up front (unfused + ATOM_USE_TRITON_MLA_SHUFFLE_KV) raises at construction rather than from inside the first chunked prefill. Tests cover the dispatch order and numerics against a per-token reference: page_size 1 and paged addressing, ragged tails, multiple and single-token sequences, chunk-size invariance, per-row and 128x128 block weight scales, and that a bf16 cache is not scaled by k_scale (the fused kernel's rule, and the failure no assertion downstream would catch). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

github-actions · 2026-09-24T02:23:36Z

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

The recipe described ATOM_UNFUSED_GATHER_KV_B_PROJ as a knob to set. On stock ATOM that env var does not exist, so following the page as written gets a server that starts, serves short prompts, and then dies on the first long prompt at concurrency -- the LLVM PHI assertion the patch exists to avoid. Say so where a reader will hit it: a row in the validated-configuration table, the prerequisites list, a dedicated section with the fetch command and the failure it prevents, an inline note on the env var itself, and the known-issues row. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

The page said to fetch the PR branch and check it out. That is the wrong instruction for this platform: the PR is based on main, while the gfx1250 image carries its own bring-up build of ATOM that is not main. Checking out the branch replaces that build and silently drops whatever bring-up deltas the image was made with, leaving you debugging a different engine than the one the recipe was validated against. Replace it with: take the diff, keep the image's ATOM. Fetch pull/2380.diff, filter it to the three runtime files (the other three are tests the image does not ship), resolve the installed package via atom.__file__ so it works for an editable checkout or a site-packages wheel, dry-run, then apply. Adds a verification step, because the failure is half silent: if envs.py alone fails to patch, ATOM_UNFUSED_GATHER_KV_B_PROJ is ignored at runtime with no complaint, which looks exactly like not having patched at all. Also notes this is container filesystem state, so docker rm discards it -- the same footgun the rocjitsu prefix has. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Copilot

Copilot review overview

🟡 Changes recommended

Multiple unresolved correctness issues and a remaining Triton code-generation dependency block approval.

Review effort: Lite
Findings: 4

Open (4)

Fallback still depends on Triton GEMM codegen · New
E8M0 scale bytes are applied without decoding · New
Packed MXFP4 weights bypass the unsupported-format guard · New
CPU KV scale causes device mismatch for FP8 cache · New

What changed in this PR

Adds an opt-in unfused Torch fallback for MLA cached-prefix KV projection when fused Triton compilation fails.

Changes:

Implements cache gathering, dequantization, projection, and K/V assembly.
Adds dispatch, layout guards, and environment configuration.
Adds numerical, dispatch, and environment tests.

File	Review summary
`tests/test_mla_unfused_gather.py`	Adds fallback numerical and layout coverage.
`tests/test_mla_flydsl_gather_dispatch.py`	Adds backend dispatch coverage.
`tests/test_envs.py`	Tests environment-variable defaults and parsing.
`atom/utils/envs.py`	Registers the feature flag.
`atom/model_ops/mla_unfused_gather.py`	Critical issues: still depends on Triton GEMM; incomplete MXFP4 rejection; incorrect E8M0 decoding; possible CPU/GPU scale mismatch. Also repeatedly dequantizes weights per chunk, causing a moderate performance issue.
`atom/model_ops/attention_mla.py`	Integrates fallback dispatch and layout validation.

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

+    try:
+        from aiter.ops.triton.gemm.basic.gemm_a16w16 import gemm_a16w16
+    except Exception:  # noqa: BLE001 -- optional kernel; torch is the fallback
+        return torch.matmul(x, w.t())
+    # gemm_a16w16 reads N from w.shape[0] and returns [M, N].
+    return gemm_a16w16(x, w, bias=None, dtype=x.dtype)


- Refuse a raw uint8 kv_b_proj weight. Where torch has no float4_e2m1fn_x2, packed MXFP4 stays uint8 and would otherwise be read as ordinary weights, halving kv_c_dim and producing plausible but wrong K/V. - Decode uint8 weight scales as E8M0 biased exponents. aiter's fp8_e8m0 falls back to uint8 where torch lacks float8_e8m0fnu, and a plain cast would read exponent 127 as 127 instead of 1.0. - Make k_scale a 0-dim device scalar once before the chunk loop, so a [1]-shaped host scale neither fails against CUDA operands nor promotes the bf16 products to fp32. - black. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Copilot

Copilot review overview

🟡 Changes recommended

A critical CPU-path issue and per-chunk weight dequantization must be addressed before approval.

Review effort: Lite
Findings: 2

Open (2)

CPU numerics incorrectly dispatch to GPU-only AITER GEMM · New
Fallback still depends on Triton GEMM codegen

Resolved since last review (3)

CPU KV scale causes device mismatch for FP8 cache
Packed MXFP4 weights bypass the unsupported-format guard
E8M0 scale bytes are applied without decoding

+    a16w16 GEMM is correct here and much faster on this shape. Called directly
+    rather than through ``tgemm.mm`` so a tuned-CSV row cannot route it back.
+    """
+    try:


Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Copilot

Copilot review overview

🔵 Needs a closer look

Unresolved fallback GEMM, scale-shape, repeated dequantization, documentation, and test-coverage issues remain.

Review effort: Lite
Findings: 2

Open (2)

CPU numerics incorrectly dispatch to GPU-only AITER GEMM
Fallback still depends on Triton GEMM codegen

…triage Measured on a 4-node gfx1250 rack (16 GPUs, 432 GiB boards) with rocm/fw-bringup:gfx1250-atom-20260918-ep8. - B0→A0 translation is image-dependent, not silicon-dependent. That image's kernels are already built for B0: the hook installs on all 16 ranks and performs zero translations, and GSM8K still scores 0.9621 +/-0.0053 over the full 1319. Adds an A0-image vs B0-image table and a 60-second test to tell the two apart, since no revision field on the board reports the stepping. - HBM per GPU is board-dependent (288 and 432 GiB both exist); every KV number now says to re-derive it. available_for_kv measured at 159 GB @0.90. - accel_state can stick at `unconfigured` permanently and look like "still training"; records the recovery and the ifoe ordering hypothesis. - Verify the patch and prefix by byte size: a 0-byte .py from an interrupted docker cp imports cleanly and passes every check in the current text. - PR #2380's envs.py hunk does not apply to this image's ATOM; says where to put the block by hand. - Adds the fabric-down signatures seen while accel_state read `active` (HELLO/IMPORT/XA, p2p_tmp.cc:372) and notes the exporter-side dmesg is empty. - Adds a minimal torch.distributed all_reduce to separate a platform fault from the intermittent ncclCommInitRank hang, plus py-spy is not in the image. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…nd B0 Corrects an error introduced earlier in this branch. There are not two images: A0 and B0 machines run the SAME image, rocm/fw-bringup:gfx1250-atom-20260918-ep8, and both need ATOM PR #2380. The image ships B0 code objects -- which is why the library is librocjitsu_gfx1250_b0_to_a0.so -- so translation is what A0 silicon needs to run them, and B0 needs nothing. The previous text described this as "A0-built vs B0-built images", which was wrong. Restructures the page around the B0 path, which is the one with numbers: - New opening states the one difference, points at the library name as the reason, and adds a "where to go" map -- the page had no way to tell where the B0 material started. - The translation chapter moves out of position two to Appendix B, "running the same image on A0 silicon", opening with "only read this if your parts are A0". Keeps the wrong-prefix table, install steps and verification intact, and keeps the behavioural test for telling the two apart, since no revision field on the board reports the stepping. - The validated configuration becomes Appendix A and is linked from the top. - Headline numbers now include the fixed-length 14k/16 con32 point: 26,207 tok/s total, TTFT median 13.2 s, TPOT median 73.9 ms. All 18 internal links verified against GitHub's anchor rules. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

XiaobingSuper

LGTM

Validated full K3 at dp8ep8 on CT2 + CT3 with rocm/atom-dev 1009 (ROCm 10.2, torch 2.14, Triton 3.9, ATOM/aiter/mori main of 2026-10-08) plus ATOM #2380 and aiter #5431: the accuracy gate passes and GSM8K scores strict-match 0.9575 +-0.0056. - New section: the image's contents and layout, how to apply the two PRs to its wheel install, the env/CLI delta from Config A, the full launch, and the results. - Known issues: MORI_EP_TOKOFF_EXT hipIpcOpenMemHandle hipError 17 on more than one node (aiter turns on mori's single-host TokOffExt by default); on newer aiter the symm-mem custom all-reduce needs #5431. - Optimized source stack: the TokOffExt note was inferred, now observed. - Fixes: the '!' known issue had A0 and B0 backwards; the Launch env comment now says the translation lines are for A0 only. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

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

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

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

zejunchen-zejun marked this pull request as ready for review September 28, 2026 03:04

Copilot AI lite review requested due to automatic review settings September 28, 2026 03:04

Copilot started reviewing on behalf of zejunchen-zejun September 28, 2026 03:05

View session

Copilot AI reviewed Sep 28, 2026

View reviewed changes

Copilot AI review requested due to automatic review settings September 28, 2026 05:56

Copilot started reviewing on behalf of zejunchen-zejun September 28, 2026 05:57

View session

Copilot AI reviewed Sep 28, 2026

View reviewed changes

Comment thread atom/model_ops/mla_unfused_gather.py

a16w16 GEMM is correct here and much faster on this shape. Called directly

rather than through ``tgemm.mm`` so a tuned-CSV row cannot route it back.

"""

try:

zejunchen-zejun requested review from gbyu-amd and zufayu September 28, 2026 06:04

zejunchen-zejun changed the title ~~feat(mla): unfused torch fallback for gather_kv_b_proj~~ [gfx1250] feat(mla): unfused torch fallback for gather_kv_b_proj Sep 28, 2026

zejunchen-zejun changed the title ~~[gfx1250] feat(mla): unfused torch fallback for gather_kv_b_proj~~ [gfx1250] feat(mla): unfused fallback for gather_kv_b_proj Sep 28, 2026

test(mla): trim unfused gather_kv_b_proj tests to the core cases

dc83d9e

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Copilot AI review requested due to automatic review settings September 28, 2026 06:06

Copilot started reviewing on behalf of zejunchen-zejun September 28, 2026 06:07

View session

Copilot AI reviewed Sep 28, 2026

View reviewed changes

zejunchen-zejun requested a review from XiaobingSuper September 28, 2026 06:47

XiaobingSuper approved these changes Sep 29, 2026

View reviewed changes

This was referenced Oct 2, 2026

recipe(kimi-k3): add MI455 B0 wide-EP optimization workflow #2446

Closed

[Triton/Gluon] [MLA] Fix fused gather_kv_b_proj on gfx1250 ROCm/aiter#6120

Closed

XiaobingSuper closed this Oct 12, 2026

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

Navigation Menu

Repository navigation

[gfx1250] feat(mla): unfused fallback for gather_kv_b_proj - #2380

[gfx1250] feat(mla): unfused fallback for gather_kv_b_proj#2380
zejunchen-zejun wants to merge 3 commits into
mainfrom
mla-unfused-gather-kv-b-proj

zejunchen-zejun commented Sep 24, 2026

Uh oh!

github-actions Bot commented Sep 24, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

XiaobingSuper left a comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

zejunchen-zejun commented Sep 24, 2026

Uh oh!

github-actions Bot commented Sep 24, 2026

🏷️ CI Guide

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🔵 Needs a closer look

Uh oh!

XiaobingSuper left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants