[ASM] [HIP] [JIT] MI355 MLA PS mode kernel add 4 wave fp8 kernel by minmengdie · Pull Request #6455 · ROCm/aiter

minmengdie · 2026-10-13T02:03:58Z

Motivation

On gfx950 the opus MLA decode path for a8w8 (fp8 Q / fp8 KV, merged d=576 buffer, per-tensor
descales) had a single kernel (16mx8_32nx1) that did not cover the shapes the asm persistent
kernels serve and was slower than asm on most of them. This PR adds three a8w8 opus kernels whose
work tiling matches the asm heuristic in csrc/py_itfs_cu/asm_mla.cu, routes every shape the asm
path serves to one of them, and makes the opus path feature-complete for production decode:

nhead 12 (native, no head fold) and the folded head counts (e.g. 48, 80, 96 x q > 6);
causal multi-token decode (MTP / speculative, max_seqlen_q > 1) with split-KV;
round-robin context parallel (DCP), matching the asm cprr kernels;
a single public entry point opus_mla_decode_fwd for bf16 and a8w8.

With AITER_MLA_USE_OPUS=1, opus is on par with or faster than asm on every benchmarked shape
(geomean 1.17x over 144 configs, up to 1.51x), with accuracy at the same level as asm.

Technical Details

New kernels (csrc/kernels/mla/opus/)

kernel	tiling	rows / workgroup	build
`mla_decode_a8w8_32mx4_64nx1`	4 waves x 32 query rows, 64-token KV tiles, 3-slot LDS ring	128 (asm qh32 / qseqlen4)	prebuilt `.co` (register-pinned, amdgpu-pin-op-dst toolchain)
`mla_decode_a8w8_32mx2_32nx2`	2 wave pairs x 32 rows, each wave of a pair takes 32 of the 64 tokens; pair-merged softmax max, 3-slot LDS ring	64 (asm qh16 / qseqlen4)	prebuilt `.co`
`mla_decode_a8w8_16mx1_32nx4`	nhead 16 x q 1: 4 waves split the KV tile 32 tokens each, Q shared through LDS, 2-slot ring	16	JIT (`module_mla_decode_opus`)

Prebuilt code objects: hsa/gfx950/mla_opus/opus_mla_decode_a8w8_{32mx4_64nx1,32mx2_32nx2}.co,
built by csrc/kernels/mla/opus/build/build_mla_decode_a8w8_*_co.sh (hipcc failures now fail
the script). Entries opus_mla_decode_a8w8_<kernel>_{nc,causal}[_large][_cprr]; _large
addresses KV with flat 64-bit loads when the cache is >= 4 GiB.
Routing (opus_a8w8_route in opus_mla_decode_fwd.cu), on nhead per KV head after mla.py's
head fold:
- 32mx4: nhead 32 & q >= 4, 64 & q >= 2, 128, 96 & q <= 6, 16 & q > 4, 12 & 12q <= 128;
- 32mx2: nhead 16 & q in {3, 4}, 32 & q in {2, 3}, 64 & q == 1;
- 16mx1: nhead 16 & q == 1.
Online softmax rescale threshold: 0 (exact running max) for 32mx2 / 16mx1; 8 (log2 units) for
32mx4, which keeps ~13% of its performance and stays within the test tolerance.

Test Plan

All on gfx950 (MI355X), AITER_MLA_USE_OPUS=1 for opus, AITER_MLA_USE_OPUS=0 for the asm
baseline, idle GPUs.

Accuracy, default -b 1 3 5 16 32 64 128 256 -c 21 64 256 512 1200 3200 5200 8192:
```
python3 op_tests/test_mla_persistent.py -d fp8 -kvd fp8 -n 12,8 16,1 16,4 16,8 16,17 32,3 32,4 64,1 64,2 96,3 96,8 128,1
```
criterion: checkAllclose(atol=rtol=0.01) of golden fp8 vs aiter_asm and of the split-KV
output (failed = more than 5% of elements out of tolerance); same run with asm for reference.

Performance vs asm (opus and asm concurrently on two idle GPUs, fixed seed, min of 3 rounds):

python3 op_tests/test_mla_persistent.py -d fp8 -kvd fp8 -n 12,8 16,1 32,2 32,4 64,1 64,2 96,1 96,4 128,2 -b 16 32 64 128 -c 4096 8192 16384 65536

Round-robin CP, full default sweep (8 shapes x 9 ctx x 5 batch x W in {2,3,4,7,8}):
```
python3 op_tests/test_mla_persistent_round_robin.py -d fp8 -kvd fp8 -lse
```
criterion: per-rank cp_ref vs aiter for out and LSE, no NaN.
Regression: bf16 opus (-d bf16 -kvd bf16 -n 16,1 64,1 128,1), non-CP causal routes after the
CP change, and the cprr / plain entry selection checked through AITER_MLA_OPUS_TRACE=1.

