fix(gemm_a8w8_blockscale): prevent scale OOB and support Triton 3.6 - #4406
Merged
Conversation
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
bingxche
force-pushed
the
bingxche/fix-a8w8-bpreshuffle-compat
branch
from
July 28, 2026 01:43
c041808 to
1051690
Compare
mengfei-jiang
force-pushed
the
bingxche/fix-a8w8-bpreshuffle-compat
branch
3 times, most recently
from
July 28, 2026 02:43
d6580d2 to
95bbaad
Compare
Contributor
Author
|
Triton test passed: https://github.com/ROCm/aiter/actions/runs/30325081807 |
Contributor
|
I am cool with the Triton part change, can you resolve the conflict? Thanks |
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
bingxche
force-pushed
the
bingxche/fix-a8w8-bpreshuffle-compat
branch
from
July 29, 2026 00:26
e71726b to
1dbf724
Compare
valarLip
approved these changes
Jul 29, 2026
This was referenced Jul 30, 2026
k50112113
pushed a commit
that referenced
this pull request
Aug 4, 2026
…ly row-major) (#4506) * [triton] fix transpose_scale in fused_rms_fp8_group_quant (was silently row-major) `fused_rms_fp8_group_quant(transpose_scale=True)` (and its sibling `fused_reduce_rms_fp8_group_quant`) is meant to return the per-group scale in column-major ([num_groups, M]) byte-order so a CK bpreshuffle GEMM consumer can read it without a `.t().contiguous()` relayout copy. The kernel allocates the scale buffer as [num_bs_cols, M] and writes it column-major, but the wrapper returned it with `out1_bs.view(M, num_bs_cols)` -- `.view()` reinterprets that buffer ROW-major, silently scrambling the values and making transpose_scale a no-op (identical layout to transpose_scale=False). Fix: `out1_bs.transpose(0, 1)` -- a true column-major view of the [num_groups, M] buffer (strides (1, M)), same logical values as the default path. Pure-Python triton wrapper, no rebuild. The two existing `*_transpose_scale` op-tests actually asserted the buggy behavior (same row-major strides + a `.transpose().contiguous().view()` scrambled reference). Rewrote both to pin the correct contract, matching the already-correct `test_fused_flatten_fp8_group_quant_transpose_scale`: transpose_scale=True has the same logical values as the default path, is column-major (strides (1, M), `.T` contiguous), and is bit-exact with the materialize (`.t().contiguous().t()`) path. These now fail on the old `.view` (for num_bs_cols >= 2, e.g. N1=7168) and pass on the fix. Validation (MI355X gfx950, rocm/sgl-dev:v0.5.16-rocm720-mi35x-20260731, torch 2.9.1+rocm7.2), full existing parametrization (M, N1/N2[/N3], SPK, dtype in {fp16, bf16}): # with the fix $ pytest test_fused_fp8_quant.py -k transpose_scale -q 93 passed, 210 deselected in 15.91s # reverting only the 2-line fix (transpose(0,1) -> view(M, num_bs_cols)), # same tests -- the rewritten assertions catch the bug: $ pytest test_fused_fp8_quant.py -k transpose_scale 40 failed, 53 passed 36 FAILED test_fused_reduce_rms_fp8_group_quant_transpose_scale 4 FAILED test_fused_rms_fp8_group_quant_transpose_scale E AssertionError: Tensor-likes are not equal! E Mismatched elements: 1790 / 1792 (99.9%) # transpose_scale=True vs default The 40 failures are exactly the num_bs_cols>=2, M>1 cases (N1 in {1536,7168}); the 53 that still pass are the degenerate num_bs_cols==1 cases and the unaffected fused_flatten_fp8_group_quant transpose test. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * [triton] fused_fp8 transpose_scale: support both scale presentations via is_x_scale_strided Per review on PR #4506: keep the API layout-agnostic instead of forcing one convention. Add is_x_scale_strided (default False) to fused_rms_fp8_group_quant and fused_reduce_rms_fp8_group_quant. Both branches return the same column-major scale bytes and differ only in stride(0), which is how the CK bpreshuffle GEMM detects the layout (PR #4406: is_x_scale_tranposed = x_scale.stride(0) != 1): - False (default): contiguous view (M, num_bs_cols), strides (num_bs_cols, 1) -> trans->contig->view - True: strided column-major view, strides (1, M) -> trans->contig->trans Dropped the earlier Triton-version gating idea (the 3.6 compiler issue in #4406 is unrelated). Tests parametrized over is_x_scale_strided with a shared _assert_transpose_scale_layout helper that also verifies, via as_strided, that reading the storage column-major recovers the reference scale in both modes. NOTE: not yet validated on GPU (torch+triton) — pending free GPUs. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
JiaoliangYu
pushed a commit
to JiaoliangYu/aiter
that referenced
this pull request
Aug 28, 2026
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 ROCm#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 ROCm#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 ROCm#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.
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>
JiaoliangYu
added a commit
that referenced
this pull request
Aug 28, 2026
Two files, both new, on top of main: op_tests/bench_gfx1250_combo.py combined bench op_tests/triton_tests/attention/test_mla_v4_triton.py MLA v4 triton reference The bench imports the top-level @benchmark sweep fns from the op_tests siblings and runs each over its own shape axes, printing only the per-op summary tables -- all the JIT/ROCTracer/import noise is silenced at the fd level. Two suites: --perf for hardware-oriented single-op numbers, --dsv4 for the DeepSeek-V4 shapes. Token sweeps come from AITER_BENCH_TOKENS; ops whose axis means something else, or whose usable range is fixed, pin their sweep in the source and say why. Everything else the branch used to carry has reached main on its own since: test_pa_sparse_prefill.py (#4926), test_mega_moe_gfx1250.py (#5052), test_flydsl_qk_norm_rope_quant.py and the mla_v4 hsa artifacts. Only these two files were ever unique to it, so this is what is left to upstream. Notes worth carrying, all measured on gfx1250 / 20260827-28: a16w16 does not pre-check the >4 GiB operand limit. That guard belongs to one fallback path -- opus_dispatch_a16w16_gfx1250 searches the tuned table first and only reaches check_shape_4g on a miss, en route to the split-K kid whose launcher builds the 32-bit gmem descriptors. A tuned 4wave_wl_co winner never gets there. Predicting it in Python skipped shapes that tuning had already made runnable, so the kernel is left to raise and the exception is recorded as a row. It also checks the error ratio the UT returns against _A16W16_MAX_ERR: all four M=65536 shapes come back err=0.96-0.99 while every other row is 0 or ~1e-5, and nothing in the UT raises or warns, so those used to print as data. a8w8_blockscale sweeps M from 1024. Below that, get_CKGEMM_config's M -> get_padded_m -> nextPow2 retry lands on #4773's M=16/M=64 gluon rows, and the UT's extra "ck strided x_scale" check (line 120, added by #4406 and gated on ck_preshuffle alone) hands triton a stride != 1 specialization that fails to compile in make_llir. The mxfp8_128 path declares its layout with is_x_scale_transposed=True and never reads the stride, so that check tests nothing there. Fixing the gate is upstream's call; starting at 1024 keeps this bench clear of it. Verified 36/36, err=0, 2207-7003 TFLOPS. mla_v4_prefill is pinned to n=1024 and mla_v4_prefill_fp8 drops nnz_prefix=8192; both are kernel/verify faults, measured and documented at the pins. Co-authored-by: Yu <jiaolyu@amd.com>
JiaoliangYu
added a commit
to JiaoliangYu/aiter
that referenced
this pull request
Sep 3, 2026
Two files, both new, on top of main: op_tests/bench_gfx1250_combo.py combined bench op_tests/triton_tests/attention/test_mla_v4_triton.py MLA v4 triton reference The bench imports the top-level @benchmark sweep fns from the op_tests siblings and runs each over its own shape axes, printing only the per-op summary tables -- all the JIT/ROCTracer/import noise is silenced at the fd level. Two suites: --perf for hardware-oriented single-op numbers, --dsv4 for the DeepSeek-V4 shapes. Token sweeps come from AITER_BENCH_TOKENS; ops whose axis means something else, or whose usable range is fixed, pin their sweep in the source and say why. Everything else the branch used to carry has reached main on its own since: test_pa_sparse_prefill.py (ROCm#4926), test_mega_moe_gfx1250.py (ROCm#5052), test_flydsl_qk_norm_rope_quant.py and the mla_v4 hsa artifacts. Only these two files were ever unique to it, so this is what is left to upstream. Notes worth carrying, all measured on gfx1250 / 20260827-28: a16w16 does not pre-check the >4 GiB operand limit. That guard belongs to one fallback path -- opus_dispatch_a16w16_gfx1250 searches the tuned table first and only reaches check_shape_4g on a miss, en route to the split-K kid whose launcher builds the 32-bit gmem descriptors. A tuned 4wave_wl_co winner never gets there. Predicting it in Python skipped shapes that tuning had already made runnable, so the kernel is left to raise and the exception is recorded as a row. It also checks the error ratio the UT returns against _A16W16_MAX_ERR: all four M=65536 shapes come back err=0.96-0.99 while every other row is 0 or ~1e-5, and nothing in the UT raises or warns, so those used to print as data. a8w8_blockscale sweeps M from 1024. Below that, get_CKGEMM_config's M -> get_padded_m -> nextPow2 retry lands on ROCm#4773's M=16/M=64 gluon rows, and the UT's extra "ck strided x_scale" check (line 120, added by ROCm#4406 and gated on ck_preshuffle alone) hands triton a stride != 1 specialization that fails to compile in make_llir. The mxfp8_128 path declares its layout with is_x_scale_transposed=True and never reads the stride, so that check tests nothing there. Fixing the gate is upstream's call; starting at 1024 keeps this bench clear of it. Verified 36/36, err=0, 2207-7003 TFLOPS. mla_v4_prefill is pinned to n=1024 and mla_v4_prefill_fp8 drops nnz_prefix=8192; both are kernel/verify faults, measured and documented at the pins. Co-authored-by: Yu <jiaolyu@amd.com>
JiaoliangYu
added a commit
that referenced
this pull request
Sep 3, 2026
* perf(gfx1250): add the combined gfx1250 microbench (#5076) Two files, both new, on top of main: op_tests/bench_gfx1250_combo.py combined bench op_tests/triton_tests/attention/test_mla_v4_triton.py MLA v4 triton reference The bench imports the top-level @benchmark sweep fns from the op_tests siblings and runs each over its own shape axes, printing only the per-op summary tables -- all the JIT/ROCTracer/import noise is silenced at the fd level. Two suites: --perf for hardware-oriented single-op numbers, --dsv4 for the DeepSeek-V4 shapes. Token sweeps come from AITER_BENCH_TOKENS; ops whose axis means something else, or whose usable range is fixed, pin their sweep in the source and say why. Everything else the branch used to carry has reached main on its own since: test_pa_sparse_prefill.py (#4926), test_mega_moe_gfx1250.py (#5052), test_flydsl_qk_norm_rope_quant.py and the mla_v4 hsa artifacts. Only these two files were ever unique to it, so this is what is left to upstream. Notes worth carrying, all measured on gfx1250 / 20260827-28: a16w16 does not pre-check the >4 GiB operand limit. That guard belongs to one fallback path -- opus_dispatch_a16w16_gfx1250 searches the tuned table first and only reaches check_shape_4g on a miss, en route to the split-K kid whose launcher builds the 32-bit gmem descriptors. A tuned 4wave_wl_co winner never gets there. Predicting it in Python skipped shapes that tuning had already made runnable, so the kernel is left to raise and the exception is recorded as a row. It also checks the error ratio the UT returns against _A16W16_MAX_ERR: all four M=65536 shapes come back err=0.96-0.99 while every other row is 0 or ~1e-5, and nothing in the UT raises or warns, so those used to print as data. a8w8_blockscale sweeps M from 1024. Below that, get_CKGEMM_config's M -> get_padded_m -> nextPow2 retry lands on #4773's M=16/M=64 gluon rows, and the UT's extra "ck strided x_scale" check (line 120, added by #4406 and gated on ck_preshuffle alone) hands triton a stride != 1 specialization that fails to compile in make_llir. The mxfp8_128 path declares its layout with is_x_scale_transposed=True and never reads the stride, so that check tests nothing there. Fixing the gate is upstream's call; starting at 1024 keeps this bench clear of it. Verified 36/36, err=0, 2207-7003 TFLOPS. mla_v4_prefill is pinned to n=1024 and mla_v4_prefill_fp8 drops nnz_prefix=8192; both are kernel/verify faults, measured and documented at the pins. Co-authored-by: Yu <jiaolyu@amd.com> * perf(gfx1250): let AITER_BENCH_TOKENS reach every op (#5082) Three token constants were hardcoded and never consulted the variable -- _INVERSE_ROPE_TOKENS, _MLA_DECODE_TOKENS, _MEGA_MOE_TOKENS -- with comments saying so ("Pinned, not env-driven", "deliberately not consulted"). The reasoning was that a global token count means the wrong thing for those ops, which is true of the default but not of an explicit request: if a caller sets the variable, that is their decision to make, including asking for a shape the op is known to fail on. All nine token constants now go through _tokens(): unset, each op runs its own default and says at its constant why that default is not the shared list; set, the variable wins everywhere and the file does not argue with it. Behaviour with the variable unset is unchanged. The module docstring said "Two ops ignore it and pin their sweep in the source" -- no longer true, and it only listed two of the five ops whose default differs. It now lists all five with the reason for each. Should have been part of #5076; amended in after that PR was pushed. Co-authored-by: Yu <jiaolyu@amd.com> * perf(gfx1250): sweep the fp4 dispatch wire beside bf16 (#5147) mori_ep measured only the bf16 wire, which is not the leg DSv4 serves on. The receiver hands the dispatched payload straight to the expert GEMM as its A operand, and that GEMM is a4w4 -- ATOM's serve script pins MEGA_WIRE=fp4 with AITER_FORCE_A8W4=0 -- so bf16 alone measures a path production does not take, and misses the wire whose payload is a quarter as wide. Sweeps bf16 and fp4 by default, one child each: bench_ep.py reads $DISP once at import and builds the transport for that dtype, so the tiers cannot share a process. $DISP still overrides, now comma-separated, and is forwarded unvalidated -- mori owns the value set, this file does not restate it. fp4 rows are labelled UNCHECKED. mori forces its own comparison off on fp4 ("fp4 combine is too lossy to compare"), so CHECK=1 is silently dropped there; without the label a passing fp4 row reads as verified next to a bf16 row that actually was. Co-authored-by: Yu <jiaolyu@amd.com> * perf(gfx1250): follow the sparse-prefill UT through #5084 (#5170) #5084 rewrote test_pa_sparse_prefill.py and mla_v4_prefill stopped running against it: --nnz-prefix is gone, so the child aborted with "unrecognized arguments" before reaching a kernel. - drop --nnz-prefix. The nnz axis no longer exists: the CSR is generated from --mode under --seed (sparse draws a random nnz per row, dense fills every row), so nnz is an outcome, not an input - delete the mla_v4_prefill_fp8 op. Its entire content was that nnz sweep, and it has no equivalent under the new UT. mla_v4_prefill already covers fp8 -- and now covers the backends too, since #5084 compares opus/asm/triton - _space_table: decide a data row by counting numeric fields instead of testing the first one. The new table leads with prec/mode (bf16, dense), so the first-field test dropped the whole table and the op reported "no result rows" while the UT itself had exited 0 Raise _MLA_PREFILL_TOKENS to 1024..16384. Re-measured on the new UT (b45-2, one process per tier): 1024 through 16384 all clean, where 4096/8192/16384 faulted on the old one. 16384 is the DSv4 prefill chunk and the reason this op exists, so this is the coverage the pin was costing us. 65536 stays out. It faults with a memory access at 0x7f2ddbec0000 and writes an 89 GB coredump doing it -- a third of the free disk on a shared host -- and it is past the chunk size the model prefills anyway. Both the tier sweep and this op run --no-verify, hardcoded at the call site, and that is now load-bearing rather than a speed choice: with verification on, even n=1024 dies at the first case (fp8/dense, fault at 0x43000), so the reference or the comparison is what breaks, not the kernel under test. These are therefore timings from an unverified kernel. The comment says so, because a16w16's M=65536 rows looked exactly this healthy until _A16W16_MAX_ERR was added and caught err=0.99. Co-authored-by: Yu <jiaolyu@amd.com> * perf(gfx1250): stop every child from shelling out to rocminfo (#5180) chip_info runs rocminfo twice per process -- once for the arch, once for the CU count -- and rocminfo takes a per-device rocm_smi mutex on its way in. One process is fine, and the nine single-GPU ops never noticed. A torchrun op starts four ranks at once, and they contend for that mutex. Both multi-GPU ops have now lost a run to it. On b45-1 a rank lost the mutex and aborted: init_mutex /rocm_smi_renderD128: unlock timed lock, ret: 1 terminate called after throwing an instance of 'std::runtime_error' what(): Allgather operation failed The allgather is the symptom, not the cause: the rank died first and took the collective with it. On b45-2 the other shape of the same problem -- four rocminfo processes sat in the mutex for minutes, one wedged in D state, the op never produced a line, and even docker stop hung waiting for the driver to let go. Detect once here, where the call is serial, and hand the answer down: GPU_ARCHS -> get_gfx_list() skips _detect_native() CU_NUM -> get_cu_num_custom_op() skips its own rocminfo Both are read from the environment before either shells out, so a child that inherits them runs no rocminfo at all. Set on os.environ in main() for the children that inherit our environment, and setdefault() in _run_child for the ones handed an explicit env -- setdefault throughout, so an exported value from the caller still wins. The value is what this process detected on this machine, not a hardcoded gfx1250, so it is the same answer the child would have computed. Verified in the child's /proc/<pid>/environ: GPU_ARCHS=gfx1250, AITER_GPU_ARCHS=gfx1250, CU_NUM=256, matching get_gfx()/get_cu_num() here. mega_moe on b45-1 after the change: rc=0, all 36 configs, zero mutex warnings, and zero rocminfo processes observed while it ran. Before it, the same op on the same host was rc=1 with 6 mutex warnings. Co-authored-by: Yu <jiaolyu@amd.com> * test_common：data generate --------- Co-authored-by: Yu <jiaolyu@amd.com>
JiaoliangYu
added a commit
that referenced
this pull request
Sep 13, 2026
* Gfx1250/microbench (#5241) * perf(gfx1250): add the combined gfx1250 microbench (#5076) Two files, both new, on top of main: op_tests/bench_gfx1250_combo.py combined bench op_tests/triton_tests/attention/test_mla_v4_triton.py MLA v4 triton reference The bench imports the top-level @benchmark sweep fns from the op_tests siblings and runs each over its own shape axes, printing only the per-op summary tables -- all the JIT/ROCTracer/import noise is silenced at the fd level. Two suites: --perf for hardware-oriented single-op numbers, --dsv4 for the DeepSeek-V4 shapes. Token sweeps come from AITER_BENCH_TOKENS; ops whose axis means something else, or whose usable range is fixed, pin their sweep in the source and say why. Everything else the branch used to carry has reached main on its own since: test_pa_sparse_prefill.py (#4926), test_mega_moe_gfx1250.py (#5052), test_flydsl_qk_norm_rope_quant.py and the mla_v4 hsa artifacts. Only these two files were ever unique to it, so this is what is left to upstream. Notes worth carrying, all measured on gfx1250 / 20260827-28: a16w16 does not pre-check the >4 GiB operand limit. That guard belongs to one fallback path -- opus_dispatch_a16w16_gfx1250 searches the tuned table first and only reaches check_shape_4g on a miss, en route to the split-K kid whose launcher builds the 32-bit gmem descriptors. A tuned 4wave_wl_co winner never gets there. Predicting it in Python skipped shapes that tuning had already made runnable, so the kernel is left to raise and the exception is recorded as a row. It also checks the error ratio the UT returns against _A16W16_MAX_ERR: all four M=65536 shapes come back err=0.96-0.99 while every other row is 0 or ~1e-5, and nothing in the UT raises or warns, so those used to print as data. a8w8_blockscale sweeps M from 1024. Below that, get_CKGEMM_config's M -> get_padded_m -> nextPow2 retry lands on #4773's M=16/M=64 gluon rows, and the UT's extra "ck strided x_scale" check (line 120, added by #4406 and gated on ck_preshuffle alone) hands triton a stride != 1 specialization that fails to compile in make_llir. The mxfp8_128 path declares its layout with is_x_scale_transposed=True and never reads the stride, so that check tests nothing there. Fixing the gate is upstream's call; starting at 1024 keeps this bench clear of it. Verified 36/36, err=0, 2207-7003 TFLOPS. mla_v4_prefill is pinned to n=1024 and mla_v4_prefill_fp8 drops nnz_prefix=8192; both are kernel/verify faults, measured and documented at the pins. Co-authored-by: Yu <jiaolyu@amd.com> * perf(gfx1250): let AITER_BENCH_TOKENS reach every op (#5082) Three token constants were hardcoded and never consulted the variable -- _INVERSE_ROPE_TOKENS, _MLA_DECODE_TOKENS, _MEGA_MOE_TOKENS -- with comments saying so ("Pinned, not env-driven", "deliberately not consulted"). The reasoning was that a global token count means the wrong thing for those ops, which is true of the default but not of an explicit request: if a caller sets the variable, that is their decision to make, including asking for a shape the op is known to fail on. All nine token constants now go through _tokens(): unset, each op runs its own default and says at its constant why that default is not the shared list; set, the variable wins everywhere and the file does not argue with it. Behaviour with the variable unset is unchanged. The module docstring said "Two ops ignore it and pin their sweep in the source" -- no longer true, and it only listed two of the five ops whose default differs. It now lists all five with the reason for each. Should have been part of #5076; amended in after that PR was pushed. Co-authored-by: Yu <jiaolyu@amd.com> * perf(gfx1250): sweep the fp4 dispatch wire beside bf16 (#5147) mori_ep measured only the bf16 wire, which is not the leg DSv4 serves on. The receiver hands the dispatched payload straight to the expert GEMM as its A operand, and that GEMM is a4w4 -- ATOM's serve script pins MEGA_WIRE=fp4 with AITER_FORCE_A8W4=0 -- so bf16 alone measures a path production does not take, and misses the wire whose payload is a quarter as wide. Sweeps bf16 and fp4 by default, one child each: bench_ep.py reads $DISP once at import and builds the transport for that dtype, so the tiers cannot share a process. $DISP still overrides, now comma-separated, and is forwarded unvalidated -- mori owns the value set, this file does not restate it. fp4 rows are labelled UNCHECKED. mori forces its own comparison off on fp4 ("fp4 combine is too lossy to compare"), so CHECK=1 is silently dropped there; without the label a passing fp4 row reads as verified next to a bf16 row that actually was. Co-authored-by: Yu <jiaolyu@amd.com> * perf(gfx1250): follow the sparse-prefill UT through #5084 (#5170) #5084 rewrote test_pa_sparse_prefill.py and mla_v4_prefill stopped running against it: --nnz-prefix is gone, so the child aborted with "unrecognized arguments" before reaching a kernel. - drop --nnz-prefix. The nnz axis no longer exists: the CSR is generated from --mode under --seed (sparse draws a random nnz per row, dense fills every row), so nnz is an outcome, not an input - delete the mla_v4_prefill_fp8 op. Its entire content was that nnz sweep, and it has no equivalent under the new UT. mla_v4_prefill already covers fp8 -- and now covers the backends too, since #5084 compares opus/asm/triton - _space_table: decide a data row by counting numeric fields instead of testing the first one. The new table leads with prec/mode (bf16, dense), so the first-field test dropped the whole table and the op reported "no result rows" while the UT itself had exited 0 Raise _MLA_PREFILL_TOKENS to 1024..16384. Re-measured on the new UT (b45-2, one process per tier): 1024 through 16384 all clean, where 4096/8192/16384 faulted on the old one. 16384 is the DSv4 prefill chunk and the reason this op exists, so this is the coverage the pin was costing us. 65536 stays out. It faults with a memory access at 0x7f2ddbec0000 and writes an 89 GB coredump doing it -- a third of the free disk on a shared host -- and it is past the chunk size the model prefills anyway. Both the tier sweep and this op run --no-verify, hardcoded at the call site, and that is now load-bearing rather than a speed choice: with verification on, even n=1024 dies at the first case (fp8/dense, fault at 0x43000), so the reference or the comparison is what breaks, not the kernel under test. These are therefore timings from an unverified kernel. The comment says so, because a16w16's M=65536 rows looked exactly this healthy until _A16W16_MAX_ERR was added and caught err=0.99. Co-authored-by: Yu <jiaolyu@amd.com> * perf(gfx1250): stop every child from shelling out to rocminfo (#5180) chip_info runs rocminfo twice per process -- once for the arch, once for the CU count -- and rocminfo takes a per-device rocm_smi mutex on its way in. One process is fine, and the nine single-GPU ops never noticed. A torchrun op starts four ranks at once, and they contend for that mutex. Both multi-GPU ops have now lost a run to it. On b45-1 a rank lost the mutex and aborted: init_mutex /rocm_smi_renderD128: unlock timed lock, ret: 1 terminate called after throwing an instance of 'std::runtime_error' what(): Allgather operation failed The allgather is the symptom, not the cause: the rank died first and took the collective with it. On b45-2 the other shape of the same problem -- four rocminfo processes sat in the mutex for minutes, one wedged in D state, the op never produced a line, and even docker stop hung waiting for the driver to let go. Detect once here, where the call is serial, and hand the answer down: GPU_ARCHS -> get_gfx_list() skips _detect_native() CU_NUM -> get_cu_num_custom_op() skips its own rocminfo Both are read from the environment before either shells out, so a child that inherits them runs no rocminfo at all. Set on os.environ in main() for the children that inherit our environment, and setdefault() in _run_child for the ones handed an explicit env -- setdefault throughout, so an exported value from the caller still wins. The value is what this process detected on this machine, not a hardcoded gfx1250, so it is the same answer the child would have computed. Verified in the child's /proc/<pid>/environ: GPU_ARCHS=gfx1250, AITER_GPU_ARCHS=gfx1250, CU_NUM=256, matching get_gfx()/get_cu_num() here. mega_moe on b45-1 after the change: rc=0, all 36 configs, zero mutex warnings, and zero rocminfo processes observed while it ran. Before it, the same op on the same host was rc=1 with 6 mutex warnings. Co-authored-by: Yu <jiaolyu@amd.com> * test_common：data generate --------- Co-authored-by: Yu <jiaolyu@amd.com> * Blyu/opus gemm benchmark update (#5250) * [test] Add bench controls to test_opus_a16w16_gemm Add configurable warmup/iteration counts, data-init modes (zero/constant/uniform/norm) with a seed, and rotating-tensor control to the opus a16w16 GEMM test, and report both TFLOPs and bandwidth (TB/s) explicitly across all three timing paths. The opus graph sweep table now also surfaces the tuned CSV's splitK and kernelName (kid) per shape, so the algorithmic winner is visible and tracks CSV updates. test_a16w16 keeps a backward-compatible signature so the gfx1250 combo bench is unaffected. Co-authored-by: Cursor <cursoragent@cursor.com> * [test] Fix ruff C408 in test_opus_a16w16_gemm Rewrite init_kwargs dict() call as a dict literal to satisfy `ruff check` (C408); black is unchanged. Co-authored-by: Cursor <cursoragent@cursor.com> * test_common：data generate * [test] Adopt shared test_common data-init in a16w16 gemm Use the cherry-picked aiter.test_common data-init API (make_generator / add_data_init_args / fill / DATA_DISTS) in test_opus_a16w16_gemm.py instead of the local DATA_INITS/_make_generator/_make_tensor helpers, so operand init and --data-init/--seed CLI match every other op test. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: binglue.lyu <blyu@heliosp-1b114-d07-2.mnb.dcgpu> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: JiaoliangYu <jiaolyu@amd.com> Co-authored-by: blyu <blyu@amd.com> * [ASM][gfx1250] enhance ASM GEMM microbench (#5251) * test(gemm): use common data initialization options for gfx1250 asm gemm Signed-off-by: Double Young <yang.yang2@amd.com> * test(gemm): expose warmup/iters/graph/rotate CLIs for gfx1250 asm gemm Signed-off-by: Double Young <yang.yang2@amd.com> * test(gemm): show algorithmic details for gfx1250 asm gemm Signed-off-by: Double Young <yang.yang2@amd.com> --------- Signed-off-by: Double Young <yang.yang2@amd.com> * test mega moe data init and warmup (#5248) * Dev/ubench gemm (#5252) * update * update * support a8w8 ubench * revert ck module * apply aiter common interface --------- Co-authored-by: yadaish <yadai@amd.com> * [Bench] Add data-init options to qk_norm_rope_quant test (#5253) Add --seed, --init {normal,uniform,zero,constant}, --init-scale, and --init-val CLI flags so callers can control how input tensors are initialised, enabling reproducible perf sweeps across data distributions. Co-authored-by: jli-melchior <Jun.Li@amd.com> Co-authored-by: Claude <noreply@anthropic.com> * [HIP] Add data-init to test_mhc (#5245) Wire test_common.fill, --data-init, and --seed so MHC inputs are reproducible. Co-authored-by: Cursor <cursoragent@cursor.com> * perf(gfx1250): add ubench init controls and SMI monitor (#5257) Co-authored-by: JiaoliangYu <jiaolyu@amd.com> * [Perf][gfx1250] Complete DSv4 attention data-init support (#5258) * perf(gfx1250): complete DSv4 data init controls * perf(gfx1250): report per-case SMI telemetry * perf(gfx1250): monitor Mega MoE per rank --------- Co-authored-by: JiaoliangYu <jiaolyu@amd.com> * perf(gfx1250): add M512 A8W8 DSv4 case (#5284) Co-authored-by: JiaoliangYu <jiaolyu@amd.com> * [Perf][gfx1250] Emit ubench summaries as JSON (#5288) * docs(gfx1250): document bundled amdsmi binding * perf(gfx1250): emit benchmark summaries as JSON --------- Co-authored-by: JiaoliangYu <jiaolyu@amd.com> * fix(gfx1250): scope bundled amdsmi import (#5291) Co-authored-by: JiaoliangYu <jiaolyu@amd.com> * perf(gfx1250): reuse process-wide amdsmi state (#5302) Co-authored-by: JiaoliangYu <jiaolyu@amd.com> * [Perf][gfx1250] Label SMI replays per benchmark call (#5304) * perf(gfx1250): label SMI replays per benchmark call * perf(gfx1250): cover saturated DSv4 MLA decode KV * perf(gfx1250): sweep DSv4 score-QK decode batches * perf(gfx1250): cover TP4 inverse RoPE shape * perf(gfx1250): add M256 blockscale coverage * perf(gfx1250): add 64K Mega MoE coverage --------- Co-authored-by: JiaoliangYu <jiaolyu@amd.com> * [ASM][gfx1250] mla v4 prefill: rebuild the sparse_pfl kernel with the head schedule and the PTL prefetch hazard fixed Replaces mla_a8w8_qh128_1tg_32mx4_32nx1_sparse_pfl.co (462a5ee0, as shipped by #4926) with a build carrying the 09-07/09-08 head work plus one correctness fix. The correctness fix. The page-table prefetch built its row offset from a VGPR written by v_writelane_b32 at the end of KV_page_prefetch_build; the consumers are the two v_lshlrev at the top of KV_page_prefetch_issue, on the far side of the _s_p_tiles branch with zero VALU in between: v_writelane_b32 v4, s60, 0 s_cmp_eq_u32 s103, 0 s_cbranch_scc0 label_1028 v_lshlrev_b32 v2, 9, v4 <- 2 SALU after the write SCHED_MODE 2 has no VALU->VALU interlock, so the shift could read the stale register (Q addr-gen leftovers). `stale << 9` off a bounds-free 64-bit base is a wild address, observed under batched prefill as Memory access fault by GPU node-3 ... Reason: Page not present No numeric test can catch this: a prefetch has no destination register, so the bad address never reaches a result. Fixed with one s_wait_alu placed before the branch, covering both the fall-through (extend) and taken (prefix) arms. +4 bytes; the core loop is byte-identical. Also in this build, relative to 462a5ee0: * head schedule: HEAD_INDPTR_EARLY, HEAD_INDEX_HOIST, HEAD_TDM_SLOT_AWARE * HEAD_KV_PTL_* page-table prefetch over tile 0's KV rows * the packed-row scale tail mask (448:511 pad), which 462a5ee0 lacks -- that pad reads as NaN once the host stops zero-filling it Source: junxiaguo/my_agent_skills @ 495b486 (0805_base/32mx4_32nx1.sp3). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * support print overlap ratio (#5311) * [ASM][gfx1250] mla v4 prefill: update sparse_pfl kernel binary Replaces the mla_a8w8_qh128_1tg_32mx4_32nx1_sparse_pfl.co binary added in e566d68 with a newer build. Validated end-to-end on DeepSeek-V4-Pro (gfx1250, tp4, expert parallel, dp attention, kv_cache_dtype=fp8, index_cache_dtype=fp8): lm_eval gsm8k, full 1319 questions, 3-shot flexible-extract exact_match 0.9507 +/- 0.006 strict-match exact_match 0.9500 +/- 0.006 No memory access faults, page-not-present or HSA_STATUS errors in the server log across the full run. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * Revert the two sparse_pfl kernel binary updates This reverts 12f16de ("[ASM][gfx1250] mla v4 prefill: update sparse_pfl kernel binary") and e566d68 ("[ASM][gfx1250] mla v4 prefill: rebuild the sparse_pfl kernel with the head schedule and the PTL prefetch hazard fixed"), restoring hsa/gfx1250/mla_v4/mla_a8w8_qh128_1tg_32mx4_32nx1_sparse_pfl.co to 462a5ee0 -- the build shipped by #4926. No other file is touched; both reverted commits changed only that binary. The kernel work is being resubmitted against main directly rather than landing here first. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * [Perf][gfx1250] Use tuned paths in combo benchmarks (#5323) * perf(gfx1250): compare MegaMoE modes in one run * perf(gfx1250): honor tuned backend in a16w16 bench * perf(gfx1250): add Mori EP inputs and telemetry * perf(gfx1250): add 16K Mega MoE coverage --------- Co-authored-by: JiaoliangYu <jiaolyu@amd.com> * style: satisfy black and ruff checks * test: keep MLA benchmark entrypoints out of pytest * test: use architecture fp8 dtype for blockscale gemm * test: unify FlyDSL grouped GEMM coverage * test: unify architecture-specific Mega MoE runners * test: preserve grouped GEMM suites behind arch gates * bench: default combo output to readable tables * test: place unified FlyDSL MoE coverage under flydsl tests * style: remove stale shebangs from merged MoE test * test: isolate FlyDSL MoE suites by runtime arch * test: split FlyDSL MoE and grouped GEMM suites * ci: leave test sharding config unchanged * test: move grouped GEMM suite under FlyDSL tests * refactor: move SMI monitor into aiter package * test: restore readable benchmark summaries * style: satisfy Black after JSON helper removal * revert: restore JSON benchmark summaries * perf: preserve profiler kernel names * perf: avoid clearing allocator cache per iteration * fix: address microbenchmark review feedback * refactor: split benchmark support helpers Signed-off-by: JiaoliangYu <jiaolyu@amd.com> * fix: serialize benchmark dtypes as strings * style: format benchmark reporting helper Signed-off-by: JiaoliangYu <jiaolyu@amd.com> --------- Signed-off-by: Double Young <yang.yang2@amd.com> Signed-off-by: JiaoliangYu <jiaolyu@amd.com> Co-authored-by: Yu <jiaolyu@amd.com> Co-authored-by: demonsan <39889339+demonsan@users.noreply.github.com> Co-authored-by: binglue.lyu <blyu@heliosp-1b114-d07-2.mnb.dcgpu> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: blyu <blyu@amd.com> Co-authored-by: Double Young <yang.yang2@amd.com> Co-authored-by: yanboshao <yashao@amd.com> Co-authored-by: yadaish <yadai@amd.com> Co-authored-by: jli-melchior <Jun.Li@amd.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: la <junchen2@amd.com> Co-authored-by: junxiaguo <JunXia.Guo@amd.com> Co-authored-by: Xin Huang <Xin.Huang@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
4 participants
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.
Summary
x_scalelayouts.<3.7, cap preshuffle configs at two pipeline stages to avoidTritonAMDGPUConvertToBufferOpsfailures.AITER-pinned Triton 3.7+ keeps its existing three-stage tuned configs unchanged.
Regression Test
The existing benchmark path remains unchanged. An additional unmeasured public-API call validates the strided scale representation against the Torch reference.
Validation
(2048,7168),(7168,4096), and(7168,16384): pass.M=32,N=2048,K=7168M=144,N=8192,K=1536No CK/ASM fallback or CSV routing changes.
Addresses #4323.