[FlyDSL] One-stage split-K for the a8w8 preshuffle GEMM - #5007
Merged
Conversation
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>
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>
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>
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>
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>
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags: |
Both arms opened with the same fx.copy; only what follows it differs. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
XiaobingSuper
marked this pull request as ready for review
August 26, 2026 07:16
github-actions
Bot
changed the title
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds one-stage split‑K support to the FlyDSL a8w8 preshuffle GEMM path by splitting K over gridDim.z, writing fp32 partials into a workspace, and having the last arriving split reduce/convert to fp16/bf16 in-kernel using a per-tile semaphore. It also retunes Kimi‑K3 CSV entries (primarily for M <= 32) to take advantage of split‑K where it improves occupancy/perf.
Changes:
- Implement split‑K in the FlyDSL preshuffle GEMM kernel, including coherent partial stores + in-kernel reduction gated by agent-scope atomics.
- Extend kernel-name parsing (runtime + AOT) to recognize an optional trailing
_ksNsuffix and propagatesplit_kinto the launcher. - Update tuner/task generation and the Kimi‑K3 tuned CSV to benchmark/select split‑K variants for qualifying shapes.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| csrc/ck_gemm_a8w8_bpreshuffle/gemm_a8w8_bpreshuffle_tune.py | Adds split‑K candidate enumeration for FlyDSL tuning tasks; disables CKTile split‑K sweep. |
| aiter/ops/gemm_op_a8w8.py | Parses _ksN from FlyDSL kernel names and passes split_k into the preshuffle launcher. |
| aiter/ops/flydsl/kernels/preshuffle_gemm.py | Implements split‑K kernel path (workspace partials + semaphore + reduction) and launches gridDim.z = split_k. |
| aiter/ops/flydsl/gemm_tune/flydsl_gemm_a8w8_bpreshuffle_common.py | Adds k_split to kernelInstance naming and shape-fit checks; introduces k_split_candidates(). |
| aiter/ops/flydsl/gemm_kernels.py | Allocates/caches split‑K workspace+semaphore and wires new launcher signature and split‑K validation. |
| aiter/configs/model_configs/a8w8_bpreshuffle_tuned_gemm_kimik3.csv | Updates tuned entries to new/faster FlyDSL configs, including split‑K variants (splitK>0 + _ksN). |
| aiter/aot/flydsl/gemm.py | Updates AOT kernel-name regex and compile inputs to include optional split‑K workspace/semaphore args. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
XiaobingSuper
changed the title
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>
github-actions
Bot
changed the title
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.
Suppressed comments (1)
aiter/ops/flydsl/gemm_kernels.py:1006
compile_preshuffle_gemmreshapes A/C with a hardM_max = 65536(seepreshuffle_gemm.py), sosplit_k > 1should reject largermvalues up front. As written,_check_preshuffle_split_capacitycan acceptm > 65536ifsplit_k * m * nfits the workspace, but the kernel’s 2D views will then be undersized and addressing becomes incorrect.
tiles = ((m + tile_m - 1) // tile_m) * (n // tile_n)
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>
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>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.
Suppressed comments (2)
aiter/ops/flydsl/kernels/preshuffle_gemm.py:842
- After determining this CTA is the last arrival, it immediately loads all split planes from the workspace. Add an acquire fence before those loads so they’re ordered after observing the semaphore update (paired with the release fence before the atomic_add_agent).
vecs_per_row = tile_n // _REDUCE_VEC
vecs_per_tile = tile_m * vecs_per_row
aiter/ops/flydsl/gemm_kernels.py:982
- This comment says “the kernel asserts they agree”, but there’s no assertion in the FlyDSL kernel for PRESHUFFLE_M_MAX. Either add a real check at compile time, or adjust the comment to just state that the constant must be kept in sync.
PRESHUFFLE_M_MAX = 65536
PRESHUFFLE_SPLIT_K_MAX_TILES = 256
coderfeli
reviewed
Aug 27, 2026
coderfeli
reviewed
Aug 27, 2026
Collaborator
|
Some comments inline. Others look good. |
xytpai
reviewed
Aug 27, 2026
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>
github-actions
Bot
changed the title
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.
Suppressed comments (1)
Previously missed (1) — in code that hasn't changed since the last review.
aiter/ops/flydsl/gemm_kernels.py:982
- PRESHUFFLE_SPLIT_K_MAX_TILES is hard-coded to 256, but k_split_candidates is parameterized by cu_num and the repo supports GPUs with cu_num=304 (gfx942). For shapes where the tile grid is in (256, 304), split-K candidates can be generated but will then fail at runtime with the semaphore-capacity check.
Consider sizing this fixed buffer to the maximum supported CU count (at least 304) so candidate generation and runtime capacity limits are consistent.
PRESHUFFLE_SPLIT_K_MAX_TILES = 256
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.
github-actions
Bot
changed the title
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.
Suppressed comments (2)
Previously missed (2) — in code that hasn't changed since the last review.
csrc/ck_gemm_a8w8_bpreshuffle/gemm_a8w8_bpreshuffle_tune.py:41
- The FlyDSL common module is imported four times inside the same try-block. This is redundant and makes it harder to see what symbols are needed vs optional; it can be collapsed into a single import statement.
from aiter.ops.flydsl.gemm_tune.flydsl_gemm_a8w8_bpreshuffle_common import (
PIPELINES as FLYDSL_PIPELINES,
)
from aiter.ops.flydsl.gemm_tune.flydsl_gemm_a8w8_bpreshuffle_common import (
k_split_candidates,
csrc/ck_gemm_a8w8_bpreshuffle/gemm_a8w8_bpreshuffle_tune.py:426
- This loop only ever runs with splitK==0 (range(1)), which obscures intent. Using an explicit singleton iterable makes it clear that split-K is intentionally disabled for cktile here.
for splitK in range(1):
xytpai
approved these changes
Aug 27, 2026
coderfeli
approved these changes
Aug 27, 2026
XiaobingSuper
added a commit
that referenced
this pull request
Aug 28, 2026
The a8w8 rows in this file were tuned before #4151 renamed the FlyDSL kernels, and that PR retuned four CSVs but not this one. All 49 of its FlyDSL rows have failed to parse ever since: _parse_flydsl_kernel_name returns None for the old five-field name and the caller quietly falls back to the default CK kernel, so those shapes have been running untuned. Retuning is what actually fixes them; the parse-failure path is silent by design and is left for a separate change. Split-K is now in the search space (#5007), and 47 of the 134 rows pick splitK > 0. Running the op under each config on the 49 shapes -- old being the CK fallback those rows really reach today, not the kernel they name -- puts the new config ahead by a median of 24.4%, 3618 -> 2722us in total, with one shape 1.1% behind. M=2 N=2624 K=6144 goes 13.9 -> 6.9us. None of the winners is cktile, though the tuner picked it for seven small-M shapes at K=512. A cktile row cannot reproduce its tuned result: gemm_a8w8_bpreshuffle_cktile takes no kernelId, so the CSV column is dead data and the kernel comes from the compile-time lookup table in gemm_a8w8_bpreshuffle_cktile.cu. That table has 18 entries and none of these shapes, so all three lookups miss and dispatch lands on rowwise_heuristic_dispatch, which returns a fixed 128x128x128 kernel regardless of M. Measured, those seven rows cost 10.6-11.7us against the 2.6-3.9us the tuner recorded. They now carry the best FlyDSL candidate from the same profile run instead, which measures 3.4-4.0us -- a 2-15% loss on paper to avoid a 3x loss in practice. Making cktile honour its tuned kernel is a backend change, not a config one. New shapes for both TP4 and TP8: a8w8 gains N=2688/K=6144 and N=6144/K=12288 (TP4) plus N=2048/K=2048, N=3072/K=6144 and N=3584/K=512 (TP8); BF16 gains N=160, N=256 and N=38720 at K=6144 (TP4) plus N=19360 (TP8). BF16 keeps only powers of two for M on the two widest new groups, where the intermediate sizes are not shapes the model runs. Two BF16 shapes stay untuned, M=384 and M=768 at N=256 K=6144: each carries ~10.6k FlyDSL candidates and the JIT runs out of code-region memory partway through, independently of host RAM or VRAM. They are left in the untuned CSV so a later run retries them. GLM-5.2 TP4 gsm8k scores 0.9704 +/- 0.0047 exact_match on both flexible-extract and strict-match, 1319/1319 answered. That run predates the cktile replacement above, which changes which kernel is selected for seven shapes but not the numerics of the op. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
XiaobingSuper
added a commit
that referenced
this pull request
Aug 28, 2026
The a8w8 rows in this file were tuned before #4151 renamed the FlyDSL kernels, and that PR retuned four CSVs but not this one. All 49 of its FlyDSL rows have failed to parse ever since: _parse_flydsl_kernel_name returns None for the old five-field name and the caller quietly falls back to the default CK kernel, so those shapes have been running untuned. Retuning is what actually fixes them; the parse-failure path is silent by design and is left for a separate change. Split-K is now in the search space (#5007), and 47 of the 134 rows pick splitK > 0. Running the op under each config on the 49 shapes -- old being the CK fallback those rows really reach today, not the kernel they name -- puts the new config ahead by 3606 -> 2702us in total, a median of 22.8%, with no shape behind by more than 0.2%. M=2 N=2624 K=6144 goes 13.9 -> 6.9us. Measuring this needs one non-obvious step. gen_instances.py compiles the tuned CSV into the lookup table that ck and cktile dispatch through, but its output is not part of the JIT build signature, so editing a tuned CSV never invalidates an existing module. Against a module built before these rows existed, all fourteen ck and cktile rows miss the table and land on rowwise_heuristic_dispatch, which returns one fixed kernel regardless of M -- the cktile rows measure 10.6-11.7us that way against the 2.6-3.9us the tuner recorded. Deleting aiter/jit/module_*.so and aiter/jit/build/module_*/ after updating a config rebuilds the table; the numbers above are from a rebuilt module. FlyDSL rows are immune because they reconstruct the kernel from kernelName at runtime. GLM-5.2 TP4 gsm8k scores 0.9704 +/- 0.0047 exact_match on both flexible-extract and strict-match, 1319/1319 answered. New shapes for both TP4 and TP8: a8w8 gains N=2688/K=6144 and N=6144/K=12288 (TP4) plus N=2048/K=2048, N=3072/K=6144 and N=3584/K=512 (TP8); BF16 gains N=160, N=256 and N=38720 at K=6144 (TP4) plus N=19360 (TP8). BF16 keeps only powers of two for M on the two widest new groups, where the intermediate sizes are not shapes the model runs. Two BF16 shapes stay untuned, M=384 and M=768 at N=256 K=6144: each carries ~10.6k FlyDSL candidates and the JIT runs out of code-region memory partway through, independently of host RAM or VRAM. They are left in the untuned CSV so a later run retries them. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
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>
XiaobingSuper
added a commit
that referenced
this pull request
Aug 28, 2026
The a8w8 rows in this file were tuned before #4151 renamed the FlyDSL kernels, and that PR retuned four CSVs but not this one. All 49 of its FlyDSL rows have failed to parse ever since: _parse_flydsl_kernel_name returns None for the old five-field name and the caller quietly falls back to the default CK kernel, so those shapes have been running untuned. Retuning is what actually fixes them; the parse-failure path is silent by design and is left for a separate change. Split-K is now in the search space (#5007), and 47 of the 134 rows pick splitK > 0. Running the op under each config on the 49 shapes -- old being the CK fallback those rows really reach today, not the kernel they name -- puts the new config ahead by 3606 -> 2702us in total, a median of 22.8%, with no shape behind by more than 0.2%. M=2 N=2624 K=6144 goes 13.9 -> 6.9us. Measuring this needs one non-obvious step. gen_instances.py compiles the tuned CSV into the lookup table that ck and cktile dispatch through, but its output is not part of the JIT build signature, so editing a tuned CSV never invalidates an existing module. Against a module built before these rows existed, all fourteen ck and cktile rows miss the table and land on rowwise_heuristic_dispatch, which returns one fixed kernel regardless of M -- the cktile rows measure 10.6-11.7us that way against the 2.6-3.9us the tuner recorded. Deleting aiter/jit/module_*.so and aiter/jit/build/module_*/ after updating a config rebuilds the table; the numbers above are from a rebuilt module. FlyDSL rows are immune because they reconstruct the kernel from kernelName at runtime. GLM-5.2 TP4 gsm8k scores 0.9704 +/- 0.0047 exact_match on both flexible-extract and strict-match, 1319/1319 answered. New shapes for both TP4 and TP8: a8w8 gains N=2688/K=6144 and N=6144/K=12288 (TP4) plus N=2048/K=2048, N=3072/K=6144 and N=3584/K=512 (TP8); BF16 gains N=160, N=256 and N=38720 at K=6144 (TP4) plus N=19360 (TP8). BF16 keeps only powers of two for M on the two widest new groups, where the intermediate sizes are not shapes the model runs. Two BF16 shapes stay untuned, M=384 and M=768 at N=256 K=6144: each carries ~10.6k FlyDSL candidates and the JIT runs out of code-region memory partway through, independently of host RAM or VRAM. They are left in the untuned CSV so a later run retries them. Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
JiaoliangYu
added a commit
that referenced
this pull request
Sep 3, 2026
* [Config] Retune the GLM-5.2 a8w8 and BF16 GEMMs for gfx950 (#5069)
The a8w8 rows in this file were tuned before #4151 renamed the FlyDSL
kernels, and that PR retuned four CSVs but not this one. All 49 of its
FlyDSL rows have failed to parse ever since: _parse_flydsl_kernel_name
returns None for the old five-field name and the caller quietly falls
back to the default CK kernel, so those shapes have been running
untuned. Retuning is what actually fixes them; the parse-failure path
is silent by design and is left for a separate change.
Split-K is now in the search space (#5007), and 47 of the 134 rows pick
splitK > 0. Running the op under each config on the 49 shapes -- old
being the CK fallback those rows really reach today, not the kernel they
name -- puts the new config ahead by 3606 -> 2702us in total, a median
of 22.8%, with no shape behind by more than 0.2%. M=2 N=2624 K=6144 goes
13.9 -> 6.9us.
Measuring this needs one non-obvious step. gen_instances.py compiles the
tuned CSV into the lookup table that ck and cktile dispatch through, but
its output is not part of the JIT build signature, so editing a tuned
CSV never invalidates an existing module. Against a module built before
these rows existed, all fourteen ck and cktile rows miss the table and
land on rowwise_heuristic_dispatch, which returns one fixed kernel
regardless of M -- the cktile rows measure 10.6-11.7us that way against
the 2.6-3.9us the tuner recorded. Deleting aiter/jit/module_*.so and
aiter/jit/build/module_*/ after updating a config rebuilds the table;
the numbers above are from a rebuilt module. FlyDSL rows are immune
because they reconstruct the kernel from kernelName at runtime.
GLM-5.2 TP4 gsm8k scores 0.9704 +/- 0.0047 exact_match on both
flexible-extract and strict-match, 1319/1319 answered.
New shapes for both TP4 and TP8: a8w8 gains N=2688/K=6144 and
N=6144/K=12288 (TP4) plus N=2048/K=2048, N=3072/K=6144 and
N=3584/K=512 (TP8); BF16 gains N=160, N=256 and N=38720 at K=6144 (TP4)
plus N=19360 (TP8). BF16 keeps only powers of two for M on the two
widest new groups, where the intermediate sizes are not shapes the
model runs.
Two BF16 shapes stay untuned, M=384 and M=768 at N=256 K=6144: each
carries ~10.6k FlyDSL candidates and the JIT runs out of code-region
memory partway through, independently of host RAM or VRAM. They are
left in the untuned CSV so a later run retries them.
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
* Upgrade gfx1250 MLA 64nx1 code objects and their host launch contract (#5065)
The 16mx4_64nx1 decode code objects returned wrong results for some combinations
of context length and KV split count. Replace all three (qh16, qh64, qh128) with
current builds.
qh128 additionally needs the host side brought in line with the new code object:
- ABI: qh128 no longer takes the legacy 288B kernarg block. Every gfx1250 MLA
kernel now uses the 120B packed preload ABI, so the qh128 exception in the
dispatch layer is removed.
- Launch strategy: for gqa=128 the two workgroups per (batch, KV split) are
now issued along x (gdx = 2) instead of along z, and z carries only the KV
split id. get_meta_param's occupancy multiplier is unchanged -- the
workgroup count per (batch, split) is still 2 -- so only its comment needed
updating to name the new axis.
Verified on gfx1250: the previously failing (context, split) combinations now
match the fp32 reference at the fp8 quantization floor (cos_diff 1.4e-4..2.3e-4,
no element outside a 6e-2 tolerance) for qh64 (36 configs), qh16 64nx1 (20
configs) and qh128 (29 configs), partially filled last pages included. qh8 and
qh32 32nx4_3p are unaffected.
Co-authored-by: Cursor <cursoragent@cursor.com>
* Tune MoE GEMM A8W8 (#5033)
* [ASM] [HIP] [CK] feat(mha): gfx950 hd256 FP8 LINEAR paged-varlen asm prefill (#4971)
* feat: add gfx950 hd256 FP8 LINEAR paged-varlen asm prefill
Select the PAGED_VARLEN asm kernel for gfx950 FP8 hd256 page_size=64, then fall back to CK.
* fix: honor use_ext_asm and tighten paged-prefill tests
Skip page64 asm off gfx950, drop the redundant page16 case, clamp empty-page seqlen_k, and use the file's FP8 threshold.
* style: match FAV3 eligibility checks in batch-prefill asm
Fold the -1 ladder into one compound if like fmha_fwd_v3, and restore the CK kUseGlobalLoad comment.
* Tune M=48 for the GLM-5.2 a8w8 bpreshuffle GEMM shapes (#5078)
Every row in this config uses a power-of-two M, so an M=48 request has no
tuned entry and get_CKGEMM_config pads it to the M=64 row. That row was
tuned for a different width, so it is only incidentally a good fit. This
adds a tuned M=48 row for eight of the nine (N,K) groups, including three
that predate #5069, so the layer is covered at M=48 rather than borrowing
from M=64.
Tuned with --libtype all -k --shape_grouped on gfx950 (cu_num=256) in a
worktree pinned to the merge commit of #5069, so the FlyDSL candidate list
matches what the config is resolved against. All eight winners are FlyDSL
with errRatio 0; the widest-K groups pick splitK 2 or 4, which is where
most of the gain comes from.
Measured against the padded-to-M=64 behaviour, three runs, per-shape
median of 100 iterations after 20 warmup, one GPU:
N=2048 K=2048 5.838 -> 5.080us +12.98%
N=3584 K=512 4.162 -> 3.138us +24.60%
N=6144 K=12288 20.505 -> 18.898us +7.84%
N=4096 K=2048 6.035 -> 5.622us +6.85%
N=3072 K=6144 9.643 -> 9.176us +4.84%
N=7168 K=512 4.232 -> 4.067us +3.90%
N=2688 K=6144 8.916 -> 8.773us +1.60%
N=2624 K=6144 8.757 -> 8.784us -0.31% (within run-to-run spread)
N=6144 K=4096 is deliberately left out. Its M=64 row uses a tile_m=32
kernel, and the FlyDSL candidate generator offers tile_m in {16, 48, 128,
256} for M=48 -- 32 is not among them. So the best of the 2208 candidates
timed for that shape (9.943us) still loses to what padding already gives
it (9.579us median), and adding the row would cost 3.07%. Leaving the
shape out keeps it on the M=64 row it uses today. Its untuned entry is
removed as well so a later re-run does not silently re-add the regression;
it is worth revisiting if the candidate set grows a tile_m=32 variant.
No existing row is modified -- the diff is eight added lines per file.
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
* [Triton/Gluon] [ASM] [HIP] Block-sparse MHAv4 with load-balancing (#5005)
* perf(mha_v4): avoid copying odd-tail FP6 V inputs
Signed-off-by: jcaraban <jcaraban@amd.com>
* feat(mha_v4): support grouped query attention
Signed-off-by: jcaraban <jcaraban@amd.com>
* feat(mha_v4): add MXFP8 raw entrypoint
Signed-off-by: jcaraban <jcaraban@amd.com>
* docs(mha_v4): clarify grouped-query attention contract
Signed-off-by: jcaraban <jcaraban@amd.com>
* feat(mha_v4): add gfx942 native FP8 kernel
Signed-off-by: jcaraban <jcaraban@amd.com>
* fix(mha_v4): canonicalize rotated FP8 preprocessing
Signed-off-by: jcaraban <jcaraban@amd.com>
* refactor(bench): simplify MHA v4 quantized runners
Signed-off-by: jcaraban <jcaraban@amd.com>
* perf(mha_v4): deploy gfx942 XCD-swizzled kernels
Signed-off-by: jcaraban <jcaraban@amd.com>
* perf(mha_v4): deploy gfx942 block kernels
Signed-off-by: jcaraban <jcaraban@amd.com>
* fix(mha_v4): handle singleton-head rotation strides
Signed-off-by: jcaraban <jcaraban@amd.com>
* perf(mha_v4): deploy retimed gfx942 I8/FP8 kernels
Signed-off-by: jcaraban <jcaraban@amd.com>
* fix(mha_v4): deploy corrected gfx942 PV LDS waits
Signed-off-by: jcaraban <jcaraban@amd.com>
* fix(fmha): deploy gfx942 V staging
Signed-off-by: jcaraban <jcaraban@amd.com>
* fix(fmha): update gfx942 I8FP8 kernel
Signed-off-by: jcaraban <jcaraban@amd.com>
* feat(mha): add bf16 to mha v4
Add raw BF16/NONE dispatch and the gfx950 block kernel to the MHA v4 manifest. Generalize launcher strides to byte units, preserve the v3 aiter_bf16 benchmark, rename v4 benchmark providers to mha4_*, and cover BF16 recipe, finite output, and compiled parity.
Signed-off-by: jcaraban <jcaraban@amd.com>
* perf(fmha): deploy optimized gfx942 block kernels
Signed-off-by: jcaraban <jcaraban@amd.com>
* style(mha_v4): apply repository formatting
Signed-off-by: jcaraban <jcaraban@amd.com>
* test(mha_v4): isolate compile parity cases
Signed-off-by: jcaraban <jcaraban@amd.com>
* fix ruff warnings
Signed-off-by: jcaraban <jcaraban@amd.com>
* fix(mha_v4): enforce contiguous rotation layout
Dense rotation kernels flatten all leading dimensions into rows, so their row stride is the last dimension width rather than stride(-2). PyTorch permits arbitrary stride metadata on singleton dimensions, which made contiguous [B, S, 1, D] inputs report a misleading head-axis stride and caused incorrect row addressing.
Require contiguous dense inputs and outputs, use canonical input/output row widths, and validate output shapes, devices, auxiliary tensors, and empty inputs. Add regression coverage for singleton heads and rejected unsupported layouts.
* fix(mha_v4): update deterministic BF16 kernel
* Revert "fix(mha_v4): handle singleton-head rotation strides"
This reverts e79b1c8 and adds rotate_activation_hd128() to mha_v4 own .cu
Signed-off-by: jcaraban <jcaraban@amd.com>
* Sparse MHAv4 initial commit
Signed-off-by: Niko Säkkinen <niko.sakkinen@amd.com>
* Enable sparse GQA. Fix rebase issues. Fix rotate_activations bug
Signed-off-by: Niko Säkkinen <niko.sakkinen@amd.com>
* Give MHA v4 its own hd128 rotation instead of calling into dsv4
The FP8 raw recipe rotated Q and K through module_dsv4_rotate_quant,
which registers no aiter_tensor_t and so rejects the instance
torch_to_aiter_pybind builds from module_aiter_core. Every mha_v4()
call with an FP8 q/k format failed on that TypeError, block-sparse
ones included. The MX quantizers here already run the same rotation
before quantizing, so hadamard_rotate_kernel stops where they diverge
and emits it in the input dtype: bitwise identical to the dsv4 kernel
it replaces, and not gated on gfx950 since the FP8 recipe also runs on
gfx942.
A new test pins the transform against an explicit Hadamard matmul. An
autouse fixture resets Dynamo per test, because the FP8 compile parity
tests no longer die early and so exhausted the shared recompile limit,
breaking whichever test compiled next.
Signed-off-by: Niko Säkkinen <niko.sakkinen@amd.com>
* Cut the fixed cost of the sorted-sparse work table
Rebuilt on every call at a cost independent of sparsity, so it came to
dominate the packed call as density dropped. Two device syncs came from
reading lut_count back to detect uniform counts, and thirteen ATen ops
packed a few hundred elements. A stable descending sort yields the
identity permutation for uniform counts without that readback, and the
packing is now one kernel. On the shape measured the sparse call is
~2.6x faster at 2% density and ~1.2x at full.
Signed-off-by: Niko Säkkinen <niko.sakkinen@amd.com>
* Build small sorted-sparse work tables in one kernel
Order and pack the table by counting each entry's rank in LDS instead of
calling ATen's sort and packing in a separate launch. The key packs the LUT
length with the slot index, so ranks are distinct and stable by construction,
and the low half is already the permutation the packing needs. Build time at
512 entries drops from ~23us to 9us. The quadratic rank count loses to ATen
past ~1024 entries, so larger tables keep the sort path.
Also expose the builder and test its ordering. A wrong order only unbalances
the waves rather than changing the result, so no attention test can see it.
Signed-off-by: Niko Säkkinen <niko.sakkinen@amd.com>
* Rank work table entries across a wave instead of a thread
Each entry's rank was counted by a single thread walking every key, which costs
O(n) per thread and lost to ATen's sort above about 1024 entries. Split the
count across a wave and reduce it, so per-lane work is n/64 and the build stays
near 6us from 512 entries to 4096. That moves the fused cutoff to 8192, which is
where a workgroup's 64KB of LDS runs out.
Wan 720p self-attention at 5 heads needs 1480 entries and so was on the fallback
branch at 25.7us; it now builds in 5.9us, taking the whole call from 176us to
157us at 1.6% density.
Signed-off-by: Niko Säkkinen <niko.sakkinen@amd.com>
* fix(mha_v4): restore BF16 dense dispatch and launcher byte strides
Re-hook mha_v4() through mha_v4_packed for BF16/NONE, reject sparse BF16
explicitly, and pass byte strides (skipping descale setup) in populate_dense_kernarg.
Signed-off-by: Niko Säkkinen <niko.sakkinen@amd.com>
* style: run black, ruff, and clang-format on block-sparse MHA v4 changes
Signed-off-by: Niko Säkkinen <niko.sakkinen@amd.com>
* docs(mha_v4): trim sparse section implementation detail
Drop kernarg offsets, bit-packing formulas, and duplicate sparse GQA
text from mha_v4.md; keep API contracts and move sparse GQA notes into
Sparse Contract.
Signed-off-by: Niko Säkkinen <niko.sakkinen@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
* Add gfx942 sorted-sparse MHA v4 kernels
The gfx942 FP8/FP8 and INT8/FP8 sparse rows use a 256x64 tile rather
than gfx950's 256x128, so sparse geometry is no longer arch-invariant:
mask shapes go through mha_v4_kv_tile(), and the key-length check reads
cfg.ts_kv instead of a literal 128.
Signed-off-by: Niko Säkkinen <niko.sakkinen@amd.com>
* fix(mha_v4): guard sparse launches by device, validate LUT contents
The sparse launcher installed HipDeviceGuard after build_sorted_work_table,
whose raw HIP kernels take the current device rather than Q's, so a launch
with Q on a non-current GPU faulted; mha_v4_sparse_work_table had no guard
and silently returned zeros. Move the guard above every device query and
launch, and add one to the work-table op.
Also reject non-bool and wrong-device block_mask, bound kv_block_indices
against the row count, and add opt-in AITER_MHA_V4_VALIDATE_LUT=1 for
device-side checks. Empty LUT rows fault in the ASM rather than acting as
no-ops, so document them as invalid. Add tests proving sparse selection
follows the LUT per tile, per head, and across query tiles.
Signed-off-by: Niko Säkkinen <niko.sakkinen@amd.com>
* fix(mha_v4): make an empty sparse LUT row write zeros
lut_count == 0 faulted the sorted-sparse ASM, so the launcher declared empty rows
illegal. Rebuild the ten sparse code objects with the prologue reads clamped and
the row's KV traversal skipped, then follow through on the host: drop the
kLutEmptyRow rejection, and relax the unconditional kv_block_indices bound, which
was derived from the one-block-per-row assumption and would now reject a valid
tightly-packed LUT. The buffer must still be non-empty, since the kernels
dereference the row base even for a row that selects nothing.
Signed-off-by: Niko Säkkinen <niko.sakkinen@amd.com>
* perf(mha_v4): rebuild the gfx942 i8fp8 sparse object without the hot-path clamp
Signed-off-by: Niko Säkkinen <niko.sakkinen@amd.com>
* test(mha_v4): cover a partial query tile with an empty sparse row
Every sparse case used a whole number of 256-row query tiles, so the tail masking
the empty-row no-op is built on was never exercised alongside a short tile. Add
one case at 64/128/200 trailing rows that checks the short tile still reads the KV
blocks its row names and that an all-False row on it returns zero.
Signed-off-by: Niko Säkkinen <niko.sakkinen@amd.com>
* fix(bench_sage): pass the sparse LUT for the MX MHA v4 recipes
mha4_mxfp4/f4f4/mxfp6/f6f4 called mha_v4_packed directly instead of the
launch_mha_v4_packed wrapper that injects the LUT kwargs, so --block-sparsity was
silently ignored and every density measured dense.
Signed-off-by: Niko Säkkinen <niko.sakkinen@amd.com>
* refactor(mha_v4): rebuild the gfx950 f4f4 sparse object with a prologue-only clamp
Signed-off-by: Niko Säkkinen <niko.sakkinen@amd.com>
* mha_v4: take the sparse KV tile from the manifest
mha_v4_kv_tile() restated ts_kv as per-arch constants while the launcher read it
from the manifest row it dispatches on. Read the CSV instead (mode=1 rows), behind
torch_compile_guard since Dynamo traces a cached body and open() broke fullgraph
on the block_mask path. Adds the compile test, and moves the work-table build
measurements into mha_v4.md.
Signed-off-by: Niko Säkkinen <niko.sakkinen@amd.com>
---------
Signed-off-by: jcaraban <jcaraban@amd.com>
Signed-off-by: Niko Säkkinen <niko.sakkinen@amd.com>
Co-authored-by: jcaraban <jcaraban@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
* [HIP] [DCP] Enable fused indexer QK preparation (#5066)
* [DCP] Enable fused indexer QK preparation
* format
* fix uncondition clamp
* modify case
* [ASM] [HIP] [CI] Mxfp6 gemms (#4859)
* deploy mxfp6 gemms
* fix mxfp6 accuracy
* fix mxfp6 source formatting
Remove trailing whitespace so the clean branch passes git diff checks.
Co-authored-by: Cursor <cursoragent@cursor.com>
* Add MXFP6 GEMM tuning and shape-based dispatch
* refactor
* fix
* ruff
* replaces per-element log2/exp2 encoding with mathematically equivalent piecewise E2M3 encoding.
* co-pilot comments fix
* improved hip quantization
* fix MXFP6 backend and buffer validation
* Apply suggestions from code review
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
* fix broken copilot suggestions
* fix A6W6 ASM default kernel selection
* cover all A6W6 kernels and padding paths in CI
---------
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
* [Triton/Gluon] Consolidate and reorganize ops/triton utils (#5061)
* [Triton/Gluon] Add two fused ops for diffusion transformer blocks (#4659)
* [triton] Add two fused ops for diffusion transformer blocks
A DiT block spends its non-GEMM, non-attention time in two patterns that torch
runs as long chains of elementwise ops. Both are memory bound, and both are
dominated by temporaries the maths does not need.
fused_rmsnorm_indexed_adaln
out[m] = rmsnorm(x[m], weight) * (1 + scale[idx[m]]) + shift[idx[m]]
Adaptive layernorm: every token indexes a small table of modulation vectors,
one row per (modality, timestep). Unfused, the normalised activation is written
and immediately re-read, and both table gathers are materialised at [M, N] --
680 MB each at a 63k-token request. One program owns a block of rows and walks
each row in column tiles, once to accumulate the sum of squares and once to
normalise and modulate, so x is read once and out written once.
Two details that matter for this workload. Rows are tiled rather than padded to
the next power of two, because a 5376-wide row would mask off a third of every
access at 8192. And a block of consecutive tokens usually shares one modulation
index -- packed sequences are laid out in runs of one modality -- so the kernel
checks for that and collapses the [BLOCK_M, BLOCK_N] gather to a single
[BLOCK_N] load broadcast in registers.
fused_qk_norm_rope_cached
q[t, h] = rope(rmsnorm(q[t, h], q_weight), cos_sin_cache[t]) (and k)
Per-head RMSNorm followed by partial NeoX RoPE, on q and k, in place. The
existing rope ops do not cover this case: they assume the rotated subspace is
the whole head or half of it, and diffusion transformers rotate fractions in
between (96 of 128 for MiniMax-H3), while the cache-write variants want a paged
KV cache that a diffusion model does not have.
One program owns a token. A token's heads are contiguous, so the [H, D] tile is
one coalesced run and the token's cos/sin row is read once for all heads rather
than being broadcast into a [T, H, D] temporary. Only each token's [H, D] block
must be contiguous, so q and k can be strided views into a packed qkv
projection and are rotated in place, never materialised.
Measured on MI355X, bf16, at MiniMax-H3's shapes:
rmsnorm + indexed adaln, 63232 x 5376 1.708 ms -> 0.329 ms 5.2x
qk norm + rope, 63232 x 56 x 128 11.690 ms -> 1.262 ms 9.3x
Both hold their speedup across the token counts one rank sees at Ulysses 1/2/4/8.
Accuracy: both keep the row in fp32 across the whole fusion, so they are nearer
the fp32 result than the op chain they replace, which rounds to bf16 at each
step. Against that chain on a real 50-layer model, one forward agrees to
cosine 1.0000000 with max relative error 1.3e-4.
58 tests: every table row exercised individually (a kernel that broadcast row 0
would pass a uniform-index test), the uniform and scattered index paths checked
against each other, q and k given different norm weights, the unrotated tail
checked for passthrough, and the strided-qkv-view case checked to leave v
untouched.
* Address review comments on the diffusion adaLN / RoPE fusions
Test fixes:
- test_uniform_and_scattered_indices_agree asserted nothing. Both index
tensors were torch.full((M,), 2), so `fast` and `slow` were the same call on
the same input and assert_close(atol=0) could not fail. That left the
kernel's `uniform = tl.min(idx) == tl.max(idx)` branch -- which broadcasts one
modulation row instead of gathering [BLOCK_M, BLOCK_N] -- with no coverage at
all. The intent was also unreachable as written: with a single index value no
arrangement is ever non-uniform.
Two table entries are now made identical, so the same modulation is reachable
both uniformly (broadcast branch) and alternating (gather branch) and the two
must agree bit for bit. Verified by breaking the uniform branch on purpose
(broadcast table row 0 rather than the block's index): the old assertion still
passed, the new one fails.
- Every call whose output is asserted on now states `eps=1e-5` rather than
leaning on the wrapper default, matching the `reference` calls beside them.
The default is that same value, so nothing was computing the wrong thing, but
the tests should not depend on it staying put. The two `pytest.raises` calls
keep the default, where eps plays no part.
Kernels and wrappers:
- 1.0 / tl.sqrt -> tl.rsqrt in both kernels. Checked rather than assumed: all
58 tests pass unchanged, including the fp32 cases at 2e-6 / 2e-5.
- Lazy %-style logging instead of eagerly built f-strings. This needed
AiterTritonLogger to forward *args -- its methods took (self, msg) only,
which is why the f-string pattern is everywhere in the Triton kernels. The
change is additive, so existing single-argument callers are unaffected.
- The one assert in the adaLN wrapper without an error message now has one.
- Dropped the `if M == 0` / `if T == 0` guards. They were speculative; no
framework hands these ops empty tensors.
- The RoPE wrapper docstring showed only the q equation; k was missing.
* fix(dist): make raw IPC input pools usable — remove init_dist_env's vestigial signal/buffer block, add explicit raw-pool override (#4924)
* fix(dist): init_dist_env no longer breaks raw IPC input pools
Under PYTORCH_HIP_ALLOC_CONF=expandable_segments:True -- the very
configuration the raw_cached input pool exists for (#4174) -- init
failed twice over in init_dist_env's signal/buffer block (#4921):
* register_input_buffer(signal) exports the signal tensor's pointer via
hipIpcGetMemHandle, but the torch.zeros signal is VMM-backed under
expandable segments and the export dies at custom_all_reduce.cu:417
with 'invalid argument';
* ca_comm.buffer = ca_comm._pool["input"].tensor raises by design,
because the raw_cached pool is a plain hipMalloc region with no
backing torch.Tensor.
The block is removed rather than repaired, because all of it was
vestigial:
* ca_comm.signal / ca_comm.buffer are never read anywhere in the tree;
* C++ register_input_buffer only inserts a pointer-translation entry
keyed by the registered tensor's own address, which is consulted only
when an allreduce is invoked with that exact tensor as input --
something that never happens for the signal tensor (open_ipc_handle's
handle cache is filled on demand, so no pre-warming is lost either);
* gfx1250 has skipped the entire block since its VMM bring-up (the
vmm_exchange rendezvous deadlocks) and works without it.
CustomAllreduce.__init__ already builds its own meta/input pools and
forces the copy-in path under expandable segments, so nothing here was
load-bearing.
get_tp_group stays imported: this module is a re-export surface
(downstream engines import set_custom_all_reduce through it).
Adds op_tests/multigpu_tests/test_init_dist_env.py: brings up
init_dist_env per rank under both allocator modes (default torch pool,
and expandable_segments -> raw_cached) and checks one allreduce. The
existing test_custom_allreduce.py performs its own init and never
executes init_dist_env, which is how the regression shipped.
Fixes #4921.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* feat(dist): AITER_CUSTOM_AR_RAW_INPUT_POOL forces the raw IPC input pool
The raw (plain-hipMalloc) input pool previously had exactly one trigger:
PyTorch expandable segments. But expandable segments break custom
allreduce later anyway -- every capture-time output is a fresh
torch.empty_like whose VMM pointer get_output_buffer_RD records for
post-capture IPC export, which then fails in get_graph_buffer_ipc_meta
-- so the raw pool's one trigger leads to an unusable configuration
(#4921, third failure mode; #4621's copy-in guard covers inputs only).
The override gives the raw pool a trigger that works: co-resident
engines on one node, where a second engine's torch.empty input pool can
fail hipIpcGetMemHandle outright. Under the default allocator everything
else (meta pool, capture-time outputs, graph flush) stays exportable, so
only the input pool needs to move to hipMalloc.
Extends test_init_dist_env.py with a raw_override mode that asserts the
flag actually selects the raw pool and allreduce stays correct.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* feat(dist): log the input-pool allocation mode at init
A silently-inert pool trigger is indistinguishable from a working one by
behaviour alone -- the engine serves fine single-engine either way, and
the failure modes this pool exists to avoid (#4921) only appear in
specific modes under specific co-residency. One INFO line per rank makes
every run self-document which pool it actually got, so a mislabeled
measurement is catchable from the log rather than by re-deriving the
allocator state.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* fix(dist): honor capture registration setting in fused AR
---------
Co-authored-by: ThomasNing <thomas.ning@amd.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
* [FlyDSL] [opt][rope] optimize qk norm rope Ep decoding case specially for T512 (#5070)
* perf: TDM prefill bandwidth opt — K=5 occupancy + position prefetch
Two changes to improve TDM prefill kernel bandwidth at small T:
1. Reduce LDS buffer count from K=6 (192KB) to K=5 (160KB) when
num_rows <= 65536. On gfx1250 with 320KB LDS per CU this allows
2 WGs/CU instead of 1, doubling occupancy.
2. Prefetch position buffer_load before the hot loop: issue the first
group's position load before TDM prologue, and each subsequent
group's position load after the prior group's last tile compute.
This overlaps the position→cos/sin serial dependency chain with
TDM tile transfers and compute, reducing loadcnt stalls by ~41%.
ATT trace confirms total stall cycles drop 36% (72K → 46K), with
loadcnt (HBM) stalls down 41% and dscnt (LDS) stalls down 70%.
Measured kernel times (gfx1250, H=128 D=512 RD=64 BF16):
T=512: 16.5us → 11.9us (+39%) 8.2 → 11.4 TB/s
T=16384: 302us → 274us (+10%) 14.3 → 15.8 TB/s
Co-Authored-By: Claude <noreply@anthropic.com>
* perf: drop TDM prefill rotation to K=4 at small T
At num_rows <= 65536 (T=512, H=128) CT=8 yields gx_q=256 workgroups for
256 CUs -- exactly one WG per CU. LDS is therefore never the limiter at
this shape, which invalidates the reasoning behind the previous K=5
choice (it was picked to keep the arena at 160 KB so two WGs would fit,
but a second WG never exists here). With K free to pick on latency
alone, K=4 measures faster.
T=512 H=128 D=512 RD=64 BF16, gfx1250, three runs each:
K=5: 15.728 15.647 15.683 -> 15.69 us (8636 GB/s)
K=4: 15.357 15.360 15.220 -> 15.31 us (8874 GB/s)
Non-overlapping ranges, ~2.4% faster. T=16384 is unaffected (it takes
the num_rows > 131072 branch at K=6): 302.6 / 307.5 us, unchanged.
The mechanism behind the shallower rotation winning is not understood --
it is not LDS or occupancy driven -- so the docstring records the
measurement and warns against extrapolating to other shapes.
Also measured and rejected on this shape:
- CT=4 to reach 2 WG/CU: 14.97 vs 14.86 us, no gain. Doubling the wave
count doubles the per-wave cold-start cost, cancelling the extra
latency hiding.
- Issuing the position load before the TDM prologue: 15.63 vs 15.68 us,
within noise. The K descriptor setups are far too few instructions to
cover a ~1700-cycle DRAM miss.
- TDM store (LDS -> global) in place of buffer_store, tried with a
reused input buffer, one dedicated output buffer, and two rotating
output buffers: 16.04 vs 15.27 us at matched K=4, ~5% slower. The
LDS round trip (ds_write plus tensorcnt sync) costs more than the
s_wait_xcnt it removes.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* docs: erratum for 53d67009b — its perf claims and attribution were wrong
53d67009b ("perf: TDM prefill bandwidth opt — K=5 occupancy + position
prefetch") is already published, so its message is left in place and
corrected here instead. Four claims in it are wrong. The code it shipped
is fine and is kept; only the reasoning and the numbers were bad.
1. "Prefetch position buffer_load before the hot loop ... reducing
loadcnt stalls by ~41%"
The prefetch is a no-op. issue_pos() is followed immediately by the
trunci that consumes it, in the same statement, so there is no
distance for the load to cover its miss:
pending_pos[0] = issue_pos(tok_of(tile_base + i + 1))
cs_cache[0], cs_cache[1] = _cs_from_pos(
fx.Int32(pending_pos[0].trunci(i32)))
An ATT capture of the shipped code shows group 1's position still
stalling 5249 cycles despite being "prefetched". The loadcnt
reduction came entirely from K=6 -> K=5.
2. "K=5 (160KB) allows 2 WGs/CU instead of 1, doubling occupancy"
At num_rows=65536 (T=512, H=128), CT=8 gives gx_q=256 workgroups for
256 CUs, so a second WG per CU never exists and LDS was never the
limiter. 034220f3a already replaced this reasoning in the
_tdm_tiles_per_wg docstring.
3. "T=512: 16.5us -> 11.9us (+39%)"
The 11.9us came from a hand-rolled L2-warm timing loop and is not
comparable to the 16.5us op_test figure it was subtracted from.
Measured on one path (op_tests/test_flydsl_qk_norm_rope_quant.py),
idle GPU, five runs each:
K=6 16.333 16.434 16.386 16.463 16.441 -> 16.41 us
K=5 15.674 15.688 15.618 15.668 15.639 -> 15.66 us (+4.6%)
K=4 15.332 15.227 15.281 15.239 15.422 -> 15.30 us (+2.3%)
So 53d67009b was worth +4.6%, not +39%, and the two commits together
are worth +6.8% (16.41 -> 15.30 us).
4. "T=16384: 302us -> 274us (+10%)"
53d67009b does not touch that path. For num_rows > 131072,
_tdm_tiles_per_wg returns (TILES_PER_WG, NUM_BUFFERS) = (40, 6),
identical to the pre-commit default of CT=40 with the builder's
num_buffers=NUM_BUFFERS. T=16384 measures ~305 us both before and
after; the reported gain is spurious.
Root cause of 3 and 4: numbers from two different timing harnesses were
compared against each other. Only same-harness, same-session, repeated
measurements are used above.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* perf+fix: TDM prefill — 8-wave workgroups, and tighten the drain-phase wait
Two changes, found together while investigating why smaller workgroups
appeared to help.
1. Correctness: the drain phase under-waited on its TDM loads.
Tile i consumes TDM load #i; loads are issued in tile order, K in the
prologue then one per iteration while i + K < CT. In steady state K+i
are outstanding, so tensor_wait(K-1) leaves exactly #0..#i retired --
correct. Once the issues stop, the issued count freezes at CT and
K-1 is too loose: #i is only guaranteed retired with at most CT-1-i
left, which has to reach 0 on the last tile. The wait is now
min(K-1, CT-1-i); both operands are compile-time constants in the
unrolled loop, so this costs nothing.
This was latent, not new. With ROWS_PER_TILE=32 the per-tile compute
happened to outlast the load, so the shipped kernel got away with it.
Shrinking the tiles exposed it: at RT=8/CT=16 the output was wrong in
exactly the last three tiles of every workgroup (tile%CT histogram
[0]*13 + [53,53,49]), with all 512 columns of those rows wrong --
i.e. the LDS input itself, not the RoPE tail. err_q 0.027 -> 5.96e-08
with the fix, same config, same build.
Measured cost at the shipped shape (T=512 H=128, three runs each):
without: 15.389 15.254 15.281
with: 15.308 15.393 15.304
2. Perf: ROWS_PER_TILE 32 -> 8, and CT 8 -> 16 for num_rows <= 65536.
gx_q = num_rows / (ROWS_PER_TILE * CT) has to stay at or above the 256
CUs. At the low end of the TDM range it did not: num_rows=32768
(T=256, H=128) gave gx_q=128, so half the CUs sat idle. RT=8 restores
full coverage there and doubles it at num_rows=65536.
T=512 H=128, five runs each:
RT=32: 15.295 15.233 15.270 15.344 15.305 -> 15.29 us
RT=8: 15.090 14.937 14.929 15.009 14.936 -> 14.98 us (-2.0%)
T=256 H=128, four runs each:
RT=32: 12.610 12.587 12.581 12.645 -> 12.61 us
RT=8: 10.904 10.499 10.587 11.065 -> 10.76 us (-14.6%)
Across the TDM path (H=128 unless noted):
T=256 -16.3% T=512 -2.5%
T=1024 -6.5% T=16384 -2.2% T=16384 H=16 -0.9%
Shapes below TDM_MIN_ROWS=32768 take the r32_w32 path and are
untouched by ROWS_PER_TILE; the +-1-3% seen on those in a sweep is
run-to-run noise.
At RT=8, GROUP = H/RT = 16 and TILES_PER_WG=40 is not a multiple of
it, so cos/sin hoisting turns off for the largest shapes. That is not
a regression -- T=16384 still improves -2.2% -- consistent with the
separate finding that the position->cos/sin chain is worth ~2.6% of
wall clock despite being 36% of stall cycles.
Validated on 20 (T, H, q_weight) combinations plus the SWA direct/paged
and decode paths: all err_q/err_kv <= 1.3e-06, 16/16 checkAllclose pass.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* perf(flydsl): optimize qk norm rope decode
* perf: fuse FP8 quant into TDM prefill and tune gfx1250 occupancy
Keep 2 WGs/WGP on the T=512 path, use 16-row tiles only for short prefill, and stream FP8 (grouped/e8m0) through the TDM kernel so Q write traffic drops without falling back to the slower direct path.
Co-authored-by: Cursor <cursoragent@cursor.com>
* perf(flydsl): TDM reads KV strided, and drops to K=2 on deep grids
Two independent changes to the gfx1250 TDM path.
1. Read KV with a row stride.
The TDM kernel indexed KV as `tok * D`, so the wrapper had to force
kv.contiguous(). The V4 call site slices KV out of a wider qkv_a tensor,
so that fired a full elementwise copy kernel on every invocation.
get_trace_perf sums all device kernels, so the copy landed inside the
number the op-test reports: 3.52us on top of a 12.20us kernel at T=512,
22% of the reported total, for nothing -- fused-kernel time is identical
whether KV arrives strided or contiguous. Thread kv_in_row_stride through,
matching what the wave32 and wave64 paths already do.
2. TDM buffer depth K=6 -> K=2 from num_rows >= 131072.
K sets the length of the load-only prologue. Once the grid is deep enough
that one workgroup's prologue overlaps another's steady state, the shallow
K=2 wins; below that a workgroup must cover its own load latency and the
deeper prologue pays for itself. Measured on top of 46ee44bc6, public API,
rotate=4, interleaved medians:
num_rows 32768 (T=256) K=2 +31.9% 131072 (T=1024) K=2 -3.0%
49152 (T=384) K=2 +9.8% 262144 (T=2048) K=2 -5.1%
65536 (T=512) K=2 +4.1% 1048576 (T=8192) K=2 -4.3%
2097152 (T=16384) K=2 -2.9%
Note the crossover sits above T=512: with 46ee44bc6's occupancy tuning in
the base, K=6 is now correct for the decode shape, so the T=512 gain here
comes from (1) alone.
T=512 on this gfx1250, same harness and rotation as the pre-change baseline:
qw off 15.72us -> 11.07us
qw on 16.48us -> 12.16us
The runperf script's own sweep reports 10.74us / 10.83us for the same two
rows; it launches through a tighter loop, so treat the pair above as the
comparable figure.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* test: runperf script sets PYTHONPATH and prints a combined summary table
`python op_tests/foo.py` puts op_tests/ on sys.path rather than the repo
root, so `import aiter` failed unless the shell already exported
PYTHONPATH. Set it from the script's own directory.
Also tee both sweeps to a log and replay every markdown table at the end
under its original heading, so the T=16384 and T=512 runs can be compared
without scrolling back through two sweeps of output.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* test(flydsl): make the qk_norm_rope %peak column arch-aware
_PEAK_BW_GBPS was a single 22000.0 labelled "MI355X HBM3e peak", but
22 TB/s is the gfx1250 figure -- MI355X (gfx950) is 8 TB/s and MI300X
(gfx942) is 5.3. The column was therefore only meaningful on gfx1250,
and silently wrong on the other two archs the file already lists in
SUPPORTED_GFX.
Look the peak up per arch instead. Unknown archs report None rather than
a fabricated percentage.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* style: apply black to qk_norm_rope_quant.py
CI runs black[colorama]==26.5.1 and this file was the only one in the
branch it wanted to reformat. Formatting only -- verified the AST is
identical before and after.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* chore: drop the local perf harness from the repo root
runperf-qknormrope-bs16-t16384.sh is a personal benchmark driver for one
shape on one machine, not something the repo should carry at its root.
It is kept locally alongside the other measurement tooling.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* perf(flydsl): let the TDM path take the fused SWA write
`use_tdm` excluded kv_write and paged outright, so any caller that passes
swa_kv fell back to the wave32 kernel. That is what the model does, so the
decode shapes were running qk_norm_rope_H128_D512_RD64_kvw_r32_w32_flydsl
and none of the TDM tuning reached them. The op-test did not show this:
its headline rows pass no swa_kv, and its SWA sweep is pinned to T=8..96
by the paged fixture's capacity, so it never reaches a TDM-eligible size.
Port the scatter into the TDM kernel's KV path. The gates are copied from
the wave32 sibling unchanged -- bid<0, pos<0, paged blk past the table,
table entry -1, resolved row past the pool -- and the row index is widened
to 64 bits before the byte multiply, as there too.
gfx1250, public API, rotate=4, interleaved medians:
wave32 TDM gain
direct T=512 17.65us 11.61us -34.2%
T=1024 32.61us 23.59us -27.7%
paged T=512 17.46us 11.39us -34.8%
T=1024 32.07us 23.72us -26.0%
Verified against the wave32 path at T=512/1024 in both modes: the pool is
byte-identical to kv_out at every resolved row, rows nobody targets stay
zero, and guard rows either side of the pool are untouched. Each skip gate
is covered by its own case.
Not covered: pos<0. It is not a legal input -- the main path indexes
cos/sin with the raw position long before the scatter -- and the wave32
path faults on it identically, so this is not a new exposure.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* fix(flydsl): address q_out past 4 GiB in the TDM path
A buffer descriptor's num_records is 32-bit, so one descriptor reaches
4 GiB. q_out crosses that at T*H*D*2 >= 4 GiB -- T>=32768 at H=128,
D=512 -- and every row past the limit was dropped or wrapped.
The failure started exactly on the boundary: at T=32768 the first bad row
was 4194303, whose last byte sits at 0xFFFFFFFF, one past num_records, so
precisely 4 elements were lost. Beyond 4 GiB it degraded fast -- 0.05% of
q_out wrong at T=32776, 38.9% at T=40960, NaN in both.
Bias the descriptor base per workgroup instead, the same trick the SWA
scatter in this file already uses. A workgroup owns CT*RT rows, so the
32-bit offset then spans 128 KB rather than the whole tensor. The bias is
computed once per workgroup, outside the tile loop.
This predates the TDM work: the wave32 path fails identically at these
sizes, it is simply unreachable there now that TDM covers num_rows >=
32768. The op-test's default sweep includes T=65540 and had been failing
on it.
T=65540, H=128: err 1.4e-07 (was garbage), 12524 GB/s
T=40960: 0 bad elements (was 1.04e9)
No measurable cost -- T=512 11.40us, T=2048 40.66us, T=16384 330.25us,
all within run-to-run spread of the previous commit.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* style: collapse the SWA store guard (ruff SIM102)
The `do_swa is not None` guard was redundant -- None is already falsy, so
the two ifs fold into the one the wave32 path next door already uses:
if const_expr(kv_write) and do_swa:
Short-circuiting still keeps the const_expr and the runtime predicate
apart: kv_write=False never evaluates do_swa, emit_q passes None so no
store is traced, and emit_kv passes the predicate so scf.if is emitted as
before. Re-ran the SWA scatter checks (10/10) and the op-test (465 passed)
to confirm codegen did not shift.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* perf(flydsl): halve the TDM workgroup when the grid is too thin
RT sets the workgroup size and therefore how many of them the grid holds:
gx_q = num_rows / (RT * CT). Below 2 workgroups per CU there is no
neighbour whose steady state can cover a workgroup's own load latency,
and RT=8 sits under that line for num_rows < 65536 -- at num_rows=32768
it yields exactly one workgroup per CU.
Halving RT there doubles the grid and pays for itself. It stops paying at
65536 (exactly 2 WGs/CU, a wash) and turns negative past it, where the
smaller workgroup costs more than the extra parallelism returns (+2.4% at
num_rows=262144), so RT=8 holds from 65536 up.
This is the same question K already answers, one level up: can a
workgroup's latency be hidden by a neighbour, or must it cover its own?
gfx1250, public API, rotate=4, interleaved medians, before -> after:
decode (fused SWA) prefill (no SWA)
T=256 6.82 -> 6.65 T=256 6.85 -> 6.66
T=384 9.55 -> 8.53 T=384 9.30 -> 8.76
T=512 11.47 -> 11.45 T=512 11.16 -> 11.16
T=2048 41.34 -> 41.27 T=2048 41.23 -> 41.19
T=16384 333.53 -> 334.22 T=16384 331.43 -> 334.31
T>=512 is untouched by construction -- the geometry it selects is
unchanged, so those rows are noise. Prefill only reaches RT=4 on prompts
shorter than 512 tokens, where it is also a win.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* style(flydsl): trim the comments this branch added
The tuning rationale had grown into paragraphs sitting on top of two-line
functions. The measurements behind each threshold are in the commits that
introduced them, so the source only needs to say what the knob does.
Also folds emit_kv's inlined position load back into a load_pos() helper
that load_cs() now shares.
Net -26 lines. No behaviour change: op-test 465 passed, SWA scatter checks
10/10, T=512 decode 11.43us.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---------
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Satya Nikhil Kodukula <nikhil.kodukula@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
* [Tune] Add GLM-5.3 BF16 GEMM configs for gfx950 (#5060)
* [Triton/Gluon] [HIP] Dev lumen (#4978)
* Add lumen triton kernels and custom ops (clean cherry-pick)
Cherry-pick of 4e19b8e3e (ZhangDanyang-AMD) onto upstream/main.
Only new files preserved; upstream existing code left untouched.
Adds: triton quant kernels, FP8/MXFP8 attention, MoE GEMM variants,
cross_entropy, fused_norm_quant_gemm, AOT precompiled kernels,
moe_sorting test cases.
Registers cross_entropy and mxfp8_attention in triton __init__.py.
* add large-M/small-N RMSNorm backward specialization
* add gfx942 (MI308X, 80CU) blockscale bpreshuffle GEMM configs
* add requant_fp8_row_to_col, chunked cross-entropy, add gfx942 per-shape GEMM configs
compile_ops type-check fix omitted — upstream already has _is_tensor_like fix.
* add gfx942 preshuffle GEMM configs for llama2-7b/13b/70b and qwen3-8b
* add MoE weight gradient Triton kernel (moe_wgrad)
Adds a fused Triton kernel for MoE weight gradients that operates
directly on sorted_token_ids/expert_ids from moe_align_block_size,
eliminating the need for sort+pad+bmm and CPU-GPU sync in backward.
Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
* add is_cdna4() arch probe for gfx950 family
Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
* add DSV4 sparse MLA training and indexer ops for DeepSeek-V4-Flash
- Sparse MLA: fused Triton fwd/bwd kernels with CSR-based dKV gather (no atomics)
- Indexer: BLAS-based scoring via torch.einsum (hipBLASLt) + PyTorch autograd
- Correctness tests: 84 sparse MLA tests + 48 indexer tests, all passing
Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
* add Triton MHC forward and backward support
Co-authored-by: Cursor <cursoragent@cursor.com>
* integrate SonicMoE pure-Triton grouped GEMM MoE with full autograd
Port SonicMoE's pure-Triton MoE implementation from sonic-moe into aiter-lumen.
Provides trainable MoE layer with fused router + grouped GEMM + activation,
supporting forward and backward passes for all 7 activation types.
New files:
- _triton_kernels/moe/sonicmoe/: 9 kernel modules (grouped GEMM, activations,
routing metadata, reduction, forward/backward autograd functions)
- aiter/ops/triton/sonicmoe.py: public API wrapper
- configs/moe/gfx942-MOE-SONICMOE-BF16.json: autotune configs for MI308X
- op_tests/test_sonicmoe.py: correctness + benchmark tests
Correctness verified on MI308X (T=64, H=128, I=64, E=4, K=2, BF16):
| Activation | output rel err | dx rel err | dw1 rel err | dw2 rel err | Status |
|------------|---------------|------------|-------------|-------------|--------|
| swiglu | 0.0097 | 0.0132 | 0.0138 | 0.0104 | PASS |
| geglu | 0.0014 | 0.0089 | 0.0100 | 0.0000 | PASS |
| reglu | 0.0014 | 0.0103 | 0.0098 | 0.0000 | PASS |
| gelu | 0.0014 | 0.0134 | 0.0140 | 0.0000 | PASS |
| relu | 0.0014 | 0.0155 | 0.0168 | 0.0000 | PASS |
| silu | 0.0014 | 0.0146 | 0.0150 | 0.0000 | PASS |
| relu_sq | 0.0014 | 0.0104 | 0.0117 | 0.0000 | PASS |
All relative errors < 2%, well within BF16 tolerance.
Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
* fix tests: call existing topk_softmax and drop redundant RMSNorm 65536x128
The cherry-picked pytest imported a non-existent softmax_topk API; retarget it at ASM topk_softmax. 65536x128 duplicated 16384/364800 coverage of the large-M/small-N bwd path.
Co-authored-by: Cursor <cursoragent@cursor.com>
* fix style: format PR Python with Black and satisfy Ruff 0.16
Unblocks Checks so check-signal can let HIP/Triton CI run. Also add missing torch/_get_activation_from_str imports in gemm_a16w16_agnostic.
Co-authored-by: Cursor <cursoragent@cursor.com>
* drop files already removed on origin/main instead of resurrecting them
Cherry-picks had re-added pre-ctypes pybind/headers, AOT hsaco, and a
redundant bpreshuffle tuner. Keep gfx942 rows in the existing CSV.
Co-authored-by: Cursor <cursoragent@cursor.com>
* move gfx942 GEMM tunes into nested config layout so they actually load
Place llama2-7b/13b/70b, llama3-8b qkv, and qwen3-8b N/K tables next to
each family's DEFAULT.json. Legacy configs/gemm/ paths are ignored once
the nested default exists.
Co-authored-by: Cursor <cursoragent@cursor.com>
* load gfx942 SonicMoE JSON at launch instead of autotuning those kernels
Pick N/K/E and H buckets from {arch}-MOE-SONICMOE-BF16.json so production shapes skip the autotune search; fall back to the old autotune lists when the file is missing.
Co-authored-by: Cursor <cursoragent@cursor.com>
* fix ruff C408 in SonicMoE launch kwargs
Rewrite dict() calls as literals so Checks reviewdog stops failing the PR.
Co-authored-by: Cursor <cursoragent@cursor.com>
* format PR mxfp8/moe GEMM modules for Black 26
Remove extra blank lines after module docstrings so psf/black@stable in Checks passes on CI.
Co-authored-by: Cursor <cursoragent@cursor.com>
* drop gfx942 CK GEMM row that duplicates DSV4 opus tune
Merge keys are gfx/cu_num/M/N/K, so ck vs opus for 2048x4096x1024 on 80 CU fails wheel prebuild. Keep the faster opus entry from the DSV4 table.
Co-authored-by: Cursor <cursoragent@cursor.com>
* load SonicMoE JSON from nested gfx942/triton/moe layout
Co-authored-by: Cursor <cursoragent@cursor.com>
* format sonicmoe_config_utils for Black 26
Co-authored-by: Cursor <cursoragent@cursor.com>
---------
Co-authored-by: ZhangDanyang-AMD <danyzhan@amd.com>
Co-authored-by: leiwu0812 <leiwu0812@users.noreply.github.com>
Co-authored-by: Claude Opus 4 <noreply@anthropic.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
* [CI] Avoid direct github.event interpolation in run: blocks (SEC-00830) (#5109)
Mythos scan finding SEC-00830 (ROCM-26711) flags GitHub Actions event
context interpolated straight into `run:` shell blocks, where the value is
pasted into the script text before the shell parses it.
aiter-test.yaml already uses the `env:`-indirection pattern in most steps
(15 `env:` blocks; `${GITHUB_EVENT_NAME}` at lines 50/537). This brings the
five remaining spots in line:
- 3x `if [ "${{ github.event_name }}" = "schedule" ]`
-> `${GITHUB_EVENT_NAME}` (GitHub's built-in, same as lines 50/537)
- 2x `BASE_REF="${{ github.event.pull_request.base.ref || github.ref_name }}"`
-> hoisted into a step-level `env:` block
After this change no `${{ github.event* }}` remains inside any `run:` block.
Note this is hardening, not a fix for an exploitable bug. The scanner's stated
attack surface (`github.event.pull_request.title`) does not appear in any
`run:` block. Of the five occurrences, three were `github.event_name` (an
enumerated value) and two were `base.ref` — the PR's *target* branch, which
this workflow constrains to `main` via `branches: [main]` and which an external
contributor cannot name. The point is to keep the pattern out of the file so a
future edit cannot turn it into a real injection.
actionlint: clean before and after.
Refs: ROCM-26711 / SEC-00830
* [CI] Drop registry credentials after jobs on persistent runners (SEC-00837) (#5110)
Mythos scan finding SEC-00837 (ROCM-26712): self-hosted runners are
non-ephemeral, so `docker login` credentials written by one job stay in
~/.docker/config.json and are readable by whatever runs next on that machine.
aiter-test.yaml has three `Docker login` steps and no `docker logout` anywhere:
build_aiter_wheels runs-on: build-only-aiter (no cleanup step at all)
standard runs-on: ${{ matrix.runner }} (has "Cleanup container")
multi-gpu runs-on: ${{ matrix.runner }} (has "Cleanup container")
This adds `docker logout` to the two existing `Cleanup container` steps and
gives build_aiter_wheels the cleanup step it was missing. All three run under
`if: always()`.
This is the immediate mitigation the ticket calls for, not the fix. It narrows
the window but does not close it: credentials still exist on disk between login
and logout, and a cancelled job may skip cleanup entirely. The actual fix is to
register the runners with `--ephemeral` (or `ephemeral: true` under
actions-runner-controller) so every job starts from a clean machine. That lives
in the runner infrastructure, not in this repository.
Existing partial mitigation, unchanged by this PR: all three `Docker login`
steps are already gated on `!github.event.pull_request.head.repo.fork`, so fork
PRs never write credentials in the first place.
actionlint: clean.
Refs: ROCM-26712 / SEC-00837
* [HIP] [ROCm][Perf] Add head_dim 512 + weightless V-norm to fused_qk_norm_rope_cache_pts_quant_shuffle (#5027)
* [ROCm][Perf] Add head_dim 512 + weightless V-norm to fused_qk_norm_rope_cache_pts_quant_shuffle
Enable the fused QK-norm + RoPE + KV-cache op for Gemma4, whose full-attention
layers use head_dim 512 and whose every attention layer applies a weightless
v_norm (RMSNorm with has_weight=false).
- rope_common.h: add warp_rms_norm_no_weight_ (RMS normalize a head with no
learned gamma); apply it to V in fused_mrope_rms_kv_kernel when the new
runtime flag v_norm is set, before the KV-cache write. Add case 512 to the
fused_rope_rms_set_kv head_size switch (VEC_SIZE=16 at 512) and relax the
head_size guard. The mrope-3D launcher is unchanged (passes v_norm=false).
- fused_qk_norm_rope_cache_quant.cu / .h / rocm_ops.hpp: thread the trailing
bool v_norm (default false) through the pts entrypoint and pybind.
- ops/fused_qk_norm_rope_cache_quant.py: add v_norm to the python wrapper.
Validated with a standalone call at head_dim 256 and 512: the 512 template
instantiates/compiles, and the V-cache matches a weightless-norm reference at
bf16 rounding tolerance for both widths.
Co-authored-by: Claude <noreply@anthropic.com>
* [ROCm][Perf][Test] Cover weightless V-norm + head_dim 512 in pts fused op
Add test_fused_qk_norm_rope_cache_pts_v_norm: exercise
fused_qk_norm_rope_cache_pts_quant_shuffle with v_norm on/off at head_dim 256
(Gemma4 sliding) and 512 (Gemma4 full). Asserts the V written to the cache is
weightless RMS-normalized when v_norm=True and raw otherwise.
Co-authored-by: Claude <noreply@anthropic.com>
* [ROCm][Perf][Test] Address review: guard shuffle K write + real v_scale coverage
Fix two issues from PR review:
- Guard the shuffle-layout K write against silent cache corruption. That
path does a single contiguous vec_t store of VEC_SIZE = head_size /
WARP_SIZE elements and get_shuffle_layout_k_base() assumes they all land
in one x-wide chunk (VEC_SIZE <= x). At head_size=512 / WARP_SIZE=32 that
is VEC_SIZE=16, which exceeds x=8 for a bf16/fp16 cache and would corrupt
K for block_size>1. Reject that config with an AITER_CHECK. An fp8 cache
(x=16) still satisfies the bound, so shuffle layout at head_dim 512 with
fp8 KV is unaffected.
- Rework the v_norm op test to follow the file convention and add real
scale coverage. The per-tensor v_scale only divides V on the fp8 quant
write path -- a same-dtype cache copies V verbatim, so the previous
bf16-cache test never exercised the scale. The test now uses @benchmark,
is wired into __main__ with a markdown summary table, and sweeps head_dim
256/512, v_norm on/off, and (bf16, fp8@1.0, fp8@0.5) cache/scale pairs to
check norm-then-quantize ordering.
Co-authored-by: Claude <noreply@anthropic.com>
* Potential fix for pull request finding
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
---------
Co-authored-by: root <root@quanta-ccs-aus-k09-19.adc.amd.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
* [HIP] fix(topk): add acquire fence for mb radix barrier last block (#4841)
* fix(topk): add acquire fence for mb radix barrier last block
`radix_kernel_persistent` uses a per-row cross-block barrier in the multi-block radix top-k path. The waiting blocks observe `pass_done` with an acquire load, which also invalidates their cache state before they reload the global histogram for the next pass. The elected "last" block, however, only publishes `pass_done` with a release store and then falls through to the same plain histogram reload without ever doing an acquire/invalidate.
On MI355X this can let the elected block reload stale histogram lines from the persistent workspace. If that block computes a different `local_len` / `local_k` from its peer blocks, it can leave the pass loop early while another block continues into the next barrier. The early-exiting block can then be elected in the self-reset epilogue and zero `pass_done` while its peer is still spinning on it, producing a permanent GPU wedge in the GLM-5.2 DSA indexer path.
Add a `__threadfence()` plus CTA sync after the release store in the elected last-block branch. This gives the last block the missing acquire-equivalent ordering before it reloads the histogram, making both sides of the barrier observe consistent global memory before computing the next pass state.
This fixes the production hang seen with GLM-5.2 TP4 + DP attention + LMCache + atomesh `dp_sticky`, where one DP rank could wedge inside `aiter::mb::radix_kernel_persistent` and then stall the whole service through DP-attention collectives while `/health` stayed green.
Validation:
- unpatched stock kernel soak reproduced 5 wedges in 213,400 launches / 27.3M row-launches.
- fixed stock kernel soak completed 1,914,200 launches / 245.0M row-launches with 0 wedges, 8.97x the baseline exposure.
- detector build changed the failure signature from `passes=[2,2,1,0]` with `STUCK` to 0 `STUCK` events over 131.4M row-launches.
- end-to-end GLM-5.2 TP4 + DPA + LMCache + atomesh aiperf run has passed warmup and entered 3600s profiling without the previous hang signature.
Signed-off-by: Phi-C <chenxjhit@163.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
* perf(topk): use acquire-only mb radix barrier fence
The elected last block needs device-scope cache invalidation before reloading the global histogram, but does not need release/writeback semantics. Use an agent-scope acquire fence to preserve correctness while avoiding the unnecessary release overhead of threadfence.
Signed-off-by: Phi-C <chenxjhit@163.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
* fix(topk): drain mb histogram atomics before barrier
Ensure every wave completes no-return histogram atomics before block arrival, then establish agent-scope visibility after relaxed polling to prevent cross-block divergence.
Co-authored-by: Cursor <cursoragent@cursor.com>
---------
Signed-off-by: Phi-C <chenxjhit@163.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
* [aiter_opus_plus] detorch (#4958)
* [FlyDSL] 1250 clean moe aux kernel codes and ir, add ut (#5112)
* [CI] Mirror PR title component tags as auto-managed labels (#5057)
* [CI] Mirror PR title component tags as auto-managed labels
* address comments
* Potential fix for pull request finding
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
---------
Co-authored-by: Xin Huang <Xin.Huang@amd.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
* Tune the new Kimi-K3 a8w8 bpreshuffle and bf16 GEMM shapes (#5124)
Adds five a8w8 bpreshuffle (N,K) groups -- 1536x1536, 2048x512, 3584x7168,
7168x1024, 7168x1792 -- and three bf16 groups -- 896x7168, 7168x35840,
20480x7168. Each group covers M as every power of two from 1 to 32768 plus
M=48, so a 48-row request resolves to its own entry instead of padding up
to M=64. None of these had a gfx950 row before: the bf16 file carries rows
for two of the three groups, but only for gfx1250.
Existing rows are untouched. The tuners ran without --all, so only the
newly added shapes were considered, and a key-wise comparison against the
pre-tune files confirms zero modified and zero removed rows.
a8w8, --libtype all -k --shape_grouped on gfx950 (cu_num=256): 85 rows,
78 FlyDSL / 4 CK / 3 CK-tile. Seventeen land on the FlyDSL 8wave pipeline
and nine use splitK, all of them on 3584x7168 where K is large enough for
the extra parallelism to pay off at small M.
Measured against today's behaviour (no tuned row, default kernel), three
runs, per-shape median of 100 iterations after 20 warmup, one GPU, with
the CK and CK-tile lookup tables rebuilt from the new config first:
N=1536 K=1536 322.2us -> 248.2us +22.99%
N=2048 K=512 241.9us -> 175.7us +27.36%
N=3584 K=7168 1797.2us -> 1373.5us +23.58%
N=7168 K=1024 876.3us -> 631.4us +27.94%
N=7168 K=1792 1192.0us -> 910.8us +23.59%
total 4429.6us -> 3339.6us +24.61%
Best single shape is M=2 N=7168 K=1792 at 16.28 -> 6.24us. Two shapes on
1536x1536 first looked like small regressions; a seven-run recheck put
both sides within 0.5% of each other, which is inside the noise for a
4.5us kernel, so they are kept.
bf16, csrc/gemm_a16w16/gemm_a16w16_tune.py without --with-hipblaslt, run
under --compare --update_improved so a row is only written when it beats
the default kernel by at least 3%. Twenty of the 51 candidate shapes
cleared that bar; the other 31 are already at what the default dispatch
picks and are left out. The largest win is M=1 N=896 K=7168 at
14.69 -> 6.64us (54.8%).
op_tests/tuning_tests/test_config_shape_collision.py and
test_csv_validation.py pass (30 tests, 37 subtests).
* Fix PR title tag workflow syntax (#5134)
* [CI] Document and automate the AITER release plan (#4424)
* Document and automate release plan
* Harden AITER release automation
* Update release notes after asset upload
* fix: harden release automation checks
* Fix manual release Docker login
* Adjust release cadence anchor
* Fix reusable release Docker login
Signed-off-by: Xin Huang <Xin.Huang@amd.com>
---------
Signed-off-by: Xin Huang <Xin.Huang@amd.com>
* [Triton/Gluon] combine routing early exit (#5053)
* [Triton/Gluon] [gfx950] gated_delta_rule: drop removed tl.make_block_ptr (#4950)
* [Triton] fix(gated_delta_rule): replace removed tl.make_block_ptr for Triton 3.8
Triton 3.8 removed block pointers. tl.make_block_ptr still exists as a symbol
but raises at trace time:
NotImplementedError: Block pointers have been removed in favor of the
tensor descriptor API
so every gated_delta_rule kernel using it fails to compile. This is an API
removal, not a GPU issue - it reproduces identically on gfx950 and gfx942, and
is what makes op_tests/test_gdn_prepare.py fail on both MI35X and MI300X.
Convert all 128 block accesses to plain pointer arithmetic with explicit bounds
masks, reproducing the previous boundary_check=(0, 1) semantics:
prefill/chunk_o.py 42 sites (6 kernels)
prefill/fused_solve_tril_recompute.py 41
utils/solve_tril.py 35
prefill/fused_cumsum_kkt.py 10
utils/cumsum.py 2
The 2-D helper in chunk_o.py takes both strides so the transposed (K, T) views
with stride (1, H * K) convert without a special case. Stores keep their
fp_downcast_rounding="rtne" behaviour.
Validation on gfx950 with triton 3.8.0+amd.rocm7.1.0.gitf6a045ff:
op_tests/test_gdn_prepare.py 28 rows, max |err| = 0.0, all shapes / all
three hidden backends (triton/flydsl/hip)
* fix(gated_delta_rule): convert remaining l2norm/wy_representation block ptrs
l2norm.py (4 sites) and wy_representation.py (11) still used tl.make_block_ptr,
which Triton 3.8 removed. Both are on live e2e inference paths that
op_tests/test_gdn_prepare.py does not exercise:
gated_delta_rule.py: l2norm_fwd(q/k) when use_qk_l2norm_in_kernel=True
prefill/chunk.py:109: recompute_w_u_fwd (non-fused w/u path)
so a real GDN forward raises NotImplementedError at trace time. Under
torch.compile this surfaces as a masked backend-compile failure.
Convert both with the same plain-pointer-arithmetic pattern. Verified on gfx950:
chunk_gated_delta_rule(use_qk_l2norm_in_kernel=True) now runs to finite output;
l2norm_fwd matches its torch reference (max |err| 9.7e-04).
* [Triton/Gluon] Gluon MXFP4 Fuse Reduce Quant (#3937)
* Initial first verison of fuse_reduce_rms_mxfp4_quant_kernel(). Included changes to api call and relevant op_test.
* Moved tensors descriptors for second phase into relevant section. Removed redundant layout descriptor. Removed placeholder comment.
* Code Style check.
* Included _mxfp4_quant_op from triton with gluon adaption. Added barrier() to sync threads. Added warning for calling gluon without proper arch.
* Ruff checks
* [Triton/Gluon] Revert Triton parts of #4978 (Dev lumen) (#5149)
Reverts everything PR #4978 (f4e7c7509) changed under `aiter/ops/triton/`
back to its pre-merge state (4ad998328), plus the top-level op_tests that
exercise only those Triton ops.
Reverted:
- aiter/ops/tr…
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.
What
Split-K for the FlyDSL a8w8 preshuffle GEMM, plus retuned Kimi-K3 configs for
M <= 32.
K is split across
gridDim.z. Every split writes its fp32 partial into a[k_split, M, N]workspace, then bumps a per-tile semaphore; the last split toarrive reduces the k_split planes, converts to bf16/fp16 and resets the
semaphore, all in the same kernel. Accumulation stays fp32 with a single
rounding, so precision matches
k_split == 1.The partials are handed between CTAs that may sit on different XCDs, each with
its own L2, so they have to reach a common point. An agent-scope fence does
that with a whole-L2
buffer_wbl2per CTA, which also evicts the A/B tilesevery other in-flight CTA is still reading. Marking just these accesses
sc0|sc1writes them through and leaves the rest of L2 alone.k_split == 1is untouched — same kernel, same cached stores — and its outputis bitwise identical to main on every shape checked.
Performance vs main
Production dispatch (
gemm_a8w8_bpreshuffle), CUDA graph capture, single idlegfx950. Weights are rotated over ~1 GB of copies so every call reads them cold
from HBM, the way
run_perftestdoes — decode sweeps every layer, so a hotweight tensor is not the case to optimise for. main and branch alternated twice
on the same GPU, best of each; repeatability across the two rounds of the same
config was median 0.37%, p90 0.94%, max 2.69%.
The wins land on long K with small N, where the tile grid does not fill the GPU
and K is long enough to split. The unchanged shapes are mostly K <= 1536. The
single regression is at the edge of the run-to-run spread.
Configs
29 rows, all M <= 32: 21 stay on flydsl with a better config, 8 move from ck to
flydsl. 20 use split-K, mostly
k_split=7. Every changed row was re-measuredold config against new on an idle GPU, and rows that were not actually faster
there were kept at main's value.
Testing
k_split == 1zero and the workspace is safe to reuse
k_split == 1output bitwise identical to mainscheduler token, xcd/lds tokens, and
_ksN