Test Result

1. Accuracy — 12 shapes x 64 configs = 768 configs: 0 failed for opus (final output and
split-KV output). Max mismatch ratio per shape is at the level of asm:

nhead,q	opus kernel	opus passed / warning / failed	opus max mismatch	asm max mismatch
12,8	32mx4	32 / 32 / 0	1.75%	1.67%
16,1	16mx1	38 / 26 / 0	3.28%	3.49%
16,4	32mx2	34 / 30 / 0	2.71%	2.25%
16,8	32mx4	35 / 29 / 0	1.84%	1.71%
16,17	32mx4	28 / 36 / 0	1.07%	1.02%
32,3	32mx2	37 / 27 / 0	0.88%	0.88%
32,4	32mx4	31 / 33 / 0	1.90%	2.28%
64,1	32mx2	38 / 26 / 0	0.75%	0.74%
64,2	32mx4	34 / 30 / 0	0.86%	0.87%
96,3	32mx4	28 / 36 / 0	0.91%	0.85%
96,8 (folded)	32mx4	25 / 39 / 0	1.60%	1.60%
128,1	32mx4	30 / 34 / 0	2.52%	2.86%

2. Performance — asm time / opus time (> 1: opus faster), 144 configs, geomean 1.173x:

nhead,q	opus kernel	配置数	opus 几何平均 (us)	asm 几何平均 (us)	asm/opus 几何平均	最小比值	最大比值	opus 更慢的配置数	opus TFLOPS 几何平均 / 最大	asm TFLOPS 几何平均 / 最大	opus TB/s 几何平均 / 最大	asm TB/s 几何平均 / 最大
12,8	32mx4_64nx1	16	130.3	159.1	1.221x	1.122x (b16 c4096)	1.342x (b128 c65536)	0	999.2 / 1542.7	818.7 / 1206.6	2.84 / 4.27	2.32 / 3.34
16,1	16mx1_32nx4	16	90.3	91.8	1.016x	0.996x (b64 c4096)	1.060x (b16 c4096)	2	240.4 / 320.4	236.5 / 320.0	4.00 / 5.31	3.93 / 5.30
32,2	32mx2_32nx2	16	108.1	115.5	1.068x	1.016x (b16 c8192)	1.105x (b128 c4096)	0	802.9 / 1251.7	751.7 / 1149.2	3.39 / 5.19	3.17 / 4.77
32,4	32mx4_64nx1	16	143.1	172.1	1.203x	1.095x (b32 c4096)	1.322x (b128 c65536)	0	1213.3 / 1963.1	1008.9 / 1542.2	2.61 / 4.08	2.17 / 3.21
64,1	32mx2_32nx2	16	107.3	114.3	1.066x	1.006x (b16 c4096)	1.101x (b16 c65536)	0	809.6 / 1238.8	759.4 / 1150.5	3.41 / 5.14	3.20 / 4.77
64,2	32mx4_64nx1	16	140.7	171.0	1.215x	1.118x (b16 c4096)	1.307x (b128 c65536)	0	1234.4 / 1952.6	1015.7 / 1550.8	2.65 / 4.06	2.18 / 3.22
96,1	32mx4_64nx1	16	130.4	159.3	1.222x	1.142x (b16 c4096)	1.334x (b128 c65536)	0	998.9 / 1533.9	817.7 / 1212.0	2.83 / 4.25	2.32 / 3.36
96,4	32mx4_64nx1	16	328.0	432.1	1.317x	1.256x (b32 c4096)	1.507x (b128 c65536)	0	1588.3 / 1975.0	1205.8 / 1464.8	1.22 / 1.39	0.93 / 1.09
128,2	32mx4_64nx1	16	207.1	262.5	1.267x	1.175x (b16 c4096)	1.424x (b128 c65536)	0	1676.7 / 2289.4	1323.1 / 1739.2	1.87 / 2.42	1.47 / 1.87
全部		144			1.173x			2

opus is slower on 2 of 144 configs (16,1 at 0.996x and 0.999x).

3. Round-robin CP — 1800 configs, 17280 per-rank checks (out + LSE): 0 failed, 0 NaN; every
causal shape loads the _causal_cprr entry, 64,1 (q = 1) the plain _nc entry.

