Repository navigation
[ASM] [HIP] [FlyDSL] Optimize gfx1250 persistent MLA decode, add BF16 and TP4/TP2 MTP - #6563
Merged
Merged
Conversation
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
github-actions
Bot
changed the title
Co-authored-by: Cursor <cursoragent@cursor.com> # Conflicts: # aiter/ops/flydsl/kernels/mla_gfx1250/mla_pagesize1_fp8_fp8.py # hsa/gfx1250/mla_dsl/mla_a8w8_qh128_1tg_16mx8_64nx1_cprr_ps.co # hsa/gfx1250/mla_dsl/mla_a8w8_qh128_1tg_16mx8_64nx1_lse_cprr_ps.co # hsa/gfx1250/mla_dsl/mla_a8w8_qh128_1tg_16mx8_64nx1_lse_ps.co # hsa/gfx1250/mla_dsl/mla_a8w8_qh128_1tg_16mx8_64nx1_msk0_lse_ps.co # hsa/gfx1250/mla_dsl/mla_a8w8_qh128_1tg_16mx8_64nx1_msk0_ps.co # hsa/gfx1250/mla_dsl/mla_a8w8_qh128_1tg_16mx8_64nx1_ps.co # hsa/gfx1250/mla_dsl/mla_dsl.csv
shengnxu
approved these changes
Oct 20, 2026
okorzh-amd
added a commit
to okorzh-amd/aiter
that referenced
this pull request
Oct 20, 2026
ROCm#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>
zejunchen-zejun
added a commit
to ROCm/ATOM
that referenced
this pull request
Oct 21, 2026
…d BF16 caches on gfx1250 too Review follow-up. - mla_min_query_heads and _MLA_MIN_HEADS move to atom/model_ops/mla_query_heads.py, which imports neither triton nor aiter (get_gfx is read inside the function). attention_mla imports _MLA_MIN_HEADS from it, so its other importers are unchanged. The test stubs aiter's chip_info instead of importorskip-ing, so its cases run on the CPU-only CI. - The docstring no longer states aiter's PS1 gate as a fixed rule: it depends on the aiter version (ROCm/aiter#6563 widened it and added BF16). 96 heads, and 16 heads up to 4 wide, pass on every version. - The gfx1250 padding no longer requires an fp8 cache. Before #6563 a BF16 cache has no gfx1250 decode either way; since #6563 the pad keeps a BF16 draft on the PS1 kernel instead of the persistent asm decode. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
okorzh-amd
added a commit
that referenced
this pull request
Oct 23, 2026
…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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
This PR moves FP8 PS1 to a four-wave design, which makes
mla_decode_fwd15–34%faster on every FP8 shape it served before. It also adds a BF16 PS1 kernel, plus MTP
for nhead 32 (q 1–4) and nhead 64 (q 1–2) on both dtypes
Technical Details
Test Result
Accuracy
All 16 runs pass: 992 cases and 4584
checkAllclosechecks (output and LSE), 0failures, no core dumps. ruff and black are clean.
mla_decode_fwdperformance (B=1536, µs per call)FP8
Submission Checklist