Repository navigation
Conversation
* kda gluon gfx1250 implementation * Updated kernel logic and benchmarked * kernel updates * benchmark file * Update op_tests/op_benchmarks/triton/bench_kda.py Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * Added TDM support, cache state support, SK & BV are no longer hard coded, & other optimizations * black / ruff * Update op_tests/op_benchmarks/triton/bench_kda.py Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * Resolving comments * copilot requested changes * copilot fixes pt2 * UT fix * helper file for activation * Adding conv + rms/gate to the kernel * gfx950 gluon kda fused kernel * 1250 optimizations * addressing comment * cast to match 1250 code * ci fixes --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Vinayak Gokhale <vinayak.gokhale@amd.com>
…op count (ROCm#5847) gemm_a8w8_blockscale_cktile chose has_hot_loop/tail_num from the full K, but each split only runs K / (k_batch * K_Tile) loops. With two loops per split the pipeline did not match and split-K outputs were wrong (69%-97% of elements with the default gfx942 kernel). Compute the loop count per split, as the other CK-Tile launchers in aiter do, and make test_splitk_correctness assert and cover short splits. Co-authored-by: Cursor <cursoragent@cursor.com>
…p8 group-quant epilogue (ROCm#5606) * Add fused Qwen3-Next GDN prefill (gfx950 Gluon) with fp8 group-quant epilogue A single fused op for the Qwen3-Next gated delta-net *prefill* chain (causal_conv1d_split_qkv -> fused gating -> fp32 chunked delta scan -> gated RMSNorm -> per-head group-128 FP8 quant), replacing the four-kernel path with a tight set of Gluon launches that fuse the work intra-kernel and drop the intermediate HBM traffic. The FP8 activations a block-FP8 out_proj consumes are emitted by the epilogue launch (no separate quant kernel). Support is reported, not asserted: fused_gdn_prefill_qkvz_supported gates on gfx950, Triton >= 3.8 (Gluon dialect), and an in-range (tokens, batch)/dtype contract, so callers fall back to the four-kernel chain instead of crashing. Prefill spans a wide token range, so it dispatches one of four autotuned M-tile schedules. The Gluon/HIP tiles import lazily on dispatch, keeping the module and the support predicate importable on any platform. Tests: pure-torch fp32 sequential-scan reference (deliberately not the chunked algorithm) at the kernel's bf16 rounding boundaries; 9/9 pass on MI355X (gfx950, Triton 3.8). Includes an op benchmark. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Address review: harden support gate, guard arch detection, conventions Dispatch-safety: fused_gdn_prefill_qkvz_supported() now validates the full contract the fixed-stride Gluon tiles assume, so a mismatched call falls back instead of reading out of bounds: - exact baked topology (k_heads=4, v_heads=8, head_dim=128) via delta_state/ conv_state shapes -- rejects any other 2:1 head layout the tiles can't index; - packed projection widths (qkvz [tokens, 3072], ba [tokens, 16]); - conv_weight [2048, 4] and conv_bias [2048] shapes (not just non-None); - contiguity and single-device for every consumed tensor; - a_log/dt_bias/norm_weight shape+dtype when supplied (kept keyword-only so the existing (..., conv_bias, quant_dtype) positional contract is unchanged). Import-safety: architecture detection is imported and called lazily inside the gate (via _arch_supported) and any detection failure is reported as unsupported, so the module stays importable on CPU-only / non-ROCm workers. Profiling/naming: every launchable Gluon entry across the four M-tiles now has a make_kernel_repr with its constexpr schedule keys, so compiled-artifact and profiler names identify the specialization (matches the flash_kda sibling). Conventions: absolute imports for the new tile package; module-level CUDA skip marker on the test (matches the fused_kda decode sibling); corrected the "single launch" wording in the public docstring and added a full Args/Returns section documenting the shapes, dtypes and in-place state contract. Formatting: the four tile modules were also run through Black -- they predated the repo's Black pre-check (pre-checks.yaml runs `black --check .` over the whole tree) and were not conformant, which accounts for the bulk of the tile diff; the reformat carries no logic change. Validated on MI355X (Triton 3.8): op tests 9/9 pass across all four M-tiles; black + ruff clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Move GDN prefill host orchestration out of the Gluon kernel modules Kernel modules under _gluon_kernels/ are meant to hold only @gluon.jit code (matches the flash_kda sibling). Each of the four M-tile modules mixed torch/ triton host code (buffer allocation + the kernel-launch sequence, plus the FP8 precision config) in with its Gluon kernels. Relocate that host orchestration to a per-tile launcher module next to the public wrapper -- aiter/ops/triton/gated_delta_net/_gdn_prefill_launch_<key>.py -- which imports the kernels from the (now torch-free) tile module. The wrapper's _load_tile imports the launcher for the dispatched tile, so only that tile's launcher (and its Gluon module) is imported on the dispatch path, preserving the lazy-import / import-safety behavior. No kernel logic or numerics change; the host code moved verbatim. The tile-package __init__ no longer re-exports host entries (the wrapper reaches the launchers directly). Validated on MI355X (Triton 3.8): op tests 9/9 pass across all four M-tiles; tiles verified torch-free; black + ruff clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Satisfy ruff 0.16.0 + black 26.5.1 pre-checks The repo's blocking style pre-checks pin newer tools than were used originally (CI installs black 26.5.1 and ruff 0.16.0, whose default rule set is wider). Bring the PR's files into line with both: - black 26.5.1: reformat the op benchmark (it predated the Black pre-check). - ruff C408: rewrite dict() calls as literals (test/bench input builders). - ruff ISC004: parenthesize the gate's multi-line f-string messages. - ruff UP006/UP045: PEP 585/604 annotations (tuple/`X | None`) in the wrapper. - ruff RUF059/SIM102: name unused unpacked vars `_` and collapse nested ifs in the tiles (behavior-preserving -- short-circuit `and`, unused values). - ruff BLE001: keep the arch-detection probe's broad except (defensive: any failure => unsupported) with a scoped noqa. Re-validated on MI355X (Triton 3.8): op tests 9/9 across all four M-tiles; `black --check` and `ruff check` clean under the pinned CI versions. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Address review: CI-run slow tiles, in-tree/fallback contract, tile-sync note Responds to PR ROCm#5606 review (three conditional items): * Test tiles no longer marked `slow`. aiter CI runs `pytest op_tests/triton_tests/` with no `-m "not slow"` filter and no other test uses the mark, so it deselected nothing and only implied the two 12289-token tiles might be skipped. Making them unconditional params is what guarantees CI executes them, so they can't rot. * Module docstring now states explicitly that this op is not adopted in aiter's in-tree GDN path and the fallback is the caller's responsibility (probe `_supported()` -> four-kernel chain on False); the `ValueError` flags misuse after a skipped probe, not a shape-dispatch fallback. * Module docstring documents how the four M-tile schedules stay in sync: they are ported autotuner output pinned to one `ref_gdn_prefill` reference that every tile is parametrized against, so a math change lands in the reference and CI fails any drifting tile. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Address review: consolidate the four tiles into one tile + one launcher module Follow-up to vgokhale's rule-5 note on the four-tile structure. Collapse the four per-schedule tile modules into a single tile module and the four per-schedule host launchers into a single launcher module, so "the kernel" and its host orchestration each live in one canonical place with all four schedules side by side. Same kernels, same perf: every schedule's @gluon.jit bodies and host launchers are byte-identical to before. The only edits are the minimal disambiguations a shared namespace forces: - tiles: b6_15's _multiply / _update_conv_state[_repr] carry a _b6_15 suffix (they collided with b1_5 / m1024_3071); all other schedules keep bare names. - launchers: the one shared FP8PrecisionConfig + FP8_E4M3_FN are hoisted once (with a _PrecisionConfig alias so m1024_3071 stays verbatim); each schedule's public entry becomes gdn_prefill_group_fp8_quant_<key> so all four coexist. - _load_tile now imports the single launcher module and getattrs the per-key entry. _select_tile_key is unchanged and still the dispatcher. Net -93 lines (dropped the duplicated per-file headers/docstrings/config blocks). Validated on MI355X (gfx950, Triton 3.8): all four schedules dispatch to the right entry and match the ref_gdn_prefill reference (test_fused_gdn_prefill_qkvz.py, all 6 cases incl. the renamed b6_15). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Format consolidated GDN-prefill modules (black 26.5.1 / ruff 0.16.0) Blank-line normalization only around the schedule section banners in the merged tile + launcher modules; no executable content changed (git diff --ignore-blank- lines is empty). Fixes the pre-checks black/ruff CI jobs. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Fold the GDN-prefill host launchers into the wrapper module Address review: remove the separate _gdn_prefill_launch.py and move its four schedule launchers into fused_gdn_prefill_qkvz.py, following the gemm/basic/gemm_a16w16.py precedent. Each launcher lazily imports its Gluon kernels inside the function body, so the wrapper -- and the fused_gdn_prefill_qkvz_supported probe on the caller's fallback path -- stays importable on older Triton / non-gfx950 / CPU; only triton.experimental.gluon and the gfx950 kernel modules are gated, loading once a tile is dispatched. _select_tile_key remains the dispatcher (_load_tile returns the module-local gdn_prefill_group_fp8_quant_<key>). Restores the canonical kernel + wrapper + test + bench layout. Schedule bodies are unchanged (AST-identical to the old launcher modulo the injected lazy imports); pytest test_fused_gdn_prefill_qkvz.py passes 9/9 across all four tiles on gfx950 (MI355X). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add BF16 OPUS route-output schedules and gfx950 MiMo TP8 tuned and untuned model configs covering 2K through 128K tokens. Extend the MoE test harness to support bias-free stage 2 while keeping reference precision selectable through --ref-dtype. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
…M/DSV4 shapes (ROCm#5839) MI325X (gfx942, 304 CUs): 286 (M, N, K) rows over 28 weight shapes that no existing model_configs table covers, tuned on main with gemm_a8w8_blockscale_tune.py --libtype ck --splitK. Keys already present in model_configs and rows slower than the heuristic default are left out. The rows go into each model's existing model_configs table; Qwen3-14B and Qwen3.5-122B-A10B get new tables. A shape two models share is kept in one table only, since the config merge rejects duplicate keys across files. Co-authored-by: Cursor <cursoragent@cursor.com>
… configs (ROCm#5896) * feat(flydsl): gfx950 MXFP8 e8m0 GEMM on FlyDSL + DeepSeek-V4/V4.1 configs - New FlyDSL gfx950 MXFP8 batched GEMM (bmm_a8w8_mxscale_gfx950): scaled MFMA 16x16x128, async LDS DMA pipeline, 32x32 / 128x128 e8m0 blocks, split-K with same-XCD last-arrival reduction, B direct to registers, XCD tile order, non-temporal B, per-stage or preloaded scale panels, K % 64 tail, column-major (blockscale) x_scale. Scale rows that are not whole dwords (e.g. 128-wide blocks at K = 384 / 768) are copied a byte per LDS slot with exact buffer bounds. check_bmm_config is the single legality source. - Arch-neutral front door flydsl.batched_gemm_a8w8 dispatching by (arch, w_scale block, x_scale layout); gfx950 module with kernelName parsing, untuned-shape heuristic and a compiled-launcher cache (host 71 -> 21 us). - gemm_a8w8_blockscale_bpreshuffle on gfx950 routes e8m0 x_scale + w_scale to it through the tuned table AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_MXSCALE_ BPRESHUFFLE; FP32 scales keep CK / asm. gemm_a8w8_blockscale with isBpreshuffled forwards native group32 operands there. The V4 wo_a batched GEMM uses the same kernel. - AOT precompiles kernelId=bmm rows (both x_scale layouts for 128x128). - Tuned gfx950 configs: DeepSeek-V4-Pro 128x128 linears (980 rows, TP1/4/8, 1.30x vs CK/asm), DeepSeek-V4.1 group32 (1224 rows, 1.22x vs Triton), V4 wo_a batched (680 rows). - Also carries the other local changes in the tree: inverse_rope_group_quant, opus policy / bmm tune, gfx1250 bmm wrapper, tensor_shim helpers. * refactor(flydsl): tidy gfx950 MXFP8 bmm kernel per FlyDSL cleanup guide - Use fx.copy instead of fx.copy_atom_call for the single-atom copies. - Build the async-LDS DMA destination from the LDS pointer with fx.add_offset + fx.to_llvm_ptr instead of ptrtoint/inttoptr with a hand-picked address space. The raw buffer_load_async_lds stays: the BufferLoadAsyncLDS atom has no 1-byte size and no cache-policy operand. - Factor the compile-and-run + leaked ir.Context recovery out of tensor_shim._run_compiled into _compile_and_run, and use it for the bmm wrapper's per-config compile cache too. Generated ISA is byte-identical on 9 reference configs.
…varlen, fixes, etc (ROCm#5798) Motivation: adds log-sum-exp output to MHA v4 gfx950 kernels (only dense variants now), so it can run under ring / context parallelism, which merges per-rank partials via LSE. Also extends block-sparse to the BF16 recipes, canonicalizes the MXFP4 rows, and improves K/V quantization accuracy. # Kernels: - New BF16 and BF16FP8 sorted-sparse kernels. - Dense LSE epilogue on all ten hd128 recipes, gated at runtime on s_lse. - f8f6, f6f4 and mxfp4 sparse rows moved to FP6-P V, matching their dense siblings. - Dense mxfp4 claims the canonical FP6-P V order; the duplicate f4f4 row is disabled. # Host: - bugfix: per-channel V amax clamped so empty heads cannot quantize to NaN. - Optional lse output on mha_v4 / mha_v4_packed, plumbed through asm_mha_v4_fwd.cu. + ABI unchanged: ptr_lse / s_lse / s_lse_Hs were already reserved in the kernarg. - _LSE_CAPABLE_QV gates the supported format pairs; sorted-sparse still raises. - K mean is subtracted before quantizing (K-smoothing), fused into the MX quantizer kernels. - Dense MHA v4 accepts per-batch key lengths (ragged seqlen_k). # Minor: - Retired MXFP4 Q/K + FP8 V and the deprecated mha_v4_mxfp8 alias. - bench_sage.py: improve input distributions, diffusion-calibrated default, BF16 sparse modes. - Split the block-sparse cases into op_tests/test_mha_v4_sparse.py.
A GPU memory fault aborts the worker process, so its task never returns a result. mp_tuner only noticed at the task timeout (1800s by default) or never without one. Record which worker started each task and, when that process is gone, fail the task and restart the pool right away, like the existing accelerator-error path. Fixes ROCm#5840. Co-authored-by: Cursor <cursoragent@cursor.com>
…als (ROCm#5905) * [MLA v4 nm] Fix _run_one_point reading packed BF16 as FP32 partials Since ROCm#4311 the v4 nm dispatcher derives out_16_nosplit from num_kv_splits and ignores the caller's value, so a single-split launch writes packed BF16 into the logits buffer. _run_one_point still read logits_buf[:, 0] as FP32 for num_kv_splits == 1, which compared every other BF16 element (plus the never-written tail of the buffer) against the reference and printed spurious `fp8_dequant_ref vs asm` checkAllclose failures in the script-mode sweep. Read output_buf, which holds the final result for every split count. Co-authored-by: Cursor <cursoragent@cursor.com> * [MLA v4 nm] Gate accuracy checks so a failed! fails the test checkAllclose only raises on a catastrophic delta; otherwise it logs `failed!` and returns the mismatch fraction. Four of the six accuracy checks in test_mla_v4_nm.py dropped that return value, which is how the packed-BF16 readback bug fixed in the previous commit passed both pytest and the script-mode CI run. Route them through _gated_allclose, which asserts the mismatch fraction against the same tol_err_ratio checkAllclose uses for `failed!`. The script-mode sweep keeps going past a failing shape, prints a summary, and exits non-zero so aiter_test.sh reports the file as failed. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: Cursor <cursoragent@cursor.com>
…OCm#5901) Follow-up to ROCm#5796. The AOT job list for the gfx950 FlyDSL FP8 flash attention was driven by a header-only aiter/configs/fmha_fp8_aot.csv merged with aiter/configs/model_configs/*_fmha_fp8_aot.csv. A single row does not justify the CSV plumbing, so the head shapes now live in a DEFAULT_SHAPES table in the module, the way mega_moe.py already does it, and the comments introduced by ROCm#5796 are trimmed. - fmha_fp8.py: DEFAULT_SHAPES replaces parse_csv()/DEFAULT_CSVS; default_jobs() replaces the CSV walk; --csv is gone (--shape still overrides). The cu_num column is dropped with the CSVs: the kernel is gfx950-only and AOT_ARCH is fixed, so every non-gfx950 row was warned about and skipped anyway. - common.py: FMHA_FP8 returns default_jobs() next to MEGA_MOE instead of going through collect_aot_jobs(). - jit/core.py: drop AITER_CONFIG_FMHA_FP8_AOT and its config-file property, now unused. - README.md: document the table instead of the CSVs. _variant_space()/jobs_for_shape() are unchanged, so coverage is unchanged: --list still emits the same 92 kernel names for Kimi-K3 TP8 (12:12:192:128 varlen_cross) as the CSV-driven version, and setup.py source builds still compile them via run_aot(). Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
…ck (ROCm#4963) * [FlyDSL] gfx942 fp8_mqa_logits: let _auto_variant choose rows_per_block _auto_variant returned f"mfma_r2_w{wpb}", so rows_per_block was pinned at 2 and seven of the nine registered variants -- including every member of the r4 family -- could never be selected. r4 amortizes each KV tile load over twice as many query rows and is faster from seq_len 8 upward. The gap is widest exactly where it costs most: vLLM chunks indexer prefill to fit VLLM_SPARSE_INDEXER_MAX_LOGITS_MB (512 MB), which caps seq_len at 1024 when seq_len_kv is 131072, so the existing seq_len >= 2048 branch cannot fire at long context and every such call took mfma_r2_w4 -- the median of the nine by speed, with the best 1.65x faster. Measured on MI325X (gfx942), seq_len_kv 131072, best variant vs the r2 pick: seq_len 1 r2 25.2 us r4 79.1 us r4 3.1x worse seq_len 4 r2 24.3 us r4 31.5 us r4 1.3x worse seq_len 8 r2 34.0 us r4 32.9 us r4 1.03x better seq_len 16 r2 56.5 us r4 47.5 us r4 1.19x better seq_len 1024 r2 2564.3 us r4 1520.3 us r4 1.69x better Below seq_len 8 the host padding of seq_len up to a multiple of RPB dominates -- at seq_len 1 an r4 kernel computes 4 rows to obtain 1 -- so r2 is kept there and behaviour for those shapes is unchanged. Logits are bitwise identical across all variants at every shape tested, so this is purely a blocking/occupancy change. End to end on 8x MI325X, TP8, GLM-5.2-FP8, 131072 in / 1024 out, concurrency 8: median TPOT improves 7.22% and output throughput 6.60%. This kernel is 16.8% of GPU time at that point. Signed-off-by: Jin Tao <jin.tao@amd.com> * [FlyDSL] gfx942 fp8_mqa_logits: pick RPB on element count, and keep it a divisor Refines the previous commit's rule after a 2-D sweep. That rule keyed RPB off seq_len alone with a crossover measured only at seq_len_kv=131072; sweeping the other contexts shows the crossover is not a seq_len threshold at all, and that a second effect was being read as one. RPB tracks the logits element count. Over seq_len 1..8192 x seq_len_kv 1024..262144 on MI325X, the boundaries land on the same element count at every context: RPB=1 wins below 2**19 elements (27/27 shapes), RPB=2 at 2**19 (6/6), RPB=4 from 2**21 up (38/38), with 2**20 a transition band split 3/3. Keying off seq_len instead put the previous rule on the wrong side at low context: at seq_len 16, seq_len_kv 1024 it chose RPB=4 and ran 1.26x slower than RPB=1. RPB must also divide seq_len. When it does not, the launcher pads with four torch.cat calls; that is a flat ~44 us of host-side overhead, independent of seq_len_kv, and it is the whole of the "small seq_len" penalty the previous commit attributed to wasted rows. At seq_len 1, seq_len_kv 131072: RPB=1 23.1 us, RPB=2 67.8 us, of which the four cats are 44.1 us and pre-padding by hand recovers all of it (21.9 us). So the penalty is not proportional to the padding -- 1 wasted row of 2 costs the same as 3 of 4 -- and it applies to every odd seq_len, which the old rule sent to RPB=2 unconditionally. Stepping down to a divisor is only right while the kernel is cheap relative to that fixed cost, so it is gated to the same 2**21 elements: at seq_len 1025, seq_len_kv 131072 the dividing RPB=1 takes 3880 us against 2601 us for a padded RPB=2. Measured on MI325X, no FLYDSL_FP8_MQA_LOGITS_VARIANT set, old pick vs new: seq_len seq_len_kv old new speedup 1024 131072 r2_w4 2598.7 r4_w4 1623.5 1.60x 1025 131072 r2_w4 2609.3 r4_w4 1594.0 1.64x 512 131072 r2_w4 1190.1 r4_w4 795.3 1.50x 700 50000 r2_w4 661.6 r4_w4 440.3 1.50x 333 12000 r2_w4 118.4 r4_w4 97.1 1.22x 16 131072 r2_w4 58.3 r4_w4 47.4 1.23x 3 131072 r2_w4 66.4 r1_w4 28.1 2.36x 1 131072 r2_w4 69.4 r1_w4 23.2 2.99x 1 1024 r2_w4 66.2 r1_w4 22.2 2.98x No shape measured regresses; the smallest gain is 1.04x. Against the best of the nine variants at each shape, pooled over held-out data (non-power-of-two shapes, a fine seq_len sweep, and head counts 16 and 64), the geometric mean cost falls from 1.45x to 1.03x and the worst case from 3.17x to 1.41x. Logits are bitwise identical across all nine variants at all 180 shapes swept (1620 timings), so this remains purely a blocking/occupancy change. WPB is deliberately left alone. It is worth a few percent at most here, and unlike RPB its optimum moves with the head count -- at 64 heads the current WPB rule costs 1.61x worst case where a fixed WPB=4 costs 1.07x -- so it needs its own sweep rather than a change fitted to one head count. Signed-off-by: Jin Tao <jin.tao@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> * [FlyDSL] gfx942 fp8_mqa_logits: trim _auto_variant comments per review Move RPB element-count thresholds into _auto_variant and shorten the docstring; logic unchanged. Signed-off-by: Jin Tao <jin.tao@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> * [FlyDSL] gfx942 fp8_mqa_logits: scope the _auto_variant step-down note The docstring stated the divisor step-down as an unconditional rule, but it only applies in the middle band: above the top threshold RPB stays 4 and the padding is accepted. Say which band it applies to, and why the top band is exempt. No logic change. Co-authored-by: Cursor <cursoragent@cursor.com> * [FlyDSL] gfx942 fp8_mqa_logits: unit-test the variant selector The shape sweep in test_flydsl_fp8_mqa_logits.py never reaches r4 -- its largest default shape is 1024 x 1560, below the 2**21 threshold -- and nothing asserted _auto_variant directly, so a regression could re-pin RPB to 2 with every correctness test still green. Cover the RPB bands at both thresholds, that the edges track seq_len * seq_len_kv rather than seq_len alone, the middle-band step-down on odd seq_len (and that it stops above the top threshold), the production 1024 x 131072 shape, the unchanged WPB rule, and _resolve_variant precedence so the auto path is confirmed to be the default. Pure shape arithmetic, so no kernel launch. Co-authored-by: Cursor <cursoragent@cursor.com> * [FlyDSL] gfx942 fp8_mqa_logits: put selector tests on the CI path CI only discovers op_tests/test_*.py, so the r1/r2/r4 selector coverage never ran. Move it there, run pytest from __main__, and add two r4 shapes to the GPU sweep so auto-selected padding is launched. Signed-off-by: Jin Tao <jin.tao@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> * [FlyDSL] gfx942 fp8_mqa_logits: fold the selector pins into the op test aiter op tests are plain scripts, not pytest, so drop test_flydsl_fp8_mqa_logits_variant.py and pin the gfx942 auto-selected variant in test_flydsl_fp8_mqa_logits.py instead: a host-only verify_auto_variant table at both RPB band edges, either side of the odd-seq_len step-down, the long-context prefill shape and the WPB switch. It runs in the default verify scenario on gfx942, and its failures count toward the same exit code as the kernel sweep. The selector import sits in the file's existing ImportError guard. Signed-off-by: Jin Tao <jin.tao@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> --------- Signed-off-by: Jin Tao <jin.tao@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: jin.tao@amd.com <jin.tao@amd.com@tus1-p15-g43.tus.tensorwave.lan> Co-authored-by: Felix Li <felix.li@amd.com>
…OCm#4493) gfx1101 (RDNA3, e.g. RX 7800 XT) ships no MHA config, so `_get_config` in `_triton_kernels/attention/mha.py` finds no `configs/gfx1101/triton/attention/mha/DEFAULT.json` and every call to `aiter.ops.triton.attention.mha.flash_attn_func` fails before a kernel runs. Of the architectures in RDNA_ARCHS, only gfx1151 ships one today. Nine of the eleven entries are taken verbatim from the gfx1151 donor (RDNA3.5, added in ROCm#3423, tuned in ROCm#3560), which is the nearest tuned architecture. Two forward entries are tuned on gfx1101 instead of inherited, and they differ from each other only in `num_warps` and `num_stages`: fwd/default BLOCK_M 128, num_warps 8, num_stages 3 (donor: 64 / 4 / 2) fwd/small_head BLOCK_M 128, num_warps 4, num_stages 1 The split uses the `small_head` bucket added in ROCm#4414, which is opt-in per architecture by the mere presence of the key, so this stays a data-only change. It is needed because one `fwd/default` cannot serve both halves on this card: measured against the donor, `M128 w4 s1` is 0.935x on head_dim 64 but 1.269x on head_dim 128, while `M128 w8 s3` is 0.914x on head_dim 128 but 1.062x on head_dim 64. Measured on Windows native ROCm, triton 3.8.0, fp16, 10 independent repeats of 20 iterations after 5 warmups, `torch.cuda.synchronize()` per iteration; a result counts only when the [min, max] intervals across repeats are disjoint. Ratios are against the gfx1151 donor entry, i.e. against what a donor-inherited config would do. head_dim 128, `default` Flux joint 0.914x / 0.891x / 0.892x on three independent torch+ROCm stacks (2.11/7.15, 2.11/10.1, 2.15/10.1), all pinned to the same triton; llama3-8B 0.882x, mixtral-7B 0.880x, kimik25-tp4 0.849x at seqlen 16384 head_dim <= 64, `small_head` SDXL self-attn 0.935x; deepseek-V3 0.658x and glm47fp8-tp4 0.857x at seqlen 16384 The LLM shapes come from `op_tests/op_benchmarks/triton/utils/model_configs.json` (prefill, causal, GQA, batch 1); the sequence lengths are not in that file and are chosen here. Not covered: batch > 1, varlen/thd, sliding window, decode. Note that the `small_head` comment in `_get_config` does not describe gfx1101. It states that 16 < d <= 64 suffers a num_stages=1 pipelining pathology which num_stages=3 cures. On this card the ordering is the opposite -- on the tuned M128/N32/w4 tile, num_stages 1/2/3 measured 2.300 / 2.399 / 2.492 ms. The bucket is still the right mechanism here, for a different reason: the two head_dim ranges want a different `num_warps`, not a different `num_stages`. Signed-off-by: Martin Domanský <ragua@email.cz> Co-authored-by: Mehmet Cagri <mehmet.kaymak@amd.com>
* [Triton/Gluon] Consolidate tuning harnesses * [Triton/Gluon] Simplify tuning harness
…OCm#5860) Addresses review feedback on ROCm#5648 (kept separate to not disturb the approved PR): - mla_gluon.py: the >2GB global_load calls now carry a bounds mask + other=0.0, so the old "No mask needed / in-bounds" comments above them were stale and gave the opposite (unsafe) guidance. Replace with an accurate one-liner. - test_mla.py: allocate the output with device=q_nope.device instead of relying on the module default device; call torch.cuda.empty_cache() before the mem_get_info() free-memory gate to avoid allocator-state flakiness; and assert exact parity (atol=rtol=0) between the >2GB and <2GB paths, which read identical KV and must be bit-identical (a loose tolerance could hide a regression). Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
…OCm#5830) * [Triton] gfx950: tuned defaults for gemm_a16w16, gmm, MHA fwd and PA decode Tuned on MI355X (gfx950) and validated for correctness and no regressions on broader shape sets than tuned. All changes are gated to gfx950 configs/arch. - gemm_a16w16: per-shape configs for N,K = 2048/2048, 4096/4096, 8192/8192, 10240/8192, 57344/8192, 8192/28672 (non-persistent and persistent). Only the tuned M bucket differs from DEFAULT.json. 1.07-2.5x. - gmm: new "large_kn" config (8 warps), selected for K, N >= 4096 with >= 256 rows per group. 1.14-1.26x there; smaller problems keep "default". - MHA fwd (Triton): new "mid_head" config (128x128, 2 stages) for bf16/fp16 with 64 < head_dim <= 128. 1.03-1.16x; d<=64 and d>128 unchanged. - PA decode (Triton): use v2 whenever there is more than one partition for bf16/fp16 KV. 1.9-5.3x for the batch/context sizes that previously hit v1. Unit tests: test_gemm_a16w16 360 passed, test_gmm 48 passed, test_pa_decode 816 passed, test_mha 2060 passed. * [Triton] pa_decode: sort imports (ruff I001) * Move gmm, MHA fwd and PA decode changes to their own PRs Per review, each kernel gets its own PR; this PR keeps only the gemm_a16w16 tuned configs. The removed changes are on branches gfx950-tuned-gmm, gfx950-tuned-mha-fwd and gfx950-pa-decode-v2. * [Triton] gemm_a16w16_persistent: gfx950 tuned config for N=K=2048
…rnel into hi = bf16(fn) and lo = bf16(fn -hi). a simple bf16 downcast for fn inside the kernel drops accuracy more and performs worse for large M. also change the concatenation of the four streams from column major to stream major (k = stream * TILE_K + column). retuned the gfx950 configs. (ROCm#5885)
…OCm#5585) * [Config] Add Qwen3.8-27B TP1 a8w8 blockscale GEMM tunings for gfx942 ROCm#3324 covered this model family at TP=2/4/8 only, so the five widths Qwen3.8-27B drives at TP1 have no tuned entries for gfx942/cu_num=304 and every call falls back to the default -- 275 "not found tuned config" messages per profile run, now 4. This adds 641 rows in one file, 249 decode (M <= 512) and 392 prefill (M 907-65536), from a shape list extracted out of real vLLM server logs rather than generated as an M ladder. Decode, against AITER's own default, measured in a real vLLM serving run -- Qwen3.8-27B-FP8 at 128K context, max-concurrency 1, on one MI325X, byte-identical images, the config bind-mount the only variable, three interleaved repeats per arm on an idle node: decode step 15.444 ms -> 12.449 ms 1.24x Per-shape decode GEMM vs the default: down 2.33x, out_proj 1.51x, in_proj 1.12x, qkv 1.00x, gate_up 1.00x. No decode shape regresses. Decode-step spread across repeats was 0.037 ms and 0.012 ms. The 392 prefill rows ship and are tuned, but their end-to-end effect is currently unmeasurable on this workload: gemm_a8w8_blockscale silently stores only M*N mod 2^31 output elements once M*N reaches 2^31, and at a 65,536-token prefill chunk the 34,816-wide gate_up projection crosses that limit. The truncation depends only on M*N and not on which kernel the config selects -- AITER's default truncates identically -- so it is neither introduced nor worsened by this change. Tuned with gemm_a8w8_blockscale_tune.py --libtype ck --splitK, decode rows then re-timed interleaved because the tuner's min() over short samples is noisy on near-ties; only rows beating the default by >=3% were kept, following ROCm#5421. test_csv_validation (15 tests) and test_config_shape_collision (17) both pass, there are zero duplicate (gfx, cu_num, M, N, K) keys against the 12 other sources in the runtime merge set on main, and max relative Frobenius error is 4.11e-3 against bf16's 2^-8 = 3.9e-3 rounding floor. Scope: only the decode rows have in-situ evidence, and MI300X shares the gfx942/304 key but is untested here. Two rows pin AITER's default kernel (qkv and gate_up at M=1) after the trace measured the tuned picks slower -- pinned, not deleted, because getPaddedM(1,N,K,0) == 16. Six decode near-ties ship no row. gate_up ships 72 of 80 prefill rows; the 8 missing have M*N > INT32_MAX and fault during tuning, but inherit the tuned M=8192 row at runtime via the same padding collapse. Signed-off-by: Pham Binh <phamhuuthanh.binh@amd.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Cursor <cursoragent@cursor.com> * [Config] Thin Qwen3.8-27B TP1 GEMM rows to the padded M ladder Lookup already rounds M through getPaddedM, so a row per traced M is redundant. Keep the power-of-two M list used by the other gfx942 blockscale tables, and omit shapes that main already ships. --------- Signed-off-by: Pham Binh <phamhuuthanh.binh@amd.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: akii96 <aakif.nawaz@amd.com>
… one module, schedule and API (ROCm#5761) Fold the gfx950 (ROCm#5332) and gfx1250 (ROCm#5656) OPUS MXFP4 paged MQA-logits ops into one implementation. The device math of both bodies is unchanged. - One JIT module, module_pa_mqa_logits_mxfp4_opus. Sources move to csrc/kernels/opus_mqa_logits/pa_mqa_logits_mxfp4/: - _opus.h: ABI, kargs, traits; - _sched.cuh: shared builder; - _gfx950.cuh / _gfx1250.cuh: arch bodies, each an empty stub on the other arch; - _kernels.cu: launcher. Namespaces are opus_logits::gfx950 / ::gfx1250, with generic names (pa_mqa_logits_mxfp4_traits and _kernel). The per-arch sources and module_pa_mqa_logits_mxfp4_gfx1250_opus are deleted. - One schedule: gfx1250's build_tiles + build_sched serve both arches. build_tiles is skipped at q_per_block == 1, where the cut is the identity, which saves one launch per forward on every gfx950 call and on gfx1250's MTP=1 decode. - One runtime-arch dispatch: fwd_sched probes the arch once per process, then dispatches on (q_per_block, block_k). An unmatched config raises. - Host-visible bounds on both arches: - num_rows is checked against q, weights and out, and against the lengths of local_ends, local_starts and row_to_batch; - block_tables width is checked in KV tiles; - both kernels drop records with row_id >= num_rows. The caller contract (four conditions) is documented in the module docstring and the header. - One Python API (aiter.ops.opus.pa_mqa_logits_mxfp4: plan_buffers / plan / pa_mqa_logits_mxfp4), with per-arch MqaLogitsVariant instances. gfx1250: qlen4_kv64 / qlen1_kv64. gfx950: qlen1_kv64 / qlen1_kv256, cta_resident 1024. - Retired: ROCm#5332's gfx950 entry points (pa_mqa_logits_mxfp4_sched / _build_sched / _sched_slots / _sched_buffer_ints). - Kept: ROCm#5656's public gfx1250 API, except the gfx950-only plan(block_k=). - -mllvm -enable-post-misched=1 is applied on gfx950 only. It is load-bearing there; on gfx1250 it slowed chunked prefill. - kargs drops six unread fields (144 -> 112 B). - One op test for both arches, op_tests/test_pa_mqa_logits_mxfp4_opus.py. It adds ATOM-shaped padded decode, host-raise and raw row-guard cases, a real NaN control, and on gfx950 a FlyDSL cross-check.
…e mori dispatch path (ROCm#5810) * feat(mega_moe/gfx1250): bind mori tokoff-ext allocator on the mori dispatch path mori's op layer builds the tokoff-ext slot allocator in EpDispatchCombineOpHip.__init__; driving EpDispatchPlan directly bypasses that and leaves EpArgs.tokOffPeers null, keeping dispatch on the serializing cco-window atomic. Build it here from mori's TokOffExt, gated exactly like mori (default on; MORI_EP_TOKOFF_EXT=0 opts out), pass its peer-pointer array to plan.launch, and free it in close(). Needs mori with a public TokOffExt (dispatch_combine_v2.hip_backend). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * style(mega_moe/gfx1250): black-format the tok_off_peers ternary Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: HaonanWang98 <hwang@amd.com>
…= 1 (ROCm#5935) mla_decode_fwd splits each sequence's KV into NUM_SEGMENTS_PER_SEQ segments and lets a reduce kernel merge the per-segment softmax max and expsum. At one segment there is nothing to merge, so the host skips the reduce kernel and hands both scratch pointers the output buffer itself: else: segm_output = out segm_max = out # dummy ptr segm_expsum = out # dummy ptr The decode kernels store M and L unconditionally, so at one segment they write the softmax state straight over the attention output. Nothing faults and nothing warns; the result is simply wrong. Guard both stores. NUM_SEGMENTS_PER_SEQ is a constexpr, so the branch costs nothing when it is greater than one. Only the gfx1250 Gluon kernel reaches this today: select_3d_config ends its gfx12 branch at max(1, ...), while the other branch floors at MIN_SEGMENTS >= 8. The segment count falls as batch x heads grows, so on gfx1250 it reaches one at large batch -- a DeepSeek-R1 TP2 serving run at --max-running-requests 256 sits in that range throughout, and its decode is silently corrupted. The plain Triton kernel shares both the unguarded store and the host-side aliasing, so guard it as well, before a future tuning change makes it reachable. test_mla_decode_fwd stays green either way, which is why this went unnoticed. Its grid does reach one segment for the larger head count, but the corruption lands only at out.flat[token * num_query_heads + head] -- one element in kv_lora_rank, measured at 0.04% to 0.12% of the output -- while the assertion allows a tol_err_ratio of 0.01. Catching it needs a case with no error budget. Checked on gfx1250. DeepSeek-R1 TP2 at page size 64, gsm8k over 2000 questions, scores 0.945 with the fix. A standalone sweep over batch 1 to 256, bf16 and fp8 e4m3 caches, sequence lengths that are not page multiples and pages scattered through the pool gives a worst relative error of 0.0033 for decode and prefill alike. Compared against a torch reference with no error budget, a one-segment batch differs at a 0.08% ratio before this change -- softmax maxima around 100 sitting where attention outputs near 0.02 belong -- and matches to the last element after it. The plain Triton kernel, forced on by disabling the gfx12 branch, matches the same reference at batch 1 to 256, confirming the guard leaves the multi-segment reduce path intact. test_mla.py is unchanged: its 128 pre-existing failures are all shuffled_kv_cache=True on the pipelined kernel, identical before and after. Signed-off-by: Lin, Soga <soga.lin@amd.com>
…erge A persistent variant of the v4 nm fp8 decode kernel (gqa=128, one query token per row). Each call is one launch of grid (2P, 1, 1): the kernel plans the KV split over P partitions from kv_indptr, runs the attention and merges the split partials itself, so callers need no host split plan and no stage-2 merge, and the launch is CUDA-graph safe with varying KV lengths. - hsa/gfx950/mla_v4: mla_a8w8_qh64_qseqlen1_gqaratio64_nm_ps.co, symbol _ZN5aiter39mla_a8w8_qh64_qseqlen1_gqaratio64_nm_psE, registered as the ps=1 row of mla_v4_asm.csv. - csrc: mla_decode_v4_ps_asm in module_mla_v4_asm, with its own documented kernarg struct (same 21 x 16 B layout, repurposed slots marked), shape / dtype / contiguity / workspace checks and its own kernel cache. - aiter.mla: MlaV4NmPsWorkspace, get_mla_v4_nm_ps_workspace(device, num_partitions=128) and mla_decode_fwd_v4_nm_ps(...). One workspace serves calls that are ordered on the GPU; allocate it outside capture. - Tests: accuracy and LSE against the fp8-dequant reference (uniform, ragged, K=0 rows, N up to 3000, K up to 8320), kv_indptr with offset and extra entries, workspace reuse across batch sizes, partition counts, CUDA-graph replay with changing KV lengths, agreement with mla_decode_fwd_v4_nm, argument checks. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
mla_a8w8_qh64_qseqlen1_gqaratio64_nm_ps.co is now reproducible from sources in the repo: nm_ps_src/gen.py patches the disassembly of the ROCm#5195 decode kernel (asm/rt.s) into the persistent kernel, splices in the in-kernel split planner and last-arriver merge (splice/*.hip, compiled with hipcc -S), renames the kernel symbol, assembles for gfx950 and can check the sha256 against the shipped .co. README.md documents provenance, rebuild steps and the kernarg layout. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
…8 prefill (ROCm#5598) Long FP8 prefill at head size 256 matches `D_GEQ_256` (`BLOCK_M=16`, no `SPLIT_UNMASKED_LOOP`) because D outranks Q in the lookup, so it never reaches the hd128 skip key. Add the same `Q_GEQ_256` family under `D_GEQ_256`, with `SW` and `SHUF` siblings left skip-off, so the existing kernel fast path fires only for non-windowed, non-shuffled prefill.
* feat(kimi-k3): add minimal large-M MoE front Signed-off-by: jiacao-amd <jiahui.cao@amd.com> * feat(kimi-k3): enable merged front for M=7 decode Signed-off-by: jiacao-amd <jiahui.cao@amd.com> * feat(kimi-k3): tune merged front for M=14 decode Signed-off-by: jiacao-amd <jiahui.cao@amd.com> * perf(kimi-k3): tune merged front decode bucket Signed-off-by: jiacao-amd <jiahui.cao@amd.com> * test(kimi-k3): validate the merged front across the MTP decode bucket test_decode_full_front_matches_reference stopped at m=16, so the shapes an MTP server actually decodes at were never checked against the torch.mm reference: with num_speculative_tokens=3 a pure-decode step submits num_seqs * 4 tokens, which is 32 and up for any server past four concurrent sequences. Extend the parametrization to 32, 48, 64, 80, 96, 112, 128 and 192 -- the shapes the companion vLLM change enables, and for which kimik3_bf16_tuned_gemm.csv already ships tuned solutions at N=6016, K=7168, gfx950, cu_num=256. No kernel or config change is needed; this only closes the correctness gap left by the old parametrization. gfx950: 9 passed -> 17 passed. Signed-off-by: jiacao-amd <jiahui.cao@amd.com> * style(kimi-k3): satisfy the pinned ruff on this PR's own files `ruff==0.16.0`, the version .github/workflows/pre-checks.yaml pins, reports two errors in files this PR introduces: RUF022 on the `__all__` list in kimi_k3_moe_front.py and I001 on the import block in its test. Both are autofixes and neither changes behaviour -- `__all__` ordering only affects `import *`, and every symbol here is imported by name. The two EXE001 findings that remain in the repo are in unrelated flydsl files and predate this PR. Signed-off-by: jiacao-amd <jiahui.cao@amd.com> * perf(kimi-k3): retune the M=16 merged-front FP32 row The decode shape is 4 * concurrency, and getPaddedM(gl=0) rounds every M <= 256 up to 16, so this row serves both the C1 (M=4) and C4 (M=16) decode buckets -- about 40% of iterations in the 8k/1k replay. It was tuned against a single weight buffer. The merged front's weight is 6016x7168 BF16 = 86.2 MB and MI355X has 256 MB of LLC, so that regime serves most of the GEMM from cache; the shipped 16.005 us is faster than anything the shape can reach when the weight is actually streamed. In the real model 92 distinct MoE layers stream 7.9 GB per decode step and nothing is resident. Re-searched all 681 hipblaslt solutions with the weight chained across 604 MB of distinct buffers (past LLC) and timed inside a CUDA graph, which is how vLLM runs decode: solidx 443935 (shipped) 22.364 us 3857 GB/s solidx 443486 (this) 19.279 us 4474 GB/s -13.8% 443486 is not a new solution -- the table already carries it at M=7. The right kernel was present and keyed to the wrong M. Max relative error vs a torch.mm FP32 reference is 4.9e-07, so this is a dispatch change only. The us/tflops/bw columns are the streamed measurement and are therefore not comparable to the cache-hot numbers in neighbouring rows. * perf(kimi-k3): size the front-GEMM config cache to the whitelist _KIMI_K3_MERGED_FRONT_TOKEN_COUNTS admits 28 distinct token counts but the config cache held 16, so the M values a mixed decode/prefill server cycles through evicted each other. Minor: the tuned table is cached upstream, so a miss costs a few get_padded_m() extension calls rather than a CSV parse, and in the graphed decode path it is paid at capture instead of replay. This just makes the cache cover the set it was sized for. * refactor(kimi-k3): separate Triton epilogue layers Move the launchable epilogue kernel under _triton_kernels/moe, keep only the public wrapper in ops/triton/moe, and add a config-aware kernel repr and CUDA Graph benchmark. Remove Kimi-specific weight packing and hipBLASLt orchestration from the Triton module. Signed-off-by: jiacao-amd <jiahui.cao@amd.com> * refactor(moe): generalize the SiTU epilogue Rename the op to describe the SiTU epilogue it actually performs rather than implying that it projects the incoming activation. Make all branch widths caller-provided, mask partial tiles for arbitrary shapes, and reuse the shared Triton tanh helper. Update the numerical and CUDA Graph coverage plus the benchmark for the generic API. Signed-off-by: jiacao-amd <jiahui.cao@amd.com> * test(triton): use Triton arch helper for SiTU epilogue Signed-off-by: jiacao-amd <jiahui.cao@amd.com> * fix(triton): avoid branch-local mask type collisions Signed-off-by: jiacao-amd <jiahui.cao@amd.com> --------- Signed-off-by: jiacao-amd <jiahui.cao@amd.com> Co-authored-by: Jiahui Cao <jiacao@crs-m2m-cpu-spur-014.us-east2-a.compute.internal>
The persistent code object never reads kernarg slot 0xb8, yet the span workspace [P, 2] was a field of MlaV4NmPsWorkspace, a required argument of mla_decode_v4_ps_asm and size-checked in C++. Remove it from the workspace, the ctypes signature and the C++ entry point, and pass nullptr at 0xb8. Document the three regions of the cnt counter buffer next to its size. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
…rnel Run mla_decode_fwd_v4_nm_ps in a subprocess at large N (up to the 32768 row limit) with every input and output placed at the very end of its own 2 MiB-rounded allocation under the non-caching allocator, so an access one element past q, qrope, the KV pools, kv_indptr (exactly N+1 entries), kv_page_indices, sink, out or lse faults the GPU. Covers the lse and no-lse paths and checks the merge counters return to zero. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
…rator Rename the spliced planner / combine kernels and their labels (t8pre -> ps_plan, t8fin -> ps_merge, T8_* -> PS_*) and regenerate the code object. .text, .rodata (kernel descriptor) and .note (metadata) are byte-identical to the previous build; only local symbol names in .symtab change. Update the expected sha256. README fixes: the row limit is N <= 32768 (out / q byte offsets wrap above it; the cnt / arange sizes come from a 65536-row layout), cnt is int32 with its three regions spelled out, and slot 0xb8 is passed as 0. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
…spatch (ROCm#5857) * [Triton] gmm: gfx950 large-K/N config with JSON-driven shape dispatch Adds a gfx950 "large_kn" gmm config (8 warps) and an optional per-arch "dispatch" rule list in the gmm config JSON. get_config() returns the first rule whose min_K / min_N / min_rows_per_group all match, otherwise "default". Thresholds live only in JSON; arches without "dispatch" are unchanged. gfx950 rule: K >= 4096, N >= 4096, >= 256 rows per group. Measured on MI355X (fp16, CUDA-graph timing): 1.18-1.26x on matching shapes; configs that looked faster elsewhere lost 3-13% on small-K or <= 128 rows/group, which the rule excludes (all 70+ validated shapes under the rule are >= 0.98x). test_gmm: 118 passed. * [Triton] gmm: rename dispatch threshold to min_avg_rows_per_group, test rules, document schema
* [Triton] Add expert-parallel (expert_map) support to moe_gemm_a16w4 Remap weight expert ids to this rank's local experts and skip blocks whose expert is not local (expert_map == -1); the wrapper zero-inits the output so the downstream top-k combine + all-reduce sum correctly. Enables EP for the monolithic A16W4 MoE on gfx942/gfx950. Claude-Session: https://claude.ai/code/session_01CaWDfSGxdbVqtEMVc9mMna * [Triton] EP: skip non-local gates in reduce instead of zero-filling output Per review: replace the wrapper y.zero_() with a per-gate validity mask threaded into reduce_grouped (indx_valid), so the top-k combine skips gates whose expert is not on this rank. routing() now exposes topk_ids on RoutingData so callers can build the mask (gate = expert_map[topk_ids] >= 0). Also wrap an over-length assert for Black. * [Triton] EP: validate expert_map is contiguous int32 on-device The kernel indexes ExpertMap as a flat pointer; reject non-contiguous / wrong-dtype / wrong-device maps before launch, matching the fused routing path's contract (raised in review). * [Triton] EP: require gate_valid whenever expert_map is set With the output zero-fill removed, non-local experts' rows are left unwritten, so the combine must skip their gates via gate_valid. Assert it is supplied so a caller passing expert_map without gate_valid fails loudly instead of reading unwritten rows (raised in review). * [Triton] EP: check expert_map length, set topk_ids on HERD paths - Reject an expert_map whose length != routing_data.n_expts_tot (the kernel reads ExpertMap[expt_id] for every global expert). - Populate RoutingData.topk_ids in the min-unique (HERD) routing paths so the EP gate-validity mask can be built there too. - Document that topk_ids is int16 (as produced by topk); callers index with .long(). * [Triton] Add expert-parallel tests for moe_gemm_a16w4 Rank-local weights + global routing vs. a reference with non-local experts zeroed, covering the top-k combine (scatter) and the sorted local rows (gather-only), plus the expert_map/gate_valid input guards.
…#5230) Signed-off-by: jjuvonen-amd <joni.juvonen@amd.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: Vinayak Gokhale <vinayak.gokhale@amd.com> Co-authored-by: Satya Nikhil Kodukula <nikhil.kodukula@gmail.com>
…it-K (ROCm#5979) The default fp8-blockscale 2-stage fused MoE returns wrong output whenever get_ksplit() asks for model_dim / 256 splits, one 256-wide K tile each. This happens at small decode batches on gfx942 (cu_num 304), for example at one token for (model_dim, inter_dim) = (4096, 256) with 288 experts and top-8, (7168, 128) with 257 experts and top-9, and (2048, 384) with 128 experts and top-8. Every output element differs from the reference, and the values change from run to run. ck_moe_stage1_gemm turns the requested number of splits into CK's KBatch, the number of KPerBlock tiles each split covers: K / (splitk * KPerBlock). With one tile per split that is 1, which DeviceMoeGemmBlockScale reads as "no split": it runs a single z-block over the whole K and skips zeroing the output. The split-K instance still accumulates into that output atomically, and fused_moe allocates it with torch.empty. StrideE also keys the split-K layout (2 * N per row) on KBatch > 1, so the rows overlap as well. ROCm#3997 fixed this in ck_moe_stage1() by turning split-K off for such requests, and the MoE refactor in ROCm#4394 dropped that guard. This fix sits at the KBatch conversion instead, so no caller reaches KBatch == 1 through split-K, whether the split comes from get_ksplit(), a tuned CSV or AITER_KSPLIT. It also keeps the split rather than turning it off: when splitk would leave fewer than two tiles per split, stage 1 uses fewer splits of at least two tiles, KBatch = 2 or the smallest divisor of the tile count that is at least 2. At the (4096, 256, E=288, top-8) one-token shape the operator then takes 32.1 us, against 41.1 us without split-K. A K of a single tile cannot be split and now fails with an error. Requests that already give two or more tiles per split are unchanged. Signed-off-by: Jin Tao <jintao12@amd.com>
* tuning k3 shapes for a16w16 * new configs * additional tuning numbers
…96) (ROCm#5933) GLM-5.3-Flash routes over 288 experts with hidden size 4096, so each of its 42 MoE layers runs a bf16 N=288, K=4096 gate GEMM. gfx950 has no per-shape file for it and falls back to DEFAULT.json, whose M_LEQ_8 and M_LEQ_32 entries set waves_per_eu=8. That caps the kernel at 64 VGPRs; it spills 16 (M<=8) and 21 (M<=32) registers and takes 27-35 us per call where a spill-free kernel takes ~5 us. The new file drops the occupancy hint (waves_per_eu=0) and the .cg load modifier for M<=64 and keeps a single kernel (NUM_KSPLIT=1): a BM=8/16, BN=16, BK=512 tile, 90-98 VGPRs, no spills. The entries were picked from an 8,640-config NUM_KSPLIT=1 sweep timed in a HIP graph over 42 rotating weights, then checked with bench_gemm_a16w16.py. M_LEQ_128 and above are copied verbatim from DEFAULT.json, so larger M resolves exactly as before.
* add mxfp8 a8w8 blockscale for gfx950 * [FlyDSL] Optimize gfx950 M=1 MXFP8 blockscale GEMM (ROCm#5497) * [FlyDSL] Optimize gfx950 M=1 MXFP8 blockscale GEMM * [FlyDSL] Fuse BF16 workspace split-K reduction * Avoid mxscale preshuffle launch regressions * Optimize fused MXFP8 split-K epilogue * Extend fused MXFP8 split-K to small M * Use padded M config for small MXFP8 shapes * Drop redundant M1 MXFP8 tuning row * Format mxscale preshuffle kernel for Black Co-authored-by: solinzby1 <174623473+solinzby1@users.noreply.github.com> * fix ci fail * refine * refine code * fix as suggestion of copilot * Enhance dimension check for tensor contiguity Added check for contiguous tensor in mxscale_preshuffle_kernels.py. Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: solin <bingzhou@amd.com> Co-authored-by: Felix Li <felix.li@amd.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: solinzby1 <174623473+solinzby1@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
topk_select refused gfx1250 at its arch allowlist (_SUPPORTED_GFX), raising "topk_select is not supported on gfx1250". ATOM main calls topk_select for greedy sampling (embed_head.py) and in the rejection sampler, so every model on gfx1250 failed in warmup, e.g. DeepSeek-V4-Pro. The selectors already handle gfx1250: they take the wave size from the arch (wave32 here) and size LDS from get_lds_capacity_bytes() (320 KiB on gfx1250), and topk_per_row lists gfx1250 among its one-block arches. Add gfx1250 to the allowlist and to the matching list in op_tests/test_topk_select.py. Tested: DeepSeek-V4-Pro with ATOM main on 4x gfx1250 (DP4/EP4, real routing) passes warmup and serves greedy completions correctly. The op tests were not run on gfx1250. Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
…d (N=50016, K=6144) at small M (ROCm#6079) * gemm_a16w16: gfx950 config for MiniMax-M3 lm_head (N=50016, K=6144) at small M The MiniMax-M3 lm_head GEMM (TP4 vocab shard, [M, 6144] x [50016, 6144]^T, bf16 out) had tuned rows only at M = 1/2/4/8/16/... . M = 3/5/6/7 missed the exact lookup and padded to the M = 16 row (Triton "auto"), whose DEFAULT M_LEQ_8 config (8 warps, waves_per_eu 8) takes ~550 us on this shape, vs ~115 us for M = 4/8. A decode step of 3, 5, 6 or 7 sequences paid ~450 us extra. - Add GEMM-A16W16-N=50016-K=6144.json: DEFAULT with M_LEQ_8 set to the M_LEQ_16 config (4 warps, waves_per_eu 6). - Point the M = 1/2/4/8 rows of minimax_m3_eagle_bf16_tuned_gemm.csv at Triton (they were flydsl hgemm at ~115 us). tgemm.mm on MI355X, this shape: M = 1..16 all 95-111 us (before: 116 us at M = 1/2/4/8, ~550 us at M = 3/5/6/7; M = 9..16 unchanged). In an ATOM MiniMax-M3 TP4 server trace the lm_head kernel per decode step went from 553-560 us to 93-96 us at bs 3/5/6/7 and from ~116 us to ~93 us at bs 1/2/4/8. * getPaddedM: fine-grained lookup pads M < 16 to the next power of two The tuned-config walk is exact M -> getPaddedM(gl=0, "fine") -> getPaddedM(gl=1, "coarse"), used by tgemm, a8w8 / a4w4 / batched, opus and the CK / cktile a8w8 C++ lookups. For M < 16 the "fine" level (round up to a multiple of 16) is coarser than the "coarse" one (next power of two): M = 3 padded to 16 before 4, M = 5..7 to 16 before 8, so a tuned M = 4 / 8 row was skipped whenever an M = 16 row existed (the MiniMax-M3 lm_head took the M = 16 Triton row at M = 3 / 5..7). gl = 0 now pads 0 < M < 16 to nextPow2(M): M = 3 -> 4, M = 5..7 -> 8. M = 0, M = 9..15 (-> 16) and every M >= 16 are unchanged.
…mul for gfx950 (ROCm#5800) * fix gluon import exception and benchmark + tune with fused_clamp_act_mul for gfx950 * address reviews, M_LEQ_x/any configs + waves_per_eu to improve latency 3-7% * Fix M_LEQ selection and retune gfx1250 fused_clamp_act_mul * add waves_per_eu to each config * update test_fused_clamp_act_mul config comments for the gfx1250 retune for consistency --------- Co-authored-by: ShaoChunLee <Shao-Chun.Lee@amd.com> Co-authored-by: Bruno Mazzotti <bruno.mazzotti@amd.com>
…d_dim 512 (Gemma 4 global layers) (ROCm#5926)
) Six tracked files that nothing reads and nothing will read again, plus the workflow fix that stops two of them returning. CI scratch the bot committed. update-split-tests.yaml wrote split-tests-update-summary.md and then pr-body.md into the workspace and handed them to peter-evans/create-pull-request, which commits every new and modified file when add-paths is unset. Both reached the repo root that way in ROCm#5131, and the tracked copies were stale the moment they landed because the workflow regenerates them on every run. The job only ever means to change one file, so add-paths now names .github/scripts/split_tests.sh, and both scratch files move to the runner temp directory so they are never dirty workspace files in the first place. The .gitignore entries remain as a backstop for local runs, where --summary-path still defaults to the working directory. Merge and POC leftovers that ship in the wheel. hsa/gfx950/pa/pa_bf16_noquant_gqa8_1tg_4w.co.orig and .co.poc_kl_merg arrived with ROCm#3211 next to the real .co. Nothing references either by name, and the ASM loaders build .co paths explicitly rather than globbing the directory, so neither is reachable. They are not merely inert: setup.py copies all of hsa/ into aiter_meta and MANIFEST.in grafts it, so 48 KB of scratch binaries go out in every wheel. *.orig is now ignored so the next conflicted merge does not leave another one behind. One-off files their PRs forgot. probe_noqa_check.py is a single line, x = 1, added in ROCm#4403 to confirm the pinned ruff config flagged it. REPO_CLEANUP_PLAN.md is a pull request description from ROCm#2167 in March. Its Phase 1 shipped -- the .pt test data is gone from main and the size rules are in .gitignore -- and its Phase 2 history rewrite needs a coordinated force push that has not happened in the six months since. Neither file is referenced anywhere in the tree. Left alone: scripts/download_test_data.sh, a no-op stub from that same plan, because docs/tutorials/index.rst still points readers at it. Retiring the script and the doc together is a separate change.
…OCm#5667) Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
* [ASM][HIP][FLYDSL] Add a gfx1250 flydsl persistent MLA decode & DCP feature support * [feat] pin vgpr for nhead128 * [asm][flydsl] opt mla kernel for dspr1(gfx1250)
…LOCK_K == 1 (ROCm#5832) * [Perf] MXFP8 MoE: keep the MFMA in FP8 when MX_SCALE_BLOCK_K == 1 The USE_FNUZ path upcasts both operands to fp32 and broadcasts both scale tiles across the K tile (plus a tl.trans) before tl.dot. That promotes an FP8 MFMA to an FP32 one at 4x operand footprint. It is never a win. ROCm#5658 asserts quant_block_size == 32 and MX_SCALE_BLOCK_K = BLOCK_K // MX_PACK_DIVISOR, so MX_SCALE_BLOCK_K is always 1: one scale per operand per K-step, constant across the dot. Scaling the [BLOCK_M, BLOCK_N] result is then equivalent to scaling the [*, BLOCK_K] operands, and keeps the inputs in FP8. Applied to the main loop and the not-EVEN_K tail. The MX_SCALE_BLOCK_K > 1 branches are preserved byte-identically; they are unreachable while the wrapper pins BLOCK_K, and they disagree on the weight-scale broadcast axis, so at most one of them is correct. gfx942, E=32/N=2048/K=6144, weights already K-major: T=8 0.99 -> 0.30 ms T=64 0.99 -> 0.31 ms T=512 0.99 -> 0.33 ms T=2048 1.00 -> 0.41 ms Correctness across 13 shape/layout configs with random E8M0 scales, including K=96, K=32 and ragged group sizes with zero-token experts: cosine >= 0.999999 vs an fp32 reference. AI assistance was used. The submitter reviewed every line, ran the tests and measurements above, and is responsible for the result. Signed-off-by: akii96 <194881519+akii96@users.noreply.github.com> * [Perf] MXFP8 MoE: equivalence test, and keep FP32 off the fast path Two follow-ups to the previous commit. Test. The existing correctness tests already fail if the hoist is wrong (verified by mutation: dropping b_sc makes test_moe_gemm_mxfp8_correctness fail). What they do not pin down is the identity the optimization rests on, that with MX_SCALE_BLOCK_K == 1 scaling the [BLOCK_M, BLOCK_N] accumulator equals scaling the [*, BLOCK_K] operands. This asserts that directly, with scales that vary along every axis on both operands and N != M, so a transposed or wrongly broadcast scale cannot pass. Two mutants fail it: swapping the scale axes, and dropping the weight scale. The MX_SCALE_BLOCK_K > 1 branches stay uncovered because no in-tree caller can reach them. moe_gemm_mxfp8 is the only caller passing USE_FNUZ=True and it pins BLOCK_K = quant_block_size = 32 while asserting K % 32 == 0, so EVEN_K is always True and MX_SCALE_BLOCK_K is always 1; the other _moe_gemm_a8w8 caller never sets USE_FNUZ. Cleanup. a_f32/b_f32 were computed before the MX_SCALE_BLOCK_K branch and left unused on the == 1 path, which reads as if the change still promotes the MFMA to FP32. Triton was eliminating them -- timings are unchanged either way -- but moving them into the else makes the code say what it does. The broadcast comment moves with them, and the length-1 tl.sum gets a note: a_sc[:, 0] is rejected as an unsupported tensor index. Output is bitwise identical before and after the cleanup across 5 shapes (including the K=96 tail and E=32/N=2048/K=6144); no arithmetic changed. Signed-off-by: akii96 <194881519+akii96@users.noreply.github.com> * [Perf] MXFP8 MoE: read the per-K-step scale with tl.reshape Address review feedback on ROCm#5832: - Read the single per-K-step E8M0 scale with tl.reshape instead of tl.sum over the length-1 axis, in both the main loop and the tail. - Drop test_moe_gemm_mxfp8_matches_per_operand_scaling; the existing correctness test already checks the kernel against an operand-scaled fp32 reference with non-uniform scales. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> --------- Signed-off-by: akii96 <194881519+akii96@users.noreply.github.com> Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> Co-authored-by: akii96 <194881519+akii96@users.noreply.github.com>
Drop hsa/gfx950/mla_v4/nm_ps_src (assembly, patch scripts and HIP splice sources). Nothing at build or run time reads it; the kernel ships as mla_a8w8_qh64_qseqlen1_gqaratio64_nm_ps.co, unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
* Stream JIT compiler progress in ATOM CI * Fix ATOM CI isolation and Kimi-K3 startup
Resolve op_tests/test_mla_v4_nm.py: _run_varlen_point takes main's _gated_allclose (ROCm#5905), which now also returns the mismatch fraction so the split-plan tests can keep asserting on it. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
dict literals for the workspace kwargs and the test inputs; one with-statement for the capture-time rejection check. No behaviour change. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
get_mla_v4_nm_num_kv_splits(num_seqs, num_heads, kv_len) returns the split count the plan helper picks, with no device work. The uniform split_indptr for it is [0, s, 2s, ...], and any prefix of a longer one is valid, so a caller can keep one constant buffer per split count and slice it instead of writing a split_indptr before every call. get_mla_v4_nm_split_plan now builds on it; its behaviour is unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
Nothing calls the plan helper or MlaV4NmSplitPlan any more: callers take the count from get_mla_v4_nm_num_kv_splits and pair it with the uniform split_indptr [0, s, 2s, ...]. The tests build that pair with a local helper, and the mla_decode_fwd_v4_nm docstring points at the count helper. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
… experts (ROCm#6125) Add tuned rows for DeepSeek-V4.1-Flash routed experts sharded across TP (model_dim 5120, per-rank inter 640 at TP4 / 1152 at TP2; target 384x6 and DSpark draft 128x3; tokens 1..32768) with MXFP8 activations on interleaved gate/up rows (A8W4), the precision the checkpoint's reference runs. The default tuner sweep (TUNE_ONLY=flydslv2) only generates the fused-fp8 stage1 + moe2_layout family. At 32-64 tokens that family loses to the heuristic's bf16-out `_gui` stage1 + non-layout stage2 by up to 1.25x, and a v2-only table raised V4.1 TP4 decode TPOT from 8.2 ms to 12.8 ms. The rows here take, per shape, the faster of a v2 sweep and a TUNE_ONLY=flydsl sweep (v1 family). Each row was then checked against the heuristic on the production fused_moe op. The 7 shapes where the tuned pick measured >3% slower there carry the heuristic's own kernel pair instead; their us is the measured op time and us1/us2/err are left empty, since the tuner never timed that pair. End to end (V4.1 TP4, DSpark-5, ISL/OSL 1024/512, conc 8): 2597 vs 2552 tok/s for the heuristic, within run-to-run spread. Also fix a build error from ROCm#5652 in topk_softmax_kernels.cu: `__shared__ alignas(64) T x[]` is rejected by ROCm 7.1.1's clang ("'alignas' attribute cannot be applied to types"), which breaks module_moe_asm and with it the MoE tuner. Use `__attribute__((aligned(64)))`, as elsewhere in csrc.
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
20 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a persistent variant of the gfx950 V4-NM fp8 decode kernel (
ps=1row inmla_v4_asm.csv). One launch per call: a fixed grid of2*Pworkgroups walks per-partition work descriptors, the split plan is computed inside the kernel from the realkv_indptrevery call, and split rows are merged by the last-arriving workgroup. Rows that are not split write the final bf16 output directly. Because nothing about the split depends on host-side shapes, it is CUDA-graph safe with varying KV lengths.hsa/gfx950/mla_v4/mla_a8w8_qh64_qseqlen1_gqaratio64_nm_ps.co+ CSV row (ps=1,lse=1).hsa/gfx950/mla_v4/nm_ps_src/: generator that rebuilds the.cofrom sources (gen.py --checkverifies the sha256). It patches the disassembly of the [ASM] [MLA v4 nm] Swap gfx950 asm decode kernel ROCm/aiter#5195 kernel (asm/rt.s) into the persistent loop and splices in the planner / merge written in HIP (splice/*.hip, compiled withhipcc -S). See its README for provenance and the kernarg layout.csrc/py_itfs_cu/asm_mla_v4.cu:mla_decode_v4_ps_asmentry (documented packed kernarg struct, shape / workspace checks).aiter/mla.py:get_mla_v4_nm_ps_workspace(device, num_partitions=128)andmla_decode_fwd_v4_nm_ps(q_packed, q_rope, kv_packed, kv_rope, kv_indptr, kv_page_indices, sink, workspace, out=None, return_lse=False, lse=None).N <= 32768, row-dense KV pools, rows with K=0 left unwritten (same as the non-persistent kernel). Calls sharing a workspace must be ordered on the GPU; allocate the workspace outside graph capture.Stacked on ROCm#5890's branch (
fix/mla-v4-nm-split-planner).Test plan
op_tests/test_mla_v4_nm.py: 47 passed (accuracy vs reference incl. K=0 rows and long K, lse, offsetkv_indptr, workspace reuse across batch sizes, CUDA-graph replay with changing KV lengths, agreement withmla_decode_fwd_v4_nm, guard-page out-of-bounds checks, argument rejection)nm_ps_src/gen.py --checkreproduces the shipped.cobyte for byte (ROCm 7.2.4 toolchain)🤖 Generated with Claude Code