Note on golden vs *_merge out for fp8: it fails for the test's own CP reference as well as for
asm's fp8 cprr kernel, because the reference quantizes P to fp8 against the rank-local max while
golden uses the global max. opus fails it on exactly the same configs as the reference (318 / 318);
LSE vs golden: 0 failed.

4. Regression — bf16 opus passes; non-CP causal routes pass and select the plain entries
(cprr=0); same-GPU int

Submission Checklist

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

github-actions · 2026-10-13T02:04:28Z

🏷️ CI Guide

Runs automatically on every PR:

✅ Pre-checks (submodule verification, code formatting)
✅ Aiter op tests (gfx942 + gfx950)
✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
`ci:gfx1250-ffm-triton`	Run the five-shard gfx1250 FFM Triton test suite
`ci:triton-300x`	Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
`ci:triton-355`	Run the full Triton test suite on MI35X, not only the tests the change affects
`ci:windows`	Windows JIT op smoke test on the gfx1201 ROCm runner
`ci:multi-gpu`	Aiter multi-GPU tests on the 8-GPU runner
`ci:sglang`	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
`ci:atom`	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
`ci:atom_full`	ATOM accuracy suite for PR and main models from ATOM `models_accuracy.json`
`ci:vllm`	vLLM benchmark: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
`ci:all`	All standard extended tests (excludes `ci:atom_full`)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 6455 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

Copilot

🟡 Changes recommended

Masking constraints, unsafe pointer construction, masked build failures, and missing automated coverage remain unresolved.

6 open findings

What changed in this PR

Adds gfx950 wave4 FP8 MLA decode kernels, prebuilt code-object dispatch, and shape-based routing.

Changes:

Adds 32mx2/32nx2 and 16mx1/32nx4 kernels and traits.
Adds code-object wrappers and build scripts.
Routes supported FP8 shapes through the new kernels.

File	Description
`mla_global_load.hpp`	Adds fixed-width integer support.
`mla_decode_traits.h`	Adds kernel arguments and variant traits.
`mla_decode_fp8_32mx4_64nx1_co.hip`	Defines 32mx4 code-object entries.
`mla_decode_fp8_32mx2_32nx2.hpp`	Implements the 32mx2 kernel.
`mla_decode_fp8_32mx2_32nx2_co.hip`	Defines 32mx2 code-object entries.
`mla_decode_fp8_16mx1_32nx4.hpp`	Implements the 16mx1 kernel.
`build_mla_decode_fp8_32mx4_64nx1_co.sh`	Builds the 32mx4 code object.
`build_mla_decode_fp8_32mx2_32nx2_co.sh`	Builds the 32mx2 code object.
`opus_mla_decode_fwd.cu`	Adds routing and launch logic.
`aiter/mla.py`	Gates Opus shapes and updates HK dispatch.
`optCompilerConfig.json`	Updates JIT source dependencies.

🧠 Review effort: Balanced

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

Copilot

🟡 Changes recommended

Unsafe assembly constraints, masked build failures, compatibility breakage, and missing kernel-path tests remain unresolved.

7 open findings

🧠 Review effort: Balanced

Copilot

🟡 Changes recommended

The FP8 gate can route unsupported shapes to Opus, removes a public API without compatibility, and lacks automated route coverage.

7 open findings

2 resolved since last review

🧠 Review effort: Balanced

Copilot

🟡 Changes recommended

The Python dispatch admits unsupported FP8 configurations, removes an exported API without compatibility handling, and lacks coverage for the new kernels.

7 open findings

1 resolved since last review

Avoid partial-output pointer arithmetic for final work items

🧠 Review effort: Balanced

Copilot

🔵 Needs a closer look

The FP8 gate can route unsupported inputs into errors, removes a public API without compatibility, and lacks CI coverage for the new kernels.

7 open findings

Previously missed (1)

In code that hasn't changed since last review

Add CI coverage for all Opus FP8 dispatcher routes

csrc/kernels/mla/opus_mla_decode_fwd.cu:218

No op_tests test enables AITER_MLA_USE_OPUS, so CI never exercises this new fp8 dispatcher or any of its three routes; the existing persistent MLA tests continue through the ASM/HK paths. These kernels contain independent causal masking, split-KV reduction, and large-cache addressing logic, so regressions can ship undetected. Add gfx950 correctness cases that activate each route and cover causal/non-causal plus final/split work items.

🧠 Review effort: Balanced

Copilot

🔵 Needs a closer look

The Opus gate accepts unsupported or scaleless FP8 inputs, removes a public symbol, and lacks correctness coverage for the new routes.

7 open findings

🧠 Review effort: Balanced

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

kaiyang-1

LGTM

