[HIP] FIX MLA the nhead fold error for cp round robin - #4964
Merged
Conversation
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags: |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR updates the MLA persistent decode “nhead folding” path to rebuild/adjust indptr tensors (and CP global-KV indptr) after reshaping Q/O into a folded (nhead=16) representation, aiming to fix failures in context-parallel (round-robin) mode.
Changes:
- Rebuilds
qo_indptrfor the folded layout in the persistent-mode nhead-fold branch. - When
g_kv_indptris provided, derives a foldedg_kv_indptrand adjustskv_indptraccordingly.
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
minmengdie
changed the title
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.
Suppressed comments (1)
aiter/mla.py:860
- In the nhead-folding CP path,
kv_indptris expanded viakv_indptr[:-1].repeat_interleave(fold_factor), which repeats the start offsets but not the per-batch lengths. That produces long runs of identical adjacent entries, sokv_indptr[i+1] - kv_indptr[i]becomes 0 for most of the folded “batches”, implying empty KV ranges for many head-groups. This looks inconsistent with the foldedg_kv_indptr(which is built from per-batch lengths) and is very likely to break KV slicing / KV-length computations under folding.
Please re-check the intended semantics here: either (a) leave kv_indptr in original batch space and have the kernel map folded batch→original batch (like metadata does via qk_batch_ratio), or (b) if the kernel truly expects folded-batch kv_indptr, fold kv_indices/kv_indptr together so each folded batch has a correct, non-zero KV segment.
if g_kv_indptr is not None:
g_len = (g_kv_indptr[1:] - g_kv_indptr[:-1]).repeat_interleave(
fold_factor
)
folded_g_kv_indptr = torch.zeros(
bs * fold_factor + 1, dtype=g_kv_indptr.dtype, device=device
)
folded_g_kv_indptr[1:] = torch.cumsum(g_len, 0).to(g_kv_indptr.dtype)
g_kv_indptr = folded_g_kv_indptr
kv_indptr = torch.cat(
[kv_indptr[:-1].repeat_interleave(fold_factor), kv_indptr[-1:]]
)
minmengdie
force-pushed
the
mmd/fix/fold_cp_round_robin
branch
from
August 25, 2026 04:02
6f0c98d to
d579b1a
Compare
github-actions
Bot
changed the title
shengnxu
approved these changes
Aug 28, 2026
JiaoliangYu
added a commit
that referenced
this pull request
Aug 28, 2026
…error check (#5075) * [FlyDSL] gfx942 a16wi4: pack f32->bf16 with lshr-16 instead of scalar (#5017) * [FlyDSL] gfx942 a16wi4: pack f32->bf16 with lshr-16 instead of scalar truncf v_cvt_pk_bf16_f32 is gfx950-only. After #4646 the gfx942 int4 fallback used f32.to(bf16)/truncf, which is much more VALU than the old moe_gemm_2stage shift-pack. Same nibble order; gfx950 packed convert and MXFP4 are unchanged. * [FlyDSL] Clarify gfx942 a16wi4 upconvert comments * ci: allow multigpu label to trigger tests (#5008) * [HIP] [CK] [MoE] Added Gelu with tanh approx for CK XDL 2-stage MoE (#4620) * [MoE] Added Gelu with tanh approx for CK XDL 2-stage MoE * applied copilot's comment for str2ActivationType * Dropping cross-activation CK configs * Applied Ying comment * Added block for run_1stage for unsuported activations * fixed test * [Triton] Move attention configs to nested layout and unify their resolution (#5019) Relocate 14 attention config files from the flat arch-prefixed layout to configs/<arch>/triton/attention/<d_type>/DEFAULT.json - mha, extend_attention, mla_decode_rope, hstu_attn_fwd and hstu_attn_bwd - retiring configs/hstu_attn/. The redundant -DEFAULT suffix is dropped from directory names (the file is already DEFAULT.json), matching the chunk_delta_attn precedent. The six reader modules resolve through the shared resolve_config_dir("attention", ...) probe instead of hand-built paths. LEANATTN is not migrated: upstream removed the lean_atten kernel and its config. * [Triton] Migrate the GMM tuned configs to the nested layout (#5020) Move configs/<arch>-GMM.json (gfx942, gfx950, gfx1250) to configs/<arch>/triton/gmm/gmm/DEFAULT.json and point the reader at it. GMM gets its own op directory instead of folding under gemm/. The doubled gmm/gmm is just the <op>/<d_type> layout: the op is "gmm" and the family's config name is "GMM", so _dtype_dir() yields "gmm" too. _triton_kernels/gmm.py now resolves the directory through the shared resolve_config_dir("gmm", "GMM", backend="triton") probe and loads DEFAULT.json from it; arch_info and AITER_TRITON_CONFIGS_PATH are dead there and are dropped. No legacy_dir is passed - the files move and the loader flips in this one commit, so every revision resolves. * [Triton] Move MOE tuned configs to the nested layout (#5022) Move the three remaining MOE tuned configs from the flat configs/moe/ directory into configs/<arch>/<backend>/<op>/<d_type>/: moe/gfx950-A8W4.json -> gfx950/triton/moe/a8w4/DEFAULT.json moe/gfx1250-A8W4.json -> gfx1250/gluon/moe/a8w4/DEFAULT.json moe/gfx1250-A4W4.json -> gfx1250/gluon/moe/a4w4/DEFAULT.json The backend directory follows the dispatch path the table actually feeds, not the arch: gfx950's a8w4 table is keyed bm<block_m>_n<N>_k<K> and is read by the Triton path, while both gfx1250 tables are bucket-keyed and read by the Gluon path. So the a8w4 family spans backends and a4w4 is Gluon-only. These three are all that is left of configs/moe/: PR #4833 removed the rest of the legacy MOE stack (utils/moe_config_utils.py, the fused sigmoid-top1 routing kernel, the moe_op/moe_op_e2e/mxfp4 variants and every configs/moe/*-MOE-*.json), so this completes the directory. The two surviving loaders are rewired onto the shared probe in the same commit. _get_a8w4_dispatch() and _get_a4w4_dispatch() now resolve their directory with resolve_config_dir("moe", "<A8W4|A4W4>") and read DEFAULT.json from it, instead of hand-building an arch-prefixed path under configs/moe/. Neither call passes backend=: because the backend differs per arch for the same family, pinning one would make the other arch's file unreachable. The documented probe order -- nested triton, then nested gluon -- picks whichever directory the running arch ships. a4w4 also moves off its private os.path.exists + json.load pair onto load_config_json(..., required=False), matching a8w4; both still return {} when no tuned file is shipped for the arch, so the safe-default fallback paths are unchanged. resolve_config_dir() lives in utils/gemm_config_utils.py and is added by the config-unification branch -- merge that one first. * [Doc][Skill] port udpated flydsl kernel code cleanup skill (#5051) * [Triton/Gluon] MOE a8w4 cudagraph updates (#5037) * [Triton] Remove legacy flat-layout support from config resolution (#4948) * [Triton/Gluon] Move gluon gemm_a8w8 kernel into _gluon_kernels/gfx950 (#4866) * [Triton] Migrate conv configs to the nested arch/backend layout (#5018) Move all 59 flat configs/conv/<arch>-<CONFIG_NAME>.json files to configs/<arch>/triton/conv/<d_type>/DEFAULT.json, the layout GEMM already uses, and point _conv_config_path() at the shared resolve_config_dir() probe instead of building the legacy path by hand. This picks up the ten tables #4869 added (CONV-PREPACK on all seven arches, CONV-3X3-NCHW on gfx1100/gfx1151/gfx1201) alongside the original 49. The renames and the loader flip land in one commit so every revision resolves conv configs from exactly one layout: no legacy_dir fallback is needed and bisect stays clean. File contents are untouched (pure renames). _conv_config_path() is the single choke point, so get_conv_config(), has_conv_config(), conv_config_uses_exact_routes() and has_exact_conv_config() all pick up the nested path; the variant-aware four-tier walk, STANDARD_M_BOUNDS and the lru_caches are untouched. * Tune MoE GEMM A8W8 blockscale (#5028) * [Triton] Migrate MHC configs to the nested arch/backend layout (#5021) Move all 15 flat configs/<arch>-MHC_*.json files to configs/<arch>/triton/mhc/<d_type>/, keeping the C=<n> specialized file stems and naming each family default DEFAULT.json, and rewire mhc_config_utils onto the shared resolve_config_dir() probe. The documented gfx942 fallback retry resolves through the probe's arch= override (added by the legacy-removal PR, which merges first); the C-bucket walk and _FALLBACK_DEV semantics are unchanged. The renames and the loader flip land in one commit so every revision resolves MHC configs from exactly one layout, and this branch touches no shared resolver code. * [Gluon] add bench for mxfp8 GEMM (#5029) * [FlyDSL] One-stage split-K for the a8w8 preshuffle GEMM (#5007) * [FlyDSL] one-stage split-K for the a8w8 preshuffle GEMM Fold the split-K reduction into the GEMM launch: every split publishes an fp32 partial, and the last one to arrive at the tile's semaphore reduces and converts in the same kernel, so split-K costs one launch rather than two. The partials cross CTAs that may sit on different XCDs, each with its own L2, so they have to reach a common point. Doing that with an agent-scope fence costs a whole-L2 buffer_wbl2 per CTA plus a buffer_inv on the reader, which also evicts the A/B tiles every other in-flight CTA is still reading -- measured 2-3x slower than not splitting at all. Marking just these accesses sc0|sc1 writes them through and leaves L2 alone. That turns split-K from a 18-180% regression into a 9-50% win over k_split=1. k_split == 1 is untouched: same kernel, same cached stores, and its output is bitwise identical to main across the shapes checked. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * [FlyDSL] drop the k_split == 2 special case It had its own epilogue: whichever split arrived first published its fp32 fragment, and the second spun on a ready flag, kept its own fragment in VGPRs and wrote the final tile -- saving one workspace plane and one round trip. It does not pay for itself. Spinning is slower than just going through the generic path: 4.7 vs 5.8 us at 1x576, 18-27% across the six shapes measured. Removing it also drops two fragments, two copy atoms, the doubled semaphore, and the split-plane special cases in the launcher and the AOT pre-compile. k_split == 2 now takes the same path as every other split count, which also fixes the per-split workspace offset: it was guarded on split_k > 2, so a k_split == 2 launch routed through the generic path would have had every split write the same plane. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * [FlyDSL] fix the split-K keyword the tuner passes to the launcher The tuner called the preshuffle launcher with k_split=, but that launcher names the argument split_k= (matching the hgemm split-K path it sits next to). Every flydsl preshuffle candidate therefore raised TypeError. The tuner records a raising candidate as rejected rather than as an error, so the run completed, kept only the 8wave candidates, and picked ck or cktile for four shapes that flydsl had previously won -- a result indistinguishable from a legitimate tuning outcome. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * [FlyDSL] trim the split-K diff to what split-K needs Three things that were not split-K: out_dtype grew an fp32 branch and a raise. Nothing needs it -- fp32 is the type of the *partial*, which "Float32 if split_k > 1" already covers, and the final output is still bf16/fp16. The bias element-type change existed only to feed that branch. Both are back to main's two-case form. The K-tile index has to gain a bid_z offset, which is genuine, but the name k_tile_base pushed several one-line fx.copy calls past the line limit and a trailing comma pinned others open, so a one-token change read as +5 -1. Renaming to k_off and dropping the magic trailing commas keeps them one-liners. The copy atom for the output no longer branches on out_elem_bytes; it picks the op from split_k directly. Kernel diff: +201 -22 -> +168 -21, with no behaviour change. k_split == 1 still compiles to bytes identical to main. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * [Config] retune Kimi-K3 a8w8 M<=32 with one-stage split-K 29 rows, all M <= 32: 21 stay on flydsl with a better config and 8 move from ck to flydsl. 20 of them use split-K, mostly k_split=7. The tuner proposed 37 rows. Each changed row was then re-measured old config against new on an idle GPU, and the 8 that were actually slower there were kept at main's value -- the tuner picks its winner while four GPUs are saturated, and for shapes where several configs sit within noise of each other that choice does not survive on an idle card. Nearly all of them were k_split=2 at N=6400, which lost 7-13%. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * [FlyDSL] hoist the partial store out of the split_k branch Both arms opened with the same fx.copy; only what follows it differs. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * [FlyDSL] address review: semaphore dtype, buffer lifetime, reduce guard The k_split == 1 path passed an empty bf16/fp16 tensor in the semaphore slot while the AOT pre-compile passed an empty int32 one. dtype is part of FlyDSL's executable cache signature, so every non-split-K preshuffle kernel missed its AOT entry and JIT-compiled at first call -- a regression across all existing tuned configs, not just split-K. Both sides now pass int32. The split-K buffers were cached per (m, n, tile, k_split). m is in the key, so a server sweeping batch sizes grows the cache without bound, and an eviction frees memory whose address a captured CUDA graph still holds. k_split_candidates only proposes split-K while the tile grid is under one CTA per CU and caps k_split * tile_count at four per CU, which bounds tile_count below CU_NUM and the workspace at 4 * CU_NUM * tile_m * tile_n floats -- so the buffers are now fixed-size and keyed on (device, stream) only, the way _get_split_k_tensors already does it, with a capacity check for anything that would exceed the bound. The reduce derives its vector count as tile_n // 4 and would have dropped the tail columns for a tile_n that is not a multiple of 4; the comment claimed the invariant but nothing enforced it. Now rejected at compile time. The semaphore reset was a plain cached store while every other cross-CTA access in that block carries sc0|sc1. It is the same cross-XCD handoff, one launch later, so it writes through too. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * [FlyDSL] guard M against the layout bound the kernel assumes The kernel views A and C through layouts with a hardcoded 65536 rows, in three places, with nothing on the host stopping a larger M from indexing past them. Named the bound, used it at all three sites, and rejected an out-of-range M in the launcher with a message that says why. gemm_kernels keeps its own copy of the literal because that module has to import without FlyDSL present; a test asserts the two agree. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * [FlyDSL] cut the comments back to what the code cannot say 57 added comment lines down to 26. Dropped the ones restating the line below them -- what k_off is, which path split_k > 1 takes, that the partial store publishes a partial -- and shortened the rest. What is left is the reasoning that is not recoverable from the code: why the partials cannot use an agent-scope fence, why the buffers are fixed size rather than shape-keyed, why the semaphore dtype has to match the AOT side, why _REDUCE_VEC is 4, and why the k_split candidates are enumerated per shape. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * Extract the flydsl split-K reduce into a reusable copy-atom epilogue Move the one-stage split-K reduction out of preshuffle_gemm into splitk_epilogue.splitk_reduce_epilogue, with the output element class as its only dtype knob so other GEMMs can reuse it. The reduce now goes through copy atoms and a buffer-tensor descriptor instead of raw buffer_ops: make_layout_tv gives each thread 4 contiguous columns, so the loads stay dwordx4 and the stores dwordx2, and the descriptor bounds cover the ragged-M tail. Resetting the semaphore with atomic_add(-split_k) rather than a plain store also drops a next-launch increment race. Verified on gfx950: rel_err matches k_split=1 for M in {1,8,64} x k_split in {1,2,7,14}, CUDA-graph replays clean, and the k_split=1 output hashes are identical to origin/main. Over the 20 tuned Kimi-K3 split-K shapes the reduce is 0-3% faster than the buffer_ops version. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * Let --splitK gate the FlyDSL candidates as it does the other backends The ck, cktile and asm task builders all take useSplitK and collapse the split-K dimension to a single splitK=0 candidate when it is off. The FlyDSL builder never received the flag, so it generated k_split candidates unconditionally: on the Kimi-K3 shape set that is 98488 extra candidates on top of 99136, roughly double overall and 2.5-3.1x over M in 1..128. It also made the flag useless as a switch. Split-K wins often enough at small M that 20 of the 176 tuned rows are split-K winners, so a run without --splitK still produced split-K configs. --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> * [FlyDSL] Retune GLM5.2 mxfp4 MoE and fix a scale-view cache leak (#5045) * [FlyDSL] Retune GLM5.2 mxfp4 MoE and fix a scale-view cache leak Retune all 64 GLM5.2 shapes (model_dim=6144, inter_dim 256..2048, E=257, topk=9) for gfx950. 27 shapes move to the coupled flydsl_mxmoe port, which the previous config only reached on 5 rows. Measured through the production fused_moe path, each shape timed on one GPU under both configs: median +6.7%, mean +8.6%; 44/64 faster by >1%, 6 slower (worst -3.5%). Small batches gain most (token<=64 median +11.4%). Six shapes (2048/256, 4/1024, and 2/16/64/128 at 2048) are left on their existing main entries rather than retuned. Two fixes fell out of the tuning runs: _mxfp4_scale_u8 was wrapped in lru_cache(maxsize=2048). Its body is a bare .view(), so the memo buys nothing, but tensors hash by identity: every per-call intermediate scale misses and is then pinned by the cache. A tuning sweep leaked ~0.75 GiB per timed iteration and exhausted a 288 GiB card. v2_stage1_dequant_cosine_err looped per sorted row, costing one .item() sync each -- ~295k syncs per timed candidate at token=32768/topk=9. Now batched in chunks, which bounds the int64 gather in mxfp4_to_f32 while keeping the equal-weight average over rows. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * [FlyDSL] Default FMoE tuning to FlyDSL v2 and update GLM5 FP4 layout configs * fix black test * Emit the non-f4out AOT job for _f4out mxmoe stage-2 rows An `_f4out` GEMM2 kernel only really runs the mxfp4-out path when both gates are open: AITER_MXFP4_INTERMEDIATE, and the shape check in fused_moe (`D_HIDDEN == 7168`). Otherwise `_f4out` is stripped from the kernel name and the plain kernel launches instead. The AOT generator skipped such rows outright, so it never pre-compiled the kernel that actually launches. GLM5 is D_HIDDEN=6144, so every `_f4out` row there falls back -- and the config only survived because an unrelated row happened to name the plain kernel and seed the same cache entry. Retuning that row to `_f4out` removed the last such seed and CI hit `FLYDSL_RUNTIME_RUN_ONLY=1 but no usable AOT cache for launch_gemm2` on token=16384, inter_dim=512. Emit the fallback job unconditionally, plus the f4out one when AITER_MXFP4_INTERMEDIATE is set. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: charlieguo1106 <cguo@amd.com> * [UT] Support a4w4 in test_mega_moe (#5052) * support a4w4 in test_mega_moe_gfx1250 * support 64K * [Triton/Gluon] [ASM] [HIP] add mla v4 prefill asm kernel (#4926) * Add MLA v4 sparse prefill asm support Integrate the gfx1250 MLA implementation and consolidate sparse prefill correctness and performance tests. Co-authored-by: Cursor <cursoragent@cursor.com> * Update op_tests/test_pa_sparse_prefill.py Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * Apply black formatting to test_pa_sparse_prefill Pure reformat, no behaviour change. Fixes the failing black CI job (black[colorama]==26.5.1, default line length). * Fix asm candidate reference in test_pa_sparse_prefill The asm candidate passed split["ref"] -- the raw input dict -- where checkAllclose expects the reference tensor, so the first asm comparison died with: TypeError: isclose(): argument 'other' (position 2) must be Tensor, not dict meaning the asm path could never run. Compute the fp8 reference the same way the opus fp8 candidate above it does. * Report per-row nnz and default the CLI to the asm comparison sweep Two test-driver changes: * nnz_prefix/nnz_extend columns now report per-row nnz instead of the pool-wide total, so they match the --nnz-prefix/--nnz-extend asked for rather than scaling with N. total_nnz still carries the full count -- the TFLOPS/TB-s figures need the real work done. * CLI defaults now describe the three-way opus/triton/asm comparison out of the box: N in [512, 1024, 2048, 4096] x nnz_prefix in [256, 1024, 4096, 8192, 16384] x nnz_extend 128, at H_Q=128 fp8 (the only shape the asm candidate registers for). --mode/--total_pages default empty so the unrelated mode sweep stays off unless asked for. Every flag still overrides. Pytest coverage is unaffected: it reads _PYTEST_SHAPES/_PYTEST_MODES, not argparse. * Accept an over-allocated CSR indptr in mla_sparse_prefill check_csr required indptr->numel() == T+1 exactly. Decode reuses this kernel with the extend region empty and sizes its CSR row-pointer buffers once at [max_batch+1], launching with the live batch, so numel > T+1 is the normal case there rather than a mistake -- and the exact test rejected it outright. The kernel reads indptr[0..T] and nothing past it, so the extra tail is inert: verified bit-identical output against the exactly-sized call. An undersized indptr is still rejected. Trade-off: an indptr built for a different T is no longer caught here. Separating that from the legitimate case needs device data (indptr[T] against the indices length), i.e. a sync per call. Callers that can slice to [:T+1] should. * Fix int32 overflow in sparse prefill query offset `_sparse_attn_prefill_kernel` derived `query_idx` from `tl.program_id(0)`, which Triton types as int32. The q/out addresses are computed as `query_idx * q_stride_t` and `query_idx * out_stride_t`, and in the V4 layout that stride is `num_heads * head_dim` = 128 * 512 = 65536. The product therefore leaves the int32 positive range at `query_idx >= 32768` and wraps to a negative offset, so the kernel reads and writes outside the q/out allocations. Observed as NaNs followed by a hard GPU page fault: Memory access fault by GPU node-2 ... Reason: Page not present Verified on gfx1250 with a fixed-pattern sparse prefill case (H=128, D=512, pool=16384, nnz_prefix=256, nnz_extend=128): N=32768 before: clean (largest size that still fits int32) N=32769 before: fault after: nan=0 inf=0 N=65536 before: fault after: nan=0 inf=0 Promoting `query_idx` to int64 moves both offsets to 64-bit address arithmetic. This mirrors the existing `slot_off` cast a few lines below, which already handles the same class of overflow on the pool index; the difference is that the wrapped pool offset stays inside the allocation and reads silently, while this one faults. --------- Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * [HIP] FIX MLA the nhead fold error for cp round robin (#4964) * fix the nhead fold error for cp round robin * fix the split test * support varlen * gqa96 qseqlen<6 not fold * [HIP] [Bugfix] Fix DSV4 FP4 KV-cache scattered row writes (#5034) * fix(dsv4): scatter FP4 KV cache by row-local offset Signed-off-by: AMD-yanfeiwang <yanfei.wang@amd.com> * test(dsv4): remove specialized KV-cache regression Keep the bug fix focused without carrying a narrow special-case test. --------- Signed-off-by: AMD-yanfeiwang <yanfei.wang@amd.com> * [ASM] [HIP] 1x32 mxfp4 asm kernel (#4890) * 1x32 mxfp4 asm kernel * Update tuned config * Upate 1x32 kernel to embedd X quant * Drop the standalone MXFP4 X quant pre-pass plumbing The FLAT MXFP4 kernels dynamic-quantize X in-kernel, so the host-side pre-pass entrypoint, its Python binding and the test helper have no caller left. Removing them also restores the per_1x32 scale-sorting condition, which still tested a pre-pass flag that no longer exists. Co-authored-by: Cursor <cursoragent@cursor.com> * Tune 1x32 kernel * Fix 1x32 race condition for O buffer clearning * SImplify zero protocall and bind it to TG0 always * Fix lm_eval utter failure with 1x32 kernel --------- Co-authored-by: Sergey Solo <ssolovye@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> * [HIP] update MHA CPP reademe (#4874) * update the supported arguments configuration * update the perf data * update the image * fix the log * fix * benchmark_fwd support opus kernel * add opus perf data * perf data * fix the comment * fix * Test FFM bringup on MI250 build runner (#5071) * perf(gfx1250): drop a16w16's 4 GiB pre-check, fail on wrong results The bench skipped any a16w16 shape whose largest operand passed 4 GiB, on the stated grounds that "the heuristic refuses" it. That reads the guard too broadly. opus_dispatch_a16w16_gfx1250 (opus_gemm_arch_gfx1250.cuh:150-183) searches the tuned table first and returns on a hit. check_shape_4g runs only after that misses, on the way to the split-K heuristic kid, and it is that kid's launcher that builds the 32-bit gmem descriptors. A tuned 4wave_wl_co winner never reaches the check: per gen_instances_gfx1250.py:770-778 the pipeline "builds no gmem descriptor at all" and clamps every dimension through TDM descriptors instead. So 4 GiB bounds one fallback path, not a16w16, and a pre-check in Python keeps skipping shapes that tuning has already made runnable. Removed; the kernel raises if it must, and the exception is recorded as a row. The 20260827 sweep shows what the fallback costs. At N=129280 the tuned 4wave_wl_co kid does M=512 in 449us (2112 TFLOPS); M=1024 has no tuned winner, drops to split-K, and takes 3343us (568 TFLOPS) -- 7.4x slower for 2x the work. 11 of 60 shapes hit a 4wave_wl_co kid; the rest are split-K, so most of the low numbers in this table measure tuning coverage rather than the hardware. Widening that coverage is a job for csrc/gemm_a16w16/gemm_a16w16_tune.py --libtype opus, not for this file. Worse, split-K is not just slow at the top of the range: all four M=65536 shapes came back err=0.96-0.99, an unrelated result, while every other row was 0 or ~1e-5. None of them trip the 4 GiB guard (M*K*2 = 896 MB, M*N*2 <= 256 MB), and the UT neither raises nor warns -- it returns the ratio and prints a number. The sweep reported them as data. a16w16 now checks the returned ratio against _A16W16_MAX_ERR and calls _note_failure, so a silent miscompare shows up in the failed-op list. a16w16 also gets its own M list. The global sweep jumps 2048 -> 65536, so the prefill chunk sizes were never measured on the BF16 linears; _A16W16_MS adds 4096/8192/16384 and AITER_BENCH_TOKENS still overrides it. The lm_head cap stays. It is a statement about what DSv4 runs -- one row per sequence -- not about what the kernel can do, and its comment no longer leans on the 4 GiB number. Separately, put a8w8_blockscale back in --dsv4 and correct its note. The note blamed #4773's gluon tuning rows for the make_llir crash. The real cause is the UT's extra "ck strided x_scale" check (test_gemm_a8w8_blockscale.py:120), added by #4406 and gated on ck_preshuffle alone. The mxfp8_128 path declares its layout with is_x_scale_transposed=True and never reads the stride, so a strided x_scale tests nothing there and only gives triton a specialization that fails to compile. A/B with that line as the only variable, over a 162-case matrix (27 default M x six (n,k)): case 2 before it dies, case 160 after -- M=16 and M=64 included, which is what #4773 covers. Fixing it properly is upstream's call; meanwhile _A8W8_BLOCKSCALE_TOKENS starts at 1024, clear of the M that reach those rows. Verified on gfx1250-atom--20260827-ubench: 36/36 cases, err=0, 2207-7003 TFLOPS. --------- Signed-off-by: AMD-yanfeiwang <yanfei.wang@amd.com> Co-authored-by: msaffari-amd <msaffari@amd.com> Co-authored-by: Xin Huang <Xin.Huang@amd.com> Co-authored-by: Alexandra Sidorova <asidorov@amd.com> Co-authored-by: Satya Nikhil Kodukula <nikhil.kodukula@gmail.com> Co-authored-by: Felix Li <felix.li@amd.com> Co-authored-by: Lukasz Burzawa <lukasz.burzawa@amd.com> Co-authored-by: Vinayak Gokhale <vinayak.gokhale@amd.com> Co-authored-by: Nidal Danial <81209936+nidal567@users.noreply.github.com> Co-authored-by: Shao-Chun Lee <Shao-Chun.Lee@amd.com> Co-authored-by: XiaobingZhang <xiaobingzhangupc@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: charlieguo1106 <cguo@amd.com> Co-authored-by: yanboshao <yashao@amd.com> Co-authored-by: junxiaguo <JunXia.Guo@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: minmengdie <memin@amd.com> Co-authored-by: AMD-yanfeiwang <yanfei.wang@amd.com> Co-authored-by: Sergey Solovyev <sergey.solovyev@amd.com> Co-authored-by: Sergey Solo <ssolovye@amd.com> Co-authored-by: Yu <jiaolyu@amd.com>
4 tasks
okorzh-amd
added a commit
to okorzh-amd/vllm
that referenced
this pull request
Sep 2, 2026
The gfx950 asm persistent decode has kernels for head tiles {16,32,64,128}.
aiter accepts any 16-aligned count, but a non-native one is *folded*:
qk_batch_ratio = nheads // 16; nheads = 16; num_batches *= qk_batch_ratio
(aiter csrc/kernels/mla/metadata/v1_2_device.cuh, mirrored in aiter/mla.py),
and every sub-pass of the fold re-reads the same KV: the batch index is
divided by qk_batch_ratio and the KV cursor only advances on sub-head 0.
get_actual_mla_num_heads rounds to the next multiple of 16, so a count that is
already 16-aligned but not native passes straight into the fold. Under decode
context parallelism the query is all-gathered to num_heads * dcp_world_size,
which makes this easy to hit: Kimi-K3 at TP8/DCP8 gathers 12*8 = 96 heads and
pays a 6x KV re-read on the largest kernel in the model. DeepSeek at TP8/DCP8
gathers 16*8 = 128 and is native, which is why this has gone unnoticed.
Rounding up to the next native tile instead trades a few zero query heads of
compute for the entire re-read.
Measured on 8x MI355X, aiter 0.1.19, fp8 q + fp8 kv, bs=52 qlen=1, 16384 KV
tokens/request, kernel identity read from rocprofv3:
NH=16 qh16_qseqlen1_gqaratio16_lse_ps 104.32 us FETCH 497.3 MB (1.01x)
NH=96 qh16_qseqlen1_gqaratio16_lse_ps 464.14 us FETCH 2970.9 MB (6.05x)
NH=128 qh32_qseqlen4_gqaratio32_lse_ps 185.26 us FETCH 517.2 MB (1.05x)
The FETCH_SIZE counter is the direct evidence of the re-read; the 1x reference
is 52 * 16384 * 576 B = 490.73 MB.
End to end on Kimi-K3 DCP8/TP8, conc 52, ISL 131072, one image and two runs
differing only by this knob, 8 ranks each:
off: qh16_gqaratio16, MLA 459.38 us mean (8-rank spread 4.2%), 18.9% of GPU
on : qh32_qseqlen4_gqaratio32, MLA 169.37 us (spread 2.7%), 9.6% of GPU
MLA kernel time -63.1%; median ITL 52.66 -> 43.33 ms (-17.7%); completed
requests and total output tokens identical in both arms. The step-time delta
(-17.3 ms/pass) is larger than the MLA delta (-7.26 ms/pass) because collective
time also fell -- mscclKernel_Sum mean 57.06 -> 28.54 us at an unchanged call
count -- i.e. ranks stopped waiting on the slowest MLA. Only the -7.26 is the
kernel.
Off by default: it changes the kernel selected for every non-native head count,
and only 96 -> 128 has been measured. 48 -> 64, 80 -> 128 and 112 -> 128 follow
the same mechanism but are not measured here, and the non-gfx950 / bf16-KV
combinations are not covered by aiter's native list at all.
The padded lanes never escape the backend: get_mla_unpadded_o and
get_mla_unpadded_lse slice them off before forward() returns, so nothing
padded reaches the DCP output combine or the LSE merge. Tests cover the
rounding table and the pad/unpad round trip for o and lse.
Complementary to ROCm/aiter#4964, which teaches aiter to take gqa=96 natively
(routing it to the same qh32_qseqlen4_gqaratio32 kernel) rather than padding to
128, worth a further ~15 us/call. That is merged to aiter main but is in no
released aiter tag; this change works on shipped aiter.
Signed-off-by: okorzh <okorzh@amd.com>
okorzh-amd
added a commit
to okorzh-amd/vllm
that referenced
this pull request
Sep 2, 2026
The gfx950 asm persistent decode has kernels for head tiles {16,32,64,128}.
aiter accepts any 16-aligned count, but a non-native one is *folded*:
qk_batch_ratio = nheads // 16; nheads = 16; num_batches *= qk_batch_ratio
(aiter csrc/kernels/mla/metadata/v1_2_device.cuh, mirrored in aiter/mla.py),
and every sub-pass of the fold re-reads the same KV: the batch index is
divided by qk_batch_ratio and the KV cursor only advances on sub-head 0.
get_actual_mla_num_heads rounds to the next multiple of 16, so a count that is
already 16-aligned but not native passes straight into the fold. Under decode
context parallelism the query is all-gathered to num_heads * dcp_world_size,
which makes this easy to hit: Kimi-K3 at TP8/DCP8 gathers 12*8 = 96 heads and
pays a 6x KV re-read on the largest kernel in the model. DeepSeek at TP8/DCP8
gathers 16*8 = 128 and is native, which is why this has gone unnoticed.
Rounding up to the next native tile instead trades a few zero query heads of
compute for the entire re-read.
Measured on 8x MI355X, aiter 0.1.19, fp8 q + fp8 kv, bs=52 qlen=1, 16384 KV
tokens/request, kernel identity read from rocprofv3:
NH=16 qh16_qseqlen1_gqaratio16_lse_ps 104.32 us FETCH 497.3 MB (1.01x)
NH=96 qh16_qseqlen1_gqaratio16_lse_ps 464.14 us FETCH 2970.9 MB (6.05x)
NH=128 qh32_qseqlen4_gqaratio32_lse_ps 185.26 us FETCH 517.2 MB (1.05x)
The FETCH_SIZE counter is the direct evidence of the re-read; the 1x reference
is 52 * 16384 * 576 B = 490.73 MB.
End to end on Kimi-K3 DCP8/TP8, conc 52, ISL 131072, one image and two runs
differing only by this knob, 8 ranks each:
off: qh16_gqaratio16, MLA 459.38 us mean (8-rank spread 4.2%), 18.9% of GPU
on : qh32_qseqlen4_gqaratio32, MLA 169.37 us (spread 2.7%), 9.6% of GPU
MLA kernel time -63.1%; median ITL 52.66 -> 43.33 ms (-17.7%); completed
requests and total output tokens identical in both arms. The step-time delta
(-17.3 ms/pass) is larger than the MLA delta (-7.26 ms/pass) because collective
time also fell -- mscclKernel_Sum mean 57.06 -> 28.54 us at an unchanged call
count -- i.e. ranks stopped waiting on the slowest MLA. Only the -7.26 is the
kernel.
Off by default: it changes the kernel selected for every non-native head count,
and only 96 -> 128 has been measured. 48 -> 64, 80 -> 128 and 112 -> 128 follow
the same mechanism but are not measured here, and the non-gfx950 / bf16-KV
combinations are not covered by aiter's native list at all.
The padded lanes never escape the backend: get_mla_unpadded_o and
get_mla_unpadded_lse slice them off before forward() returns, so nothing
padded reaches the DCP output combine or the LSE merge. Tests cover the
rounding table and the pad/unpad round trip for o and lse.
Complementary to ROCm/aiter#4964, which teaches aiter to take gqa=96 natively
(routing it to the same qh32_qseqlen4_gqaratio32 kernel) rather than padding to
128, worth a further ~15 us/call. That is merged to aiter main but is in no
released aiter tag; this change works on shipped aiter.
Signed-off-by: okorzh <okorzh@amd.com>
okorzh-amd
added a commit
to okorzh-amd/vllm
that referenced
this pull request
Sep 9, 2026
AITER's MLA metadata planner folds a head count it does not natively
support down to 16 heads:
// csrc/kernels/mla/metadata/v1_2_device.cuh (v0.1.21.post1)
if (!natively_supported && (num_heads % 16 == 0)) {
qk_batch_ratio = num_heads / 16;
num_heads = 16;
num_batches *= qk_batch_ratio;
}
Every pseudo-batch of that fold re-reads the whole KV cache -- the batch
index is divided by qk_batch_ratio and the KV cursor only advances on
sub-head 0 -- so a folded decode costs num_heads/16 passes over the cache.
MLA decode is KV-bandwidth bound, so removing the fold is worth far more
than the extra query heads a pad costs.
natively_supported is not a head-count set. It is a disjunction over
(arch_id, q_is_fp8, kv_is_fp8, num_heads, max_seqlen_qo), and it has
changed in every AITER release vLLM has shipped: v0.1.20 collapsed the
qo-gated 32-head clauses, v0.1.21 (ROCm/aiter#4964) added
gfx950 && fp8 && 96 heads && max_seqlen_qo <= 6. So this does not hardcode
a tile list. It mirrors the clauses in Python and admits each one only if
its literal is present in the shipped JIT source, extending the probe
idiom already in this file (_aiter_mla_native_h24_metadata_supported,
vllm-project#51647). A clause AITER removes stops being claimed here with no vLLM
change; a clause AITER adds is picked up the same way.
Two things gate a pad target beyond the planner's verdict:
- reduce.cu's HEAD_DIM 512 instantiation list. The reducer and the planner
dispatch independently -- the same split that forced the two-probe H24
check -- and the planner's blanket gfx950/bf16 clause claims every head
count while the reducer stops at 128.
- the shipped asm kernel table, hsa/<arch>/mla/mla_asm.csv. A native
planner verdict does not imply mla_decode_fwd can dispatch:
get_heuristic_kernel_mla filters on the LSE flag exactly, and gfx942
ships gqa=128 at lse=0 only. A DCP rank asks for the LSE, so padding to
128 there would be "cannot find suitable kernel" at the first decode.
gqa=64 ships both flags, so the same rank can still pad 48 -> 64. All
four config remaps in asm_mla.cu are guarded on gfx950, where every
persistent fp8 shape lands on gqa 16 or 32 and both carry LSE variants,
so the table is consulted for gfx942 only.
The padded count is resolved once in the metadata builder's __init__,
below the dtype block, and carried on
AiterMLADecodeMetadata.padded_num_heads. It has to be a per-run constant:
max_seqlen_qo varies per build() and a head count that changed between
passes would resize o and break cudagraph capture. A target is accepted
only if it is native at *every* query length the run can produce, not at
the largest: the fold factor is num_heads/16 with no qlen term, so a
target that is non-native at some reachable qlen folds harder there than
not padding at all.
Measured on 8x MI355X (gfx950), AITER v0.1.21.post1, fp8 q + fp8 kv,
bs=52, qlen=1, 16384 KV tokens/request, kernel identity from rocprofv3:
heads AITER time padded to delta
16 native 136.6 us -- --
32 native 153.6 us -- --
48 fold 3x 349.0 us 64, 171.4 us -50.9%
64 native 171.4 us -- --
80 fold 5x 440.1 us 96, 231.7 us -47.4%
96 native 231.7 us not padded --
112 fold 7x 658.2 us 128, 247.2 us -62.4%
128 native 247.2 us -- --
The reachable win is 48 decode heads on gfx950 with an fp8 KV cache:
Kimi-K3 (96 MLA heads) at TP8+DCP4, TP4+DCP2, TP16+DCP8 or TP2+DCP1.
80 and 112 need a model with 80/112 total MLA heads; none exists in tree,
so those rows are mechanism evidence only.
What this explicitly does not do:
- It does not pad 96 heads on a current AITER. v0.1.21+ takes gqa=96
natively on gfx950/fp8 at max_seqlen_qo <= 6, and
docker/Dockerfile.rocm_base pins v0.1.21.post1 (vllm-project#52826). Padding
96 -> 128 there measures 231.7 -> 247.2 us, a 6.7% regression. An
earlier revision of this change did pad it, against v0.1.19 where 96
folded 6x; the probe is what stops that from silently rotting again.
- It does not pad under a bf16 KV cache on gfx950. The blanket
(gfx950 && !q_fp8 && !kv_fp8) clause marks every head count native, so
there is no fold to remove and a pad would be pure waste plus a
materialized q.repeat(...).contiguous(). That is the default
--kv-cache-dtype auto configuration.
- It claims no clause on an arch AITER does not name. Each clause matches
its arch guard, not just the head-count fragment, so a future AITER that
narrows one cannot produce a false positive on the other arch, and an
arch outside {gfx942, gfx950} gets the unpadded rule.
- It changes nothing at or above 129 heads, and nothing for non-multiples
of 16 (120 heads already 16-aligns to 128, which is native). AITER's
python dispatcher folds only nhead in range(32, 128+1, 16) and asserts
otherwise, so there is no target above 128; the resolver logs
warning_once instead of silently accepting the shape.
- Sparse/DSA MLA and the FP8 PS prefill are untouched. Both are different
kernel pairs with no qk_batch_ratio fold to remove.
Env: VLLM_ROCM_AITER_MLA_PAD_TO_NATIVE_SHAPE=auto|off|force, default off.
"off" is bit-identical to the previous next-multiple-of-16 rule, asserted
over every head count in 1..256. "auto" pads only when the installed AITER
reports the target native for this arch, KV dtype and every reachable
query length, has a kernel for it, and the pad stays within
_AITER_MAX_PAD_RATIO. "force" ignores that cost cap. The default flip to
"auto" is deferred until there is a gfx942 sweep and a short-KV point --
the measured curve is one operating point, and the fold/pad crossover
moves toward the fold at short KV.
The padded lanes never escape the backend: get_mla_unpadded_o and
get_mla_unpadded_lse slice them off before forward() returns, so nothing
padded reaches the DCP output combine or the LSE merge. Both now take the
resolved count explicitly, since pad and unpad branch on
m % num_heads == 0 independently and a disagreement returns wrong values
rather than raising.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Oxana Korzh <okorzh@amd.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 updates the MLA persistent decode “nhead folding” path to rebuild/adjust indptr tensors (and CP global-KV indptr) after reshaping Q/O into a folded (nhead=16) representation, aiming to fix failures in context-parallel (round-robin) mode.
Technical Details
Rebuilds qo_indptr for the folded layout in the persistent-mode nhead-fold branch.
When g_kv_indptr is provided, derives a folded g_kv_indptr and adjusts kv_indptr accordingly.
Test Plan
Grid, per mode:
Commands (one process per
(mode, nhead, qlen)pair, 7 GPUs in parallel):Configurations actually executed: 384 non-CP and 576 CP
Test Result
golden fp8 vs aiter_asmgolden fp8 vs aiter_asmcp_ref vs aiter(1728 per-rank checks)cp_ref vs aiter lse(1728 per-rank checks)No configuration fails its authoritative comparison, and no configuration
produced NaN, a GPU memory fault, or a crash. Everything below is about the
size of the residual error and about two reference-side artifacts.
Where the residual error sits (
golden fp8 vs aiter_asm, run 2)Worst max-abs-delta over all batches, by qlen x ctx:
Worst max-abs-delta over all batches and ctx, by qlen x nhead:
Submission Checklist