…4 rule Review follow-ups on get_mla_decode_shape_support, so the query answers what mla_decode_fwd does for the same inputs: - New cp_world_size keyword (default 1, ValueError below 1). The planner used to get 2 if cp_round_robin else 1, while mla_decode_fwd passes the real value, so CP > 1 without g_kv_indptr showed the gfx1250 FlyDSL PS1 route that the launch rejects. The PS1 asm kernel lookup now keys cprr on cp_world_size > 1 as mla_ps1_fp8_asm_fwd does; stage-1 asm keeps keying it on g_kv_indptr (cp_round_robin). - HK: ROCm#6455 made the route real (aiter.hk_mla_v32_decode_fwd), so it is no longer reported as undefined. It answers None, or False when the launch would be wrong: return_lse (no final LSE for unsplit work), causal=False at max_seqlen_qo > 1 (always masked) and round-robin CP (no g_kv_indptr). The reasons note that HK ignores q_scale/kv_scale. - gfx1250 stage-1 asm needs page_size == 64. The rule moves into mla_asm_decode_remap_gfx1250 (NoKernel, so the dispatch keeps its own message), and mla_decode_asm_query takes page_size. - gfx1250 A0: stage-1 asm routes answer False (require_gfx1250_asm raises). - The gfx1250 LSE reason says LSE is valid only for gqa64 or kv_split > 1. - Docstrings: only kv_layout="token_major" is described; fp8 Q is assumed to come with scales; persistent=True is the persistent route, and gfx950 bf16 16-head qlen-1 decode can still downgrade at large batch. Tests (op_tests/test_mla_shape_support.py, host-only): query vs the head plan mla_decode_fwd actually requests, captured before launch, at cp_world_size 1/4 with and without round-robin; PS1 cprr kernel choice; HK verdicts; gfx1250 page size. Each new test fails when its fix is reverted. Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…acked by shared C++ head-fold/backend rules (#6107) * [MLA] A1: move decode head-count planning into mla_decode_shape.h Add the host-only pure header csrc/include/mla_decode_shape.h with MlaPlanner, MlaHeadPlan, MlaPlannerFlags, MlaDecodeHeadPlan and mla_decode_head_plan() (per-planner v1_0/v1_1/v1_2 bodies). The bodies are literal moves of the inline head-fold / native-support / acceptance logic from the v1_0, v1_1 and v1_2 metadata planners, which now call the header; get_gpu_arch() and the env reads stay at the call sites and the AITER_CHECK messages are unchanged. mla_metadata_cluster_multiplier also moves into the header. The v1_2 text is moved as of #5844, whose gfx1250 FlyDSL PS1 clause also covers 96 heads at any max_seqlen_qo. No behavior change: planner/reduce/decode golden capture on gfx950 is bit-identical to the pre-change baseline. Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [MLA] A2: move the mla_reduce (heads, head_dim) table into mla_reduce_cases.h Move the _AITER_MLA_REDUCE_FOR_EACH_ROW support matrix (#5481) out of mla_reduce_kernels.cuh into the preprocessor-only csrc/include/mla_reduce_cases.h, with the same 24 rows in the same order; mla_reduce_kernels.cuh includes it, so MLA_REDUCE_ROUTER, the extern template block and reduce_instances/ are unchanged. The one edit to the moved text: the comment's "the extern template block below" now reads "the extern template block in mla_reduce_kernels.cuh", so it stays true in the new file. Add mla_reduce_v1_supports(num_heads, head_dim) to mla_decode_shape.h. It expands the same table, and it does not check dtype or work-group count. Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [MLA] A3: factor the ASM decode remap into mla_asm_decode_config Move the stage1 head/qseqlen remap into a pure mla_asm_decode_config(arch, q_type, kv_type, num_heads, nhead_kv, max_seqlen_q, persistent, causal, has_lse, has_cprr) that returns an MlaAsmDecodeCfg with a status (Ok=0, RemapReject=1, NoKernel=2, ArchNotBuilt=3), a reason and the table key. The remap block is moved literally; its three AITER_CHECKs become RemapReject with the same text. gfx1250 gets its own remap (persistent/lse/cprr -> NoKernel; the supported_variant list moved verbatim). Add try_get_heuristic_kernel_mla() returning "" on miss; get_heuristic_kernel_mla wraps it with the identical AITER_CHECK. Stage1 and the gfx1250 dispatch call the new function and throw the same messages as before. The provably dead stage1 "cannot find suitable kernel" check is removed. The cfg_mla_dsl lookup of mla_ps1_fp8_asm_fwd (#5844) moves verbatim into mla_ps1_fp8_asm_config(), which returns nullptr on a miss; the entry point keeps its AITER_CHECK. #6132's gfx950 non-persistent fp8 gqa64 qlen1 clause moves with the remap; its total_q * kv_split <= 32768 check needs launch sizes, so the remap records the bound in MlaAsmDecodeCfg::max_q_x_kv_split and stage1 raises the same AITER_CHECK right after the remap. Notes: status is enum class MlaAsmStatus : int (doc types it as int; a ctypes packer must cast). A gfx1250 variant miss is RemapReject and gfx1250 has_cprr is NoKernel (both unspecified in the doc). causal is now bool, so a raw caller passing a value other than 0/1 gets the causal kernel instead of a miss. The equally dead gfx1250 "cannot find suitable gfx1250 kernel" check is kept. No behaviour change: golden compare.py 0 unexpected diffs (gfx950), device-free old-vs-new sweep identical for gfx942/gfx950/gfx1250. Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [MLA] A4: expose the decode shape query and drive mla_decode_fwd from it Add get_mla_decode_head_plan_v1 / mla_reduce_v1_supports pybinds in module_mla_metadata, the mla_decode_asm_query ctypes entry in module_mla_asm (bound with ctypes.CDLL, NULL stream, device-free with an explicit arch), and the public aiter.get_mla_decode_shape_support query (design doc §5.2-5.4). The CDLL path is build_ctypes_module, factored out of _ctypes_call with main's JIT_LIB_EXT (.pyd on Windows); the entry points export via AITER_C_ITFS. mla_decode_fwd's native/fold gate and FlyDSL/opus/HK backend selection now come from the same C++ helpers the planner uses (C++ arch, atoi env parse), with a new fast_mode keyword (default True, today's behaviour) so the fold follows the planner mode the caller ran. Intended changes: >128-head folds (fp8 144 now runs), fp16, the gfx1250 fallthrough, the arch source, and "2"-style env values. A byte KV buffer (3-buffer layout) is planned as fp8 like its callers plan it; non-native fp8 variants and byte KV always take asm. is_experimental_enabled and the PS1 FlyDSL parses use atoi semantics. AITER_MLA_USE_OPUS now uses atoi semantics too (previously only "1"). also takes 96 heads at any max_seqlen_qo and round-robin CP (new cp_round_robin input), and 96/128 heads go to the new MlaBackend.PS1_FP8_ASM (mla_ps1_fp8_asm_fwd) unless AITER_MLA_DECODE_PS1_ASM (atoi, default 1) is 0; the query checks that route against cfg_mla_dsl through the new mla_ps1_fp8_asm_query entry. The query reports a FlyDSL PS1 or PS1 asm route as unsupported when the planner folds the heads (fast_mode=False), since those kernels read unfolded heads; mla_decode_fwd's dispatch is unchanged there. With persistent=False the planner and reducer verdicts do not gate supported, since that path never plans; launch-size bounds such as #6132's total_q * kv_split <= 32768 are not checked. Adds op_tests/test_mla_shape_support.py (device-free planner verdict, asm LSE, non-persistent asm and gfx1250 PS1 backend checks, GPU planner and decode agreement) and passes fast_mode at test_mla_persistent's mla_decode_fwd call sites. Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [MLA] A5: size metadata info and occupancy from the C++ planner rules get_mla_metadata_info_v1 (dense path) now takes qk_batch_ratio, kernel_num_heads and packed_qo_len_per_wg from the C++ head plan (get_mla_decode_head_plan_v1) instead of its Python mirror of the planner. For now each buffer is the max of the plan-derived size and the former mirror size, so no size shrinks on any shape (doc 5.5 rows 3-4; the shrink the doc's 2.4 row 6 allows is left to a follow-up). The sizes grow where the mirror under-sized the planner (v1_0/v1_1 folds, fp16 at msq 5-7). The num_head_qo % 4 assert now fires only where the planner rejects the shape, so N=10 bf16/fp16 fast_mode is sized instead of asserting. The sparse path stays on the moved mirror (_mla_legacy_max_qo_tiles_per_batch, which still resolves the arch with get_gfx()); it could not be proven to match the head plan. #5844's gfx1250 FlyDSL PS1 96-head sizing clause is honored by the plan-based dense sizing (the head plan makes 96 heads native, so max_seqlen_qo tiles per batch); the moved mirror keeps the clause, with the atoi parse, for the floor and the sparse path. get_mla_decode_fwd_occupancy / _max_splits call the new pybind mla_metadata_cluster_multiplier_v1 with the C++ arch and env parse, and gain mla_version and rope dtype parameters with backward-compatible defaults. mla_decode_fwd passes fp8 for a byte KV buffer, the dtype the planner is given. Separate fix: with fast_mode and intra_batch_mode both on, the dense work_info_set and reduce_partial_map are now at least the v1_2 (non-intra) bounds; tile_cnt * num_kv_splits under-counts the v1_2 planner, which num_kv_splits does not cap. The sparse path has the same under-count and is left unchanged. Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [MLA] A6: take decode_update_mla_metadata_v1's head fold from the planner decode_update_mla_metadata_v1 kept a hand-written copy of the v1_2 head fold rule, resolved with get_gfx() and an exact dtypes.fp8 check, which had drifted from metadata.cu. Where it disagreed it mapped works to the wrong request (e.g. gfx950 bf16 N=32: request 1's KV range became request 0's) without raising. It now asks the C++ head plan (get_mla_decode_head_plan_v1 via _mla_metadata_head_plan: C++ get_gpu_arch(), atoi env parse, original max_seqlen_qo, fast_mode=True) for qk_batch_ratio and the kernel head count. Changelog (behavior change for external callers): - The fold now matches what get_mla_metadata_v1 (fast_mode=True) planned. Shapes no longer folded: gfx950 fp8 N=32 at max_seqlen_qo != 4, gfx950 fp8 N=64, gfx950 bf16/fp16 at any N > 16 with N % 16 == 0, gfx942 fp8 N=64 at max_seqlen_qo == 1, and the AITER_ENABLE_EXPERIMENTAL HK (hk_mtp_experimental) shapes. - gfx1250 fp8 N=96 with AITER_MLA_DECODE_PS1_FLYDSL is no longer folded either: #5844 made it native in the planner but not in this copy. - float8_e4m3fnuz on gfx950 now counts as fp8, as in the planner. - Callers must pass the max_seqlen_qo they planned with (>= 1, before the qlen > 1 -> 1 rewrite; now asserted) and the dtype_q_nope / dtype_kv_nope given to get_mla_metadata_v1. The commented-out caller in op_tests/test_mla_persistent.py that passes max_seqlen_qo=1 for decode_qlen > 1 is wrong under this contract. - vLLM main (7e583e615c) does not call this function; it only uses get_mla_metadata_v1. - The unused batch_size local is gone; the signature is unchanged. Known limitations: only v1_2 (fast_mode=True) metadata is handled, as before (v1_0's N=128 seqlen fold is not undone). The arch is resolved on torch.cuda.current_device(), not the metadata tensors' device. Test (op_tests/test_mla_shape_support.py): on the GPU the update leaves the planner's work_info request/KV columns unchanged when no token is added. It fails on the previous rule. Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [MLA] Group mla_reduce_v1_supports with the compile_ops shape-query ops mla_reduce_v1_supports is a pybind op bound through compile_ops, but in NONE_WRAPPED_OP it sat after the ctypes.CDLL-bound mla_decode_asm_query. Drop mla_decode_asm_query from the list: it is a plain Python function, not a compile_ops op, so its entry was never consulted. The compile_ops shape-query entries now form one group. Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [MLA] Shape query: take cp_world_size, model HK and the gfx1250 page64 rule Review follow-ups on get_mla_decode_shape_support, so the query answers what mla_decode_fwd does for the same inputs: - New cp_world_size keyword (default 1, ValueError below 1). The planner used to get 2 if cp_round_robin else 1, while mla_decode_fwd passes the real value, so CP > 1 without g_kv_indptr showed the gfx1250 FlyDSL PS1 route that the launch rejects. The PS1 asm kernel lookup now keys cprr on cp_world_size > 1 as mla_ps1_fp8_asm_fwd does; stage-1 asm keeps keying it on g_kv_indptr (cp_round_robin). - HK: #6455 made the route real (aiter.hk_mla_v32_decode_fwd), so it is no longer reported as undefined. It answers None, or False when the launch would be wrong: return_lse (no final LSE for unsplit work), causal=False at max_seqlen_qo > 1 (always masked) and round-robin CP (no g_kv_indptr). The reasons note that HK ignores q_scale/kv_scale. - gfx1250 stage-1 asm needs page_size == 64. The rule moves into mla_asm_decode_remap_gfx1250 (NoKernel, so the dispatch keeps its own message), and mla_decode_asm_query takes page_size. - gfx1250 A0: stage-1 asm routes answer False (require_gfx1250_asm raises). - The gfx1250 LSE reason says LSE is valid only for gqa64 or kv_split > 1. - Docstrings: only kv_layout="token_major" is described; fp8 Q is assumed to come with scales; persistent=True is the persistent route, and gfx950 bf16 16-head qlen-1 decode can still downgrade at large batch. Tests (op_tests/test_mla_shape_support.py, host-only): query vs the head plan mla_decode_fwd actually requests, captured before launch, at cp_world_size 1/4 with and without round-robin; PS1 cprr kernel choice; HK verdicts; gfx1250 page size. Each new test fails when its fix is reverted. Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [MLA] gfx1250 stage 1: reject round-robin CP instead of running unmasked mla_decode_gfx1250_dispatch takes no g_kv_indptr, so a round-robin CP call ran without the global-position mask and returned wrong results. It now fails with the same check mla_asm_decode_config (and so the shape query) already applies. Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [MLA] test_mla_persistent: add --metadata_mode to cover planner v1_1 -nps selects fast_mode=False, intra_batch_mode=True (v1_0), and nothing ran v1_1 (both flags off). -mm/--metadata_mode {fast,v1_0,v1_1} picks the planner for get_mla_metadata_info_v1, get_mla_metadata_v1 and the mla_decode_fwd calls; the default stays v1_0 with -nps, else fast. Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [MLA] Add a numerical regression for decode_update_mla_metadata_v1 The function's block in test_mla_persistent.py is commented out, so nothing checked the metadata it rewrites. The new test plans get_mla_metadata_v1 (fast_mode) for 4 requests with different KV lengths (1100, 37, 530, 2049), moves them to new lengths with decode_update_mla_metadata_v1 (+1 each, and +1/0/-1/0 through num_reject_tokens), and decodes with the updated plan: - fp8/fp8 and bf16/bf16 at 16, 32, 48 (folded on gfx950 fp8), 64 and 128 heads; planned max_seqlen_qo 1 and 2 (the qlen > 1 -> 1 rewrite); the parallel and the serial (AITER_MLA_META_USE_PARALLEL=0) planner fill; - the updated works must tile every (pseudo-)request's new KV range, and at least one request must span several splits; - O and LSE are compared with a fresh plan at the new lengths and with an fp32 torch reference. On MI355X: 46 passed, 28 skipped (shapes the query does not report as supported, serial fill on folded or qlen > 1 plans, and qlen-2 plans that split per token), 6 xfailed. The xfails are fp8 with 32 heads: the persistent qh32 qseqlen1 lse kernel leaves the final LSE of a request whose work is not split unwritten, with fresh metadata as well, and on main too. O is still checked for them. Restoring main's hand-written fold rule fails 30 of the cases. Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [MLA] Cut the host cost of the decode gate and of get_mla_metadata_info_v1 An eager-mode A/B against main on MI355X (same image, same GPU, ABBA order) showed what moving the gate into C++ cost on the host: mla_decode_fwd issued about 8.7 us more per call (46.6 -> 55.3 us, fp8 16 heads, batch 32) and get_mla_metadata_info_v1 went from 0.6 to 5.9 us. Graph replay, capture, GPU kernel time and buffer sizes were unchanged. Most of the cost was five env-flag reads parsed with a regex, two device-arch lookups and three separate cache lookups per call. - env_flag_atoi memoises the parse per raw value. The variable is still read on every call, so changing it at run time still takes effect. - mla_decode_fwd's persistent gate is one call, _mla_decode_launch_plan. It reads the four flags as raw strings, and one cache keyed on them and on the device returns both the head plan and the split bound. - get_mla_metadata_info_v1 caches its result on its arguments, the arch and the two env flags it depends on. - The occupancy cache is keyed on torch dtypes rather than converted ids. The issue time is now 47.9 us against main's 46.1. The rest is the four env reads that main skips when it short-circuits. get_mla_metadata_info_v1 takes 1.2 us. The CP query-vs-launch test now captures _mla_decode_launch_plan. Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [MLA] test_mla_persistent: take the reference's head fold from the planner torch_mla_extend_split_kv kept its own copy of the v1_2 head-fold rule to undo the planner's 16-head pseudo-batches. Nothing else used that copy, and it only matched fast_mode plans with 32..128 heads. With -mm v1_1, and with the >128-head folds (144 heads, fold 9) that mla_decode_fwd now runs, the reference read work_info batch ids it had not folded and raised IndexError. That happened after the O check had already run, so the split-level comparison was lost. The reference now asks aiter.get_mla_decode_shape_support for the fold. It passes the dtypes and the planner mode (fast_mode/intra_batch_mode) the metadata was built with, which test_mla threads through. The 3-buffer wrapper plans its byte KV as fp8, as mla_decode_fwd does. On MI355X the fold the query reports matches the work_info batch ids for every (dtype, heads, qlen, planner) combination probed. Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [MLA] Shape query: report opus CP launches that abort opus_mla_decode_fwd raises when cp_world_size > 1 and Q is not fp8, or when it gets no g_kv_indptr. The query answered None (no availability table) for both, the same as for a launch that runs. It now answers False with the launcher's reason. The CP query-vs-launch test also asserts that CP without round-robin on gfx1250 is unsupported, not only that it is routed to ASM. Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [MLA] test_mla_shape_support: pin #6563's FlyDSL PS1 rules #6563 widened the gfx1250 FlyDSL PS1 route (32/64 heads at qlen <= 4 with heads x qlen <= 128, bf16/bf16 on the FlyDSL JIT kernel, fp8 asm code objects only) and narrowed 16 heads to qlen <= 4. On rebase those rules were ported into mla_decode_head_plan_v1_2 and mla_decode_backend; these asserts check the query reports them. Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

minmengdie requested review from a team and a balanced review from Copilot October 13, 2026 02:03

github-actions Bot changed the title ~~add wave4 kernel~~ [ASM] [HIP] [JIT] add wave4 kernel Oct 13, 2026

github-actions Bot added ASM HIP JIT labels Oct 13, 2026

Copilot started reviewing on behalf of minmengdie October 13, 2026 02:05

View session

Copilot AI reviewed Oct 13, 2026

View reviewed changes

add wave4 kernel

9f5cdcf

Copilot AI balanced review requested due to automatic review settings October 13, 2026 03:04

zhiding512 force-pushed the mmd/dev/mla_fp8_opus_n16 branch from 0438c31 to 9f5cdcf Compare October 13, 2026 03:04

Copilot started reviewing on behalf of minmengdie October 13, 2026 03:05

View session

Copilot AI reviewed Oct 13, 2026

View reviewed changes

Comment thread aiter/ops/attention.py Resolved

Copilot AI balanced review requested due to automatic review settings October 13, 2026 06:14

Copilot started reviewing on behalf of minmengdie October 13, 2026 06:15

View session

Copilot AI reviewed Oct 13, 2026

View reviewed changes

Comment thread aiter/mla.py Resolved

Comment thread csrc/kernels/mla/opus_mla_decode_fwd.cu Resolved

minmengdie added 2 commits October 13, 2026 06:23

fix the comments

8a46877

fix the black

06a1662

Copilot AI balanced review requested due to automatic review settings October 13, 2026 07:12

zhiding512 force-pushed the mmd/dev/mla_fp8_opus_n16 branch from 6c4d096 to 06a1662 Compare October 13, 2026 07:12

Copilot started reviewing on behalf of minmengdie October 13, 2026 07:13

View session

Copilot AI reviewed Oct 13, 2026

View reviewed changes

Comment thread aiter/mla.py Resolved

Copilot AI balanced review requested due to automatic review settings October 13, 2026 07:40

Copilot started reviewing on behalf of minmengdie October 13, 2026 07:41

View session

Copilot AI reviewed Oct 13, 2026

View reviewed changes

add opus nhead12 support

003136b

Copilot AI balanced review requested due to automatic review settings October 13, 2026 07:49

zhiding512 force-pushed the mmd/dev/mla_fp8_opus_n16 branch from 2ef1c51 to 003136b Compare October 13, 2026 07:49

Copilot started reviewing on behalf of minmengdie October 13, 2026 07:50

View session

Copilot AI reviewed Oct 13, 2026

View reviewed changes

support fold nhead

9ba6016

zufayu requested a review from kaiyang-1 October 14, 2026 01:29

zufayu assigned kaiyang-1 Oct 14, 2026

support cp round robin

e2c0f56

Copilot AI balanced review requested due to automatic review settings October 14, 2026 02:25

Copilot AI reviewed Oct 14, 2026

fp8 change to a8w8

8192c5c

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

Copilot AI reviewed Oct 14, 2026

minmengdie changed the title ~~[ASM] [HIP] [JIT] add wave4 kernel~~ [ASM] [HIP] [JIT] MI355 MLA PS mode kernel add 4 wave fp8 kernel Oct 14, 2026

minmengdie added 2 commits October 15, 2026 09:20

change pin agpr to named-reg-pin-std

8bf3862

modify the opus_mla_decode_kargs

83051be

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

Copilot AI reviewed Oct 16, 2026

kaiyang-1 approved these changes Oct 16, 2026

View reviewed changes

kaiyang-1 merged commit 4fac080 into main Oct 16, 2026

kaiyang-1 deleted the mmd/dev/mla_fp8_opus_n16 branch October 16, 2026 05:51

Conversation

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

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

github-actions Bot commented Oct 13, 2026

🏷️ CI Guide

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

kaiyang-1 left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

minmengdie commented Oct 13, 2026 •

edited

Loading