Repository navigation
Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
github-actions
Bot
changed the title
yuychang
changed the title
…BF16 GEMM tunings (ROCm#5974) Tuned on MI308X (gfx942, 80 CUs) for the TP2 shapes of Qwen3.8-Flash-Next-PTPC-FP8, M = 1..16384. None of these shapes had a tuned entry, so each ran AITER's default kernel. - qwen38_flash_next_a8w8_bpreshuffle_tuned_gemm.csv (60 rows, CK / asm): GDN in_proj_qkvz and out_proj, QSA qkv_proj and o_proj, shared expert gate_up / down. CK-tile candidates are left out: the tuner timed them far below their runtime cost (N=2560, K=3072, M=1: 7.2 us tuned vs 30 us served). - qwen38_flash_next_bf16_tuned_gemm.csv (74 rows, hipBLASLt): MoE router gate, GDN in_proj_ba, QSA indexer projection, hyper-connection projections. - The matching untuned files list the same shapes, in the same order. Previously shipped and registered from ATOM (ROCm/ATOM#2418). Co-authored-by: Cursor <cursoragent@cursor.com>
…OCm#5763) * Add gfx950 MXFP4 SiTUv2 FLAT 16x192/16x32 kernels and KK3 fake tuner CSVs. Route fused_moe 1-stage Situv2 through the new situv2 hsaco and relax asm MXFP4 model_dim to multiples of 512 so remote GPU checks can pin token 8/32 via AITER_CONFIG_FMOE. Co-authored-by: Cursor <cursoragent@cursor.com> * fmoe situv2 16x192: rebuild .co with the LDS DS-offset fix The vendored 16x192 binary predated the D_LDS_RD_DIRECT / D_LDS_RD_REBASED fix in FMOE_MXFP4_G1U1_1TG_4W_32mx1_64nx4_FLAT_SUBGU192.sp3. The DS read offset is a 16-bit immediate, so at PF_D=4 / SUB_GU=192 the fourth gemm2 D-prefetch buffer (LDS byte 73728) truncated to 8192 and read buffer 0's weights: every fourth 256-dim output slice was garbage. 16x32 was unaffected -- its whole group segment is 46080 B, so all four prefetch buffers stay under the 65536 immediate limit. Its .co is instruction-identical to a fresh build and is not touched here. Rebuilt with: X_BF16=1 ACT_TYPE=2 MARKER_ZERO=1 \ poc_kl/scripts/fused_moe/build_fp4_64nx4_flat_subgu192.sh aiter test_moe_2stage.py -q 4 -a situv2 -e 896 -k 16 -p t, logits_diff: 3584x384 t=1 0.238 -> 0.00017 3584x384 t=2 0.215 -> 0.00030 3584x192 t=1 0.2115 -> 0.00026 * Remove temporary csv files * remove redundant asm path * Improve C4 scenario for KK3 * Load U as nt to avoid cache contention with G * Update 16x128 emsort kernel with deeper D prefetch * Remove redundant entry for token=3 from the config * Cleanup kernels and csv files * Update csv and enable nt for U for 16x192 and 16x32 * Fix typo in csv manifest --------- Co-authored-by: Sergey Solo <ssolovye@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
K3's KDA gate projection `f_b_proj` is `[M, 128] x [128, 128 * heads/tp]`.
K is the KDA head dim, 128 at any tp; N is `num_local_heads * 128`, so the
96 KDA heads give N=1536 at tp8 and N=3072 at tp4. Only 1536 was tuned, so
tp4 misses the table on every token.
A miss here is not a slower kernel. `gemm_a8w8_bpreshuffle` falls through
to `gemm_a8w8_bpreshuffle_ck`, whose `rowwise_heuristic_dispatch` has
branches only for `K >= 512` and `K >= 192 && K % 64 == 0`, so K=128 hits
TORCH_CHECK(false, "Unsupported K for heuristic dispatch: 128.
Supported K greater than 192 and K % 64 == 0.")
and the process raises on its first token. Tuning confirms there is no CK
kernel for this shape at all: every `ck` candidate reported "This GEMM is
not supported!", and the winners are all cktile or flydsl.
Same M grid the N=1536 rows carry (64-step between 1024 and 4096), so the
padded-M lookups behave the same at both widths. 62 shapes, no failures;
5 pick cktile at small M and 57 flydsl above, which is the split N=1536
has. Measured on gfx950 / 256 CU:
tp8 M=1..16384 N=1536 K=128 -> OK (before and after)
tp4 M=1..16384 N=3072 K=128 -> FAIL before, OK after
The untuned CSV gets the same 62 shapes so the rows can be regenerated.
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: gbyu-amd <Guanbao.Yu@amd.com>
* initial kernel and tests * fix bench to not use random avg context length * simplify fp8_paged_mqa_logits: operator style for i32 arithmetic, flatten variant scaffolding * perf: upgrade paged MQA logits to mfma_scale_f32_32x32x64 (K_STEPS 8→2, 6-13% faster) * review: un-hardcode ChunkK / block_kv in paged mqa logits host Address PR review feedback (ROCm#4136): ChunkK was hardcoded to _BLOCK_KV in the public flydsl_fp8_paged_mqa_logits call despite compile_fp8_paged_mqa_logits already accepting block_kv. Now ChunkK is forwarded, defaulting to _BLOCK_KV (128) and validated to be a multiple of MFMA_N before dispatch. * style: apply ruff + black to flydsl mqa-logits files * feat(flydsl): paged mqa-logits KVBlockSize>=1 + correctness-only test Kernel: extend fp8_paged_mqa_logits from KVBlockSize==1 to KVBlockSize>=1 (block-flat co-packed layout: KVB fp8 key rows grouped first, then KVB f32 scales). Per-lane gather resolves each logical KV column independently through the block table (kv_indices[b, col // KVB] -> block, col % KVB -> token), so KVB is a compile-time constant that folds to the per-token slot layout at KVB==1. i32 gather-offset ceiling now accounts for KVB*index_dim. Test: add --kv-block-size sweep axis (1, 64) and reference block-flat dequant; convert to a pure correctness test -- drop run_perftest / the fn() wrapper for a direct kernel call, and remove the timing-derived us/TFLOPS output columns (gate stays exact -inf mask + calc_diff < 1e-3). Verified on gfx950 from a fresh compile cache across KVB {1,64}, dtypes {fnuz,fn}, SplitKV {auto,1,4}, heads/head_dim {64,128}, short/long context. * feat(flydsl): add Preshuffle support to paged fp8 mqa-logits + isolation bench Support the production shuffle_weight(16x16) KV data layout in the FlyDSL paged fp8 MQA-logits kernel (Preshuffle=True, KVBlockSize % 16 == 0), making it a true drop-in for deepgemm_fp8_paged_mqa_logits(Preshuffle=True). Only the intra-block key-byte offset changes; the block-table gather, co-packed/token-ordered scale gather, MFMA, and causal mask are unchanged. - _mqa_logits_common: add load_pack_v8i32_preshuffle (closed-form offset into the shuffle_weight(16x16) data section; 8-byte halves stay contiguous). - fp8_paged_mqa_logits: thread a compile-time preshuffle flag through the kernel builder (with a _ps name tag), compile_* cache key, and the public API; replace the NotImplementedError guard with the KVBlockSize % 16 == 0 assert. - test: add a preshuffle path (oracle from the unshuffled cache, kernel fed the shuffled copy) and a small dedicated case set (kvb 16/64, D 64/128, FN patch). - bench_flydsl_vs_gluon: note the A/B is fixed at KVBlockSize=1 (Gluon faults at KVBlockSize>1); KVBlockSize>1 / Preshuffle coverage lives in the flydsl test. - add bench_flydsl_paged_mqa_logits.py: standalone single-config bench (no Gluon) for rocprofv3 isolation profiling, covering both block-flat and preshuffle paths. Validated on gfx950: block-flat and preshuffle correctness pass the unchanged gate (exact -inf mask + calc_diff < 1e-3); rocprofv3 traces the _ps kernel. * refactor(flydsl): move paged mqa-logits kernel into mqa_logits subpackage Upstream relocated the mqa-logits kernels into aiter/ops/flydsl/kernels/mqa_logits/. Move the paged FP8 kernel and its shared helper module there too for consistency: kernels/fp8_paged_mqa_logits.py -> kernels/mqa_logits/fp8_paged_mqa_logits.py kernels/_mqa_logits_common.py -> kernels/mqa_logits/_mqa_logits_common.py Fix the sibling .tensor_shim imports to ..tensor_shim (now one level up) and update the public __init__ import path. No behavior change; public API (flydsl_fp8_paged_mqa_logits) is unchanged. * fix(flydsl): size paged mqa-logits bench KV pool from real block demand Both paged benchmarks sized the block pool from max_model_len and then handed blocks out with `pool[counter % num_blocks]`, so the block table wrapped and every sequence reused the same physical blocks. At B=16, kv_len=32768 that is a 1024-block pool serving 8192 block references: an 8.65 MB KV working set standing in for 69 MB, small enough to live in cache. Every number measured against it was therefore a cache benchmark. That flattered the kernel by ~13% on its own, and distorted the head-to-head against production Gluon unevenly -- given the honest footprint FlyDSL loses 31-51% while the reference loses only 1-4%, since the reference is slow enough that memory is not its constraint. The reported speedup falls from 1.3-4.2x to 1.7-3.1x. Size the pool from the blocks the batch actually needs, build the block table vectorized rather than one element per GPU write, and report the touched footprint alongside the requested bytes. --pool-blocks shrinks the pool deliberately for cache-residency studies, and the vs-gluon bench now fails with a clear message when a shape would exceed the kernel's 2^31 i32 gather-offset limit instead of tripping an assert inside the launch path. Co-authored-by: Cursor <cursoragent@cursor.com> * perf(flydsl): retune paged mqa-logits launch geometry (9.3% across shapes) SplitKV was auto-computed with the production Triton host formula, which ignores how many waves a CTA contains and so drifts with shape: it lands 2.5 waves per SIMD at B=16 but 10 at B=128. Resident wave count is the lever that matters here -- this kernel is bound by memory throughput rather than per-wave latency, so too few waves leaves the memory system idle and too many just queue. Measured across six shapes (batch 4..128, kv_len 4k..128k) the curve is flat between 2.0 and 2.5 waves/SIMD and rises steeply either side: 1.0/SIMD costs 33.9us and 4.0/SIMD 29.4us against 27.6us at the optimum. Target WavePerEU waves per SIMD directly, accounting for the CTA's wave count, and default to paged_w1. At equal resident wave count, splitting a CTA across 2 or 4 waves is slower on every shape measured: a 1-wave CTA schedules more freely, and the waves have no work to share anyway since each owns a disjoint set of KV columns with no barrier or LDS between them. Worth 4.6-14.4% per shape, 9.3% aggregate, and it removes the shape sensitivity -- all six shapes now land within 1% of 27.7us for the same total KV work, where the old defaults spread them over 29.2-32.5us. Verified against the full 870-case correctness sweep (max error 5.3e-4, tolerance 1e-3); the block-flat KVBlockSize=1 path is unchanged. Two public-surface notes for reviewers. This deliberately diverges from the deepgemm_fp8_paged_mqa_logits host formula it previously mirrored. And WavePerEU changes meaning, from a multiplier in the old formula to the target waves per SIMD; its default of 2 is where the curve bottoms out. FP8_PAGED_MQA_LOGITS_DEFAULT_VARIANT changes paged_w4 -> paged_w1. Co-authored-by: Cursor <cursoragent@cursor.com> * refactor(flydsl): use FlyDSL operator sugar in paged mqa-logits The kernel was written against raw MLIR builders -- 58 _to_raw() calls and 65 arith.* calls -- for arithmetic FlyDSL already overloads. The worst case was the store predicate, 17 lines of nested arith.andi/arith.cmpi to say `(lane_div_N == 0) & (col <= q_limit)`. Route unsigned division and modulo through new udiv/umod helpers rather than the bare `//` and `%` operators: fx.Int32 is signed, so `//` lowers to floordivsi, which expands to several instructions to get negative-operand rounding right. Every quantity divided here is non-negative, so going via fx.Uint32 keeps the single-instruction divui/remui the hand-written calls emitted. Verified codegen-neutral: the ATT-decoded ISA is byte-identical to the previous commit, 655 instructions in the same order. 870/870 correctness cases pass at 5.3e-4, timing is unchanged at 28.0us, and ruff reports the same findings as before. Co-authored-by: Cursor <cursoragent@cursor.com> * refactor(flydsl): drop dead helpers and tighten mqa-logits docs _mqa_logits_common carried three helpers nothing referenced: fn_to_fnuz_i64 (45 lines -- the dense kernel defines its own local copy), load_pack_i64 (likewise), and i32_add. Removing them is the bulk of the shrink. Comments and docstrings were 30% of the two files, much of it restating the code or duplicating the PR description. Trimmed to what a reader cannot derive: the co-pack layout contract, the ReLU scale-hoist identity, the i32 gather ceiling, the D-reg-to-head mapping, and the warning that a Python `if` on KVB would become a runtime scf.if and drop branch-local addresses. 932 -> 770 lines across the two files. ISA byte-identical to the previous commit (655 instructions), 870/870 correctness at 5.3e-4, 27.8us unchanged, no new ruff findings. Co-authored-by: Cursor <cursoragent@cursor.com> * refactor(flydsl): drive the paged mqa-logits tile loop with scf.for_ Replaces the hand-built scf.ForOp plus `with ir.InsertionPoint(body)` and trailing scf.YieldOp with the scf.for_ generator, which owns the insertion point and terminator. Same pattern the fmha_gfx1250 kernel already uses. Only two lines, but it removes the last raw MLIR region plumbing from the kernel's main loop. The masked store keeps scf.IfOp: FlyDSL exposes no `if_` helper, and switching to a predicated store would change codegen. ISA byte-identical (655 instructions), 870/870 correctness at 5.3e-4, 27.7us. Also investigated expressing the preshuffle as a FlyDSL layout, as examples/04-preshuffle_gemm.py does in two lines. Not adopted: that example indexes a whole contiguous tensor, whereas our access is a paged gather off a per-lane runtime base. GTensor._linear_offset is element-granular, and while every preshuffle stride happens to be dword-aligned, a 5-D view would just replace one multiply-add chain with another for about two lines. It only pays off alongside FlyDSL-native tensors and copy atoms. Co-authored-by: Cursor <cursoragent@cursor.com> * refactor(flydsl): express the KV preshuffle as a layout, not a code path The paged gather hand-computed byte offsets and branched on Preshuffle to pick between two of them. Both layouts are actually the same coordinate space -- (block, token//16, token%16, K-step, sub-group, lane half) -- and differ only in the stride vector: block-flat (blk, 4D, D/4, 16, 4, 2) token c at c*D shuffled (blk, 4D, 4, 256, 64, 2) token c at (c//16)*16D + (c%16)*16 So the choice moves into make_kv_key_view and the kernel loses its `if preshuffle` branch entirely; one load_pack_kv serves both. The co-packed scale tail becomes a (block, token) view whose base absorbs the KVB*D key offset, replacing the manual scale_byte arithmetic. load_pack_kv slices the view down to the `sub` axis once rather than indexing all six coordinates per load. Without that the invariant part of the address is re-derived in each of the four loads and the kernel grows 49 instructions; with it, 14. Trade, measured rather than assumed: the kernel file drops 15 lines and a branch, the shared module gains 21 for the two view constructors, so the pair is +6 overall. Codegen is 14 instructions larger (669 vs 655) because the hand-written version exploited disjoint bit-fields to fuse addresses into v_or3_b32/v_add3_u32 that the general dot product does not reach for. Speed holds on every path (3 reps each, layout vs previous): preshuffle 27.58 vs 27.83us, block-flat KVB=64 9.78 vs 9.98us, KVB=1 10.20 vs 10.53us. 870/870 correctness at 5.3e-4, no new ruff findings. Co-authored-by: Cursor <cursoragent@cursor.com> * test(flydsl): curate the paged mqa-logits matrix, 870 cases -> 29 The sweep crossed seven axes into 864 cases that compiled only 12 distinct kernels, re-running each about 72 times. The kernel specialises at compile time on (heads, head_dim, KVBlockSize, Preshuffle); batch, next_n, context length, SplitKV and the Q dtype are runtime values, and none of the axes interact -- nothing couples head_dim to SplitKV or the dtype to KVBlockSize -- so the product was testing interactions that cannot exist. q_dtype alone cost 432 cases while selecting no different code path on gfx950, where convert_q_fn is always False. Now: a full cross of the compile-time axes (16 cases, one per distinct kernel -- more preshuffle coverage than before, which only hand-picked 4), each runtime axis varied alone off a base, and four interaction cases where the tile-range arithmetic genuinely couples. --exhaustive still runs the original 870 for release validation. Validated by mutation testing rather than by the suite passing, since a suite cut too far also passes. Eleven deliberate kernel bugs -- swapped preshuffle and block-flat strides, causal off-by-one, dropped tail clamp, wrong lane half, wrong block-table index, dropped kv_scale, wrong head weight, floor-instead-of-ceil split (needs SplitKV>1), batch-blind output row (needs batch>1), next_n-blind causal bound (needs next_n>1). The full suite catches 11/11; the curated matrix also catches 11/11. 29 cases in 8.7s against 870 in 79s. Harness in att_profiles/mutation_check.py. Co-authored-by: Cursor <cursoragent@cursor.com> * test(flydsl): cover the wave-split and ChunkK axes Neither was ever exercised: every case ran the default variant at ChunkK=128, so paged_w2 and paged_w4 had no coverage at all, and the Tier-1 change of DEFAULT_VARIANT silently moved which one CI touched. That hole was real, not theoretical. A mutation of the wave-to-column-tile assignment (`wave * N_TILES_PER_WAVE` -> `wave * (N_TILES_PER_WAVE - 1)`) is correct by construction at one wave per CTA, where `wave` is always 0, and wrong for any multi-wave variant. It survived both the curated suite *and* the original 870-case exhaustive sweep. Adds variant and chunk_k parameters plumbed through to the kernel, reported in the results table, and six cases: paged_w2 and paged_w4 at the default tile width, ChunkK 64 and 256 against the variants that divide them, and a preshuffle x paged_w4 cross. 35 cases in 9.0s, and the mutation set now runs 12/12 where the exhaustive suite manages 11/12. Co-authored-by: Cursor <cursoragent@cursor.com> * test(flydsl): condense the paged mqa-logits test prose and plumbing Docstrings trimmed to what a reader cannot derive -- the oracle's formula and mask rule, the co-pack byte layout, what the preshuffled copy is for -- dropping the restatement and the cross-references to other files. _build_inputs returned an 8-tuple that was unpacked over ten lines at the call site; it now returns a NamedTuple, so the test body reads inp.q_fp8 rather than positional unpacking, and the contract is documented by the type. Its block table is also built vectorized instead of one element per GPU write, which is why the suite went from 9.0s to 7.5s. Same values: the per-case errors are unchanged to the digit. The pool there is still sized from max_model_len, so sequences share blocks once the batch outgrows it. Left alone deliberately, with a comment saying why: harmless for correctness and arguably extra coverage of block reuse, unlike in the benchmarks where it shrank the measured working set into cache. Still 35 cases, still 12/12 on the mutation set. Co-authored-by: Cursor <cursoragent@cursor.com> * docs(flydsl): drop benchmark figures from the paged mqa-logits source Timings and percentages go stale as soon as anything around them changes, and the source is the wrong place to keep them -- the numbers behind these choices are in the commit history and the PR description. Removes the measured wave-count curve from _auto_split_kv, the "4-14% slower" claim from the DEFAULT_VARIANT note, and the "where the measured curve bottoms out" aside from the WavePerEU docs, keeping what the formula computes and why resident wave count is the lever. In the test, the block table comment keeps its warning that shared blocks are deliberate but no longer cross-references what that would cost a benchmark. Comments only; 35/35 correctness cases unchanged at 3.97e-4. Co-authored-by: Cursor <cursoragent@cursor.com> * docs(flydsl): fix stale --variant default in the paged bench help The help text still advertised paged_w4 after DEFAULT_VARIANT moved to paged_w1. Read it from the constant so it cannot drift again. Co-authored-by: Cursor <cursoragent@cursor.com> * refactor: move the FlyDSL benchmarks out of op_benchmarks/triton They were sitting in the Triton benchmark directory alongside 73 unrelated files, which is misleading -- neither uses Triton. op_benchmarks/flydsl/ now holds both, matching how op_tests/flydsl_tests/ already separates the FlyDSL correctness tests. No __init__.py needed: op_benchmarks subdirectories are namespace packages, which is also why the vs-gluon benchmark keeps importing kv_cache_cast_to_fp8 from op_benchmarks.triton.bench_deepgemm_attention -- that file is genuinely a Triton benchmark and stays put. Paths updated in the bench's own rocprofv3 example and in pr.md. Both run clean from the new location: microbenchmark 27.8us with its correctness gate passing, head-to-head 3.11x at B=16. Co-authored-by: Cursor <cursoragent@cursor.com> * style(flydsl): satisfy the Black and Ruff pre-checks on PR ROCm#4221 CI reported five failures; only two are real. check-signal waits on the pre-checks workflow and Aiter Test Gate depends on check-signal, so the gate failed in four seconds without running a single GPU test. Fixing style unblocks all three. These findings were invisible locally because I had been diffing ruff against the branch's own history. CI runs it through reviewdog with filter-mode=diff_context against upstream/main, where none of these files exist -- so every line is an added line and every finding is in the diff. Seven auto-fixed: import ordering in five files plus __all__ ordering in aiter/ops/flydsl/__init__.py. Five needed a decision, taken to match what the repo already does rather than inventing a convention -- it carries 214 BLE001, 21 B008 and 19 PLR0124 suppressions, including this exact `fx.Stream(None)` default in moe_route_maps.py and this exact NaN-check idiom in test_fused_qk_norm_rope_group_quant.py. So noqa for the blind except, the FlyDSL stream default and the self-comparison, and the two unused `heads` unpackings renamed to `_heads`. Black reformatted one file. Verified as CI runs them, repo-wide: `black --check .` leaves 1063 files unchanged, and `ruff check .` leaves five findings, all in upstream files this PR does not touch. Behaviour unchanged: 35/35 correctness at 3.97e-4, 27.93us, and the public exports still resolve after __init__ was resorted. Co-authored-by: Cursor <cursoragent@cursor.com> * refactor(flydsl): drop dead conversion plumbing from the paged kernel convert_q_fn and convert_kv_fn were threaded through the kernel builder, the compile cache and the host, and neither is ever read inside the emitted kernel -- they only appended a name tag. convert_kv_fn was additionally hardcoded False at every call site. Only scale_mul reaches codegen, so it is now the sole carrier of the FN->FNUZ compensation and tags the kernel name itself. That exposed a docstring that was simply untrue: compile_fp8_paged_mqa_logits claimed these flags "mark an FP8 FN operand whose -0 (0x80) byte the kernel patches to FNUZ +0". No such patch exists on this path -- the vector<8xi32> operand has no byte-patch, which the kernel body's own comment admitted. The host comment now states the actual gfx942 behaviour, including that FN -0 decodes to FNUZ NaN unpatched. Flagged rather than fixed: it needs gfx942 to validate. Also: mfma_head_reduce now derives res_ty and f32_0 and defaults fm_fast, so the call site is one line instead of eleven and the kernel drops two locals; make_out_row_view takes the i32 stride and widens it itself. Both are safe because _mqa_logits_common turns out to have exactly one consumer -- upstream's dense fp8_mqa_logits carries its own copies and does not import it, so the module docstring claiming otherwise is corrected too. 569 -> 535 lines in the kernel. ISA byte-identical to HEAD (669 instructions), 35/35 correctness at 3.97e-4, 27.92us, black and ruff clean. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(flydsl): pass buffer load/store aux as an i32 SSA operand upstream/main treated aux as an optional IntegerAttr, but RawPtrBufferLoadOp expects it as the fourth operand Value. With cache_modifier=0 that left aux None and broke JIT compilation of kernels using GTensor buffer loads. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(flydsl): align buffer_ops and MFMA helpers with flydsl 0.3.1 Revert buffer load/store aux to the upstream IntegerAttr form expected by flydsl 0.3.1, and route scaled MFMA through the expr wrapper so cbsz/blgp attributes lower correctly during paged MQA kernel JIT. Co-authored-by: Cursor <cursoragent@cursor.com> * perf(flydsl): software-pipeline block-table and K0 prefetch in paged MQA logits Prefetch kv_indices one tile ahead (skip reload when ind_off unchanged) and carry K-step-0 across the inner column loop to hide VMEM latency on the production kvb64/preshuffle decode path. Co-authored-by: Cursor <cursoragent@cursor.com> * style(flydsl): fix Black/Ruff on paged MQA logits prefetch helpers Bind loop SSA values via default args (B023) and apply Black formatting. Co-authored-by: Cursor <cursoragent@cursor.com> * address comments by Sami * style(flydsl): fix Black/Ruff on paged MQA logits files Apply Black formatting and resolve Ruff PLR0124/I001 so CI pre-checks pass. Co-authored-by: Cursor <cursoragent@cursor.com> * feat(flydsl): add gfx950 H64 16x16x128 paged FP8 MQA logits kernel Land the production indexer mapping (resident Q, depth-4 NT B-ring, SplitKV) without the v2/v3 experiments or local profiling artifacts. Co-authored-by: Cursor <cursoragent@cursor.com> * feat(flydsl): specialize gfx950 paged MQA logits for H=32 and H=64 Compile a separate 16x16x128 kernel per head count so the indexer matches gluon's ChunkQ=heads set without changing the FlyDSL page pipeline. Co-authored-by: Cursor <cursoragent@cursor.com> * feat(flydsl): compile-time next_n for gfx950 paged MQA logits Drop the unused second query row on next_n=1 so MFMA/VALU and store waits match the actual decode shape. Co-authored-by: Cursor <cursoragent@cursor.com> * feat(flydsl): dual-bank page pipeline for gfx950 paged MQA logits Keep next-page K in fixed register banks instead of scf.for yield copies, store one wave-wide page of logits, and prefetch the page table so gather latency stays overlapped. Co-authored-by: Cursor <cursoragent@cursor.com> * feat(flydsl): dispatch gfx950 paged MQA logits as the default decode path Route H32/H64 D128 KVB64 preshuffle next_n in {1,2} through the 16x16x128 kernel, keep the 32x32x64 implementation as the generic fallback, and satisfy Black/Ruff. Co-authored-by: Cursor <cursoragent@cursor.com> * refactor(flydsl): make gfx950 the only paged FP8 MQA logits path Drop the generic 32x32x64 kernel and fold helpers into the gfx950 decoder so H32/H64 stays on the tuned 16x16x128 pipeline, with i64 output-row addressing and idiomatic FlyDSL MMA/scheduling that still matches the previous ISA. Co-authored-by: Cursor <cursoragent@cursor.com> * style(flydsl): black-format gfx950 paged MQA logits and tests Co-authored-by: Cursor <cursoragent@cursor.com> * dynamic next_n scheduling WIP. * cleanup * feat(flydsl): accept compact Q[B, next_n, 32, 128] for gfx950 paged MQA logits Match the Gluon contract so next_n can be 1..8 without padding Q to 8 rows; optional next_n_lens still selects live rows per sequence. Co-authored-by: Cursor <cursoragent@cursor.com> * perf(flydsl): raise wave priority around scaled MFMA Give compute waves s_setprio 3 during gemm so they win over helper G2L. Co-authored-by: Cursor <cursoragent@cursor.com> * perf(flydsl): hold s_setprio across both scaled MFMAs per row Avoid dropping priority between the two 16x16x128 gems. Co-authored-by: Cursor <cursoragent@cursor.com> * perf(flydsl): hint waves_per_eu=6 to match LDS-bound occupancy Compiler was told 2 waves/EU while the CTA actually lands 6. Co-authored-by: Cursor <cursoragent@cursor.com> * perf(flydsl): unroll next_n scoring by two rows Issue a pair of Q/weight LDS packs per iteration of the row loop. Co-authored-by: Cursor <cursoragent@cursor.com> * perf(flydsl): unroll next_n scoring by four rows Give the compiler a larger Q/weight LDS pack window per tile. Co-authored-by: Cursor <cursoragent@cursor.com> * perf(flydsl): fully unroll next_n scoring over MAX_NN Replace the 4/2/1 while ladder with predicated constexpr rows. Co-authored-by: Cursor <cursoragent@cursor.com> * perf(flydsl): cache paged MQA logit stores for the following topk Write logits with buffer cache policy 0 so L2 can retain the row for the next kernel; keep KV page DMA nontemporal so those loads do not evict it. Co-authored-by: Cursor <cursoragent@cursor.com> * refactor(flydsl): drop dead gfx950 paged-MQA helpers and restore scalar logit stores Scalar GTensor writes already lower to cached buffer_store; the explicit store path and unused arith/B_RING bits did not change ISA. Share MAX_NN from the kernel in tests. Co-authored-by: Cursor <cursoragent@cursor.com> * test(flydsl): keep paged MQA driver out of pytest collection Co-authored-by: Cursor <cursoragent@cursor.com> * cleanup(flydsl): remove stale paged MQA test paths Co-authored-by: Cursor <cursoragent@cursor.com> * test(flydsl): put paged MQA logits on the aiter op_test template Move the gfx950 indexer tests to op_tests/test_*.py with @benchmark candidates, TFLOPS/TB/s, a torch-only reference, and one markdown table per signature so CI runs python3 op_tests/test_flydsl_fp8_paged_mqa_logits.py. Co-authored-by: Cursor <cursoragent@cursor.com> * test(flydsl): fold paged MQA pytest coverage into the op_test tables Keep a single gfx950-gated test_*.py as the source of truth, grade wide output against torch, and leave the flydsl_tests files as shims. Co-authored-by: Cursor <cursoragent@cursor.com> * test(flydsl): drop paged-MQA shims and add the probe_1079 decode table Keep a single op_test as the source of truth, with the production scattered-pool contract (mean Nq, next_n × ctx) graded against torch. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: yashagar <yashagar@amd.com> Co-authored-by: Folkert Huizinga <fhuizing@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
* fix * shrink ut
* [CI] Triton board: drive PR status from the workflow (draft=Backlog, open=In review) * [CI] Triton board: work from the current PR state; remove items that lose their label/assignee * Clean up * [CI] Triton board: remove only on a lost qualifier; drop the enable variable - `unassigned` takes an item off the board only when the person removed is a member of @ROCm/triton_kernels_team (and no rule still applies). Unassigning someone unrelated used to remove hand-added items. - When the team lookup fails, nothing is removed: without the team we cannot tell whether an assignee still qualifies the item. The item is left alone with a warning; adding through the other rules still works. - The job no longer waits for the repository variable AITER_TRITON_BOARD_ADD.
yuychang
force-pushed
the
feat/kimi-k3-kda-decode-flydsl
branch
from
September 30, 2026 14:26
0caa013 to
1d48028
Compare
github-actions
Bot
changed the title
yuychang
force-pushed
the
feat/kimi-k3-kda-decode-flydsl
branch
from
September 30, 2026 14:41
1d48028 to
145a1d8
Compare
- Add fused KDA decode and head-local f_b projection kernels for the fixed Kimi-K3 TP8 layout - Use operation-based module, kernel, public API, and test names without compatibility aliases - Export decode APIs and the gfx950 support predicate through the FlyDSL package - Add reference correctness, cache behavior, input validation, and public API tests
yuychang
force-pushed
the
feat/kimi-k3-kda-decode-flydsl
branch
from
September 30, 2026 14:53
145a1d8 to
301f09d
Compare
yuychang
added a commit
to yuychang/sglang
that referenced
this pull request
Sep 30, 2026
…decode - Match yuychang/aiter#2, which upstreams the kernels as aiter.ops.flydsl.kda_decode - Load the AITER source from aiter.ops.flydsl.kda_decode and its flydsl_kda_decode* entry points - Rename the vendored module to kda_flydsl/kda_decode.py so both sources expose one API - Rename the kernel builders to kda_decode_gfx950.py and kda_decode_fused_projection_gfx950.py - Rename the public functions, builders, compiled kernel names and registry op ids to the AITER spelling - Keep the vendored batch-2 fused-projection tuning as optional builder keywords - The vendored and AITER kernels match bitwise on gfx950 at batch 1, 2, 8 and 16
* feat(triton): add SiLU multiply backward Signed-off-by: DaiXindi-AMD <xdai@amd.com> * fix(triton): address SiLU multiply backward review Signed-off-by: DaiXindi-AMD <xdai@amd.com> * fix(triton): narrow SiLU backward support to gfx950 Signed-off-by: DaiXindi-AMD <xdai@amd.com> --------- Signed-off-by: DaiXindi-AMD <xdai@amd.com>
* Fix gfx950 gluon unified attention compile failure with int64 seq_lens * Remove int64 metadata unified attention test
github-actions
Bot
changed the title
yuychang
changed the title
- Sort KDA decode imports to satisfy Ruff I001. - Sort public exports to satisfy Ruff RUF022 without changing kernel behavior.
github-actions
Bot
changed the title
- Rebase the state and conv-state buffer descriptors on the selected slot using 64-bit byte offsets. - Pass the state and conv-state slot strides to the kernels and launchers as fx.Int64. - Keep per-access buffer offsets 32-bit and slot-relative, which fixes wrapped reads/writes past 4 GiB. - Add a large-cache regression test for both decode entry points that checks results and untouched slots.
…weights (ROCm#6092) * Fix Triton Conv2D/Conv3D caching for inference tensor weights * Review 1 Simplify inference-tensor weight handling * Update Conv2D test inventory count
* Add per token scale to a8w8 moe gemm * cleanup
…gfx950 (ROCm#4383) * Merge branch 'ROCm:main' into quant_mxfp4 * Update bench_quant_mxfp4.py * Resolve formatting issues * Resolve formatting issues * Resolve formatting issues * Resolve formatting issues * Resolve typing deprecation and import formatting * Fix Ruff import and __all__ ordering * Replace deprecated typing.Tuple in quant module * Use PEP 604 unions for optional tensor annotations * Include instruction for mxfp4 qunatization. * Change hueristic Added pre fetch slight modification in bench script * Consolidate comments and fix formatting/duplicate imports * Modifying shapes to be tested * Remove the inline asm part as the test triton patch provided by compiler team, now successfully uses the correct instruction * Removing unwanted files * Removing json file * Migrate gfx950 Gluon mxfp4/mxfp8 quant tuning to JSON configs Also update mxfp4/mxfp8 quant test files with representative shapes covering the benchmark shape sweep. * COnfig file format change * Address PR review comments: bucket-key config lookup, kernel repr, lazy Gluon import, config dir rename, trim redundant test shapes - Replace the rule-tree select_tuned_config with lookup_tuned_config, a flattened bucket-key resolver ({schema, "M_GEQ_33.N_LEQ_16384": {...}, "any": {...}}) matching unified_attention_utils's existing pattern. - Add make_kernel_repr-based repr= to both gfx950 Gluon quant kernels. - Move the gfx950 Gluon kernel imports in quant.py out of the module header into the arch-gated dispatch branches that use them. - Rename configs/gfx950/gluon/quant/{mxfp4,mxfp8}_quant -> {mxfp4,mxfp8} to match the sibling PR's naming convention. - Drop one of three redundant large-non-square shapes from the mxfp4/ mxfp8 quant parametrize lists. * docs(quant): shorten config lookup and gfx950 kernel comments * Move gfx950 Gluon quant shape-derived block sizes into JSON and retune mid-M configs - Drop the Python BLOCK_SIZE_M/N overrides in dynamic_mxfp4_quant and dynamic_mxfp8_quant; launch configs now come only from the JSON tables. - Tune the 33<=M<=1024 buckets (more, smaller CTAs) for MXFP4 and MXFP8. - Compute e8m0 scale bias in int32 before narrowing to uint8. * Use lookup_config for gfx950 Gluon quant configs to match the sibling PR Replace lookup_tuned_config with lookup_config(table, axes, **values) and drop the "schema" key from the gfx950 mxfp4/mxfp8 JSON tables. * Refactor gfx950 Gluon quant: shared helpers, extended select_leq_config, backend argument - Share the e8m0 scale math (common/utils.py), the tile store and the warp layout between the MXFP4 and MXFP8 Gluon kernels. - Extend select_leq_config to many axes and reject bad calls; add get_quant_config and drop lookup_config. - Add a keyword-only backend argument (None, "triton", "gluon") to dynamic_mxfp4_quant and dynamic_mxfp8_quant; the benchmark gets a gluon provider and the triton provider forces the Triton kernel. - Drop unused SCALING_MODE and waves_per_eu=0 from the Gluon launches; add num_warps and EVEN_M_N to the kernel repr keys. - Add CPU-only tests for the config lookup and the Gluon quant JSON tables; update the config docs. * Fix black formatting in test_quant_configs
…mm_a8w8_blockscale (ROCm#6071) Both benchmarks crash before printing a single result on current main. bench_moe_gemm_int8_smoothquant passes add_residual=False to moe_gemm_int8_smoothquant, but the MXFP8 GEMM and A8W4 MoE optimization (ROCm#3292) renamed that argument to swiglu_add_residual. The meaning is unchanged: both compute s * (linear + 1) when set. The call now fails with "unexpected keyword argument 'add_residual'". bench_moe_gemm_a8w8_blockscale builds the routing GEMM inputs in fp32 (xg = x.to(torch.float32), and wg with the default dtype) and passes them to gemm_a16w16. Since ROCm#4860, gemm_a16w16 asserts fp16/bf16 activations and weights. Build both in bf16, as bench_moe_gemm_a8w8 and bench_moe_gemm_int8_smoothquant already do.
The benchmark cannot run: it imports generate_gemm_afp4wfp4_pre_quant_inputs from op_tests.triton_tests.gemm.basic.test_gemm_afp4wfp4_pre_quant_atomic, a module that no longer exists. The DS FP4 fusion work (ROCm#1371, ROCm#1434) renamed that test to test_gemm_a16wfp4.py, so every run fails with ModuleNotFoundError. ROCm#2601 tried to restore the helper. It was closed because gemm_afp4wfp4_pre_quant is a deprecated wrapper, superseded by gemm_a16wfp4 and slated for removal, and the bench should be deleted. This does that, and drops the bench's entry from bench_schema.yaml.
* [Triton/Gluon]: Add Inkling Model GEMM Configs * Switch to using FlyDSL --------- Co-authored-by: Rahul Batra <rahbatra@amd.com>
* DSV4.1 ATOM-port local fixes (flydsl MoE LDS-DMA wait, fp4 prefill cache, MoE tuner GPU datagen, torch.Stream, tuned FMoE CSV) - flydsl mixed MoE stage1: full vmcnt drain + sched_barrier before the LDS-DMA barrier (partial waits let LDS-DMA loads cross it; ROCm#5561) - fp4 prefill mqa logits: unbounded kernel cache (lru_cache(32) evicted and recompiled per shape) - MoE tuner: generate the E x 2I x H test weights on the GPU (CPU fp4 packing took tens of minutes per shape at E=385) - cpp_itfs torch_utils: accept any torch.Stream (derive the HIP handle when cuda_stream is missing) - add DSV4.1 TP2 shared-experts-fusion fp8/fp4 tuned FMoE config * Merge DSV4.1 TP2 fp8/fp4 FMoE tuning into dsv41_fp4_{tuned,untuned}_fmoe.csv - move the 23 shared-experts-fusion rows (E385 topk7, E129 topk4, fp8 activation x fp4 weight) into dsv41_fp4_tuned_fmoe.csv and their shapes into dsv41_fp4_untuned_fmoe.csv instead of a separate file - drop the torch_utils, fp4 prefill cache and flydsl stage1 wait changes --------- Co-authored-by: wunhuang <wunhuang@amd.com>
…d 64-row pages (ROCm#6145) * feat(flydsl): row-group FP4 paged MQA logits for 8- and 64-row pages Add pa_mqa_logits_fp4_rowgroup, an FP4 paged MQA logits kernel in which a wave scores a group of a sequence's query rows off each key load (the existing kernel reads every key once per row), and route flydsl_pa_mqa_logits_fp4 to it for kv_block_size 8 and 64. - Page 64 keeps the existing preshuffled cache layout and its writers; page 8 (DeepSeek-V4.1's indexer pages) has its own layout (pack_kv_cache), so a candidate layer's kept blocks are its block table with no gather. - Work is shared out by each sequence's real length inside the kernel (static split, or a length-balanced share when that walks shorter); the grid depends on the batch and max_seq_len only, so a captured graph keeps it. No schedule launch; cta_info / block_k / num_warps / parallel_unit_num do not apply on this path. - Keys and scales are read through 64-bit addresses (pools past 4 GiB); sequences are handled in runtime chunks of 64 (one build per shape for any batch). - A given `out` is left as it was past each row's bound (no -inf fill). - compute_varctx_schedule moves to pa_mqa_logits_fp4_common. Cold, MI355X, against the original FlyDSL kernel (its schedule prebuilt): page 64 1.07-3.0x (B = 64, H 32/64, next_n 1/2/6, up to 100K keys) and 1.25-2.9x at B = 1-16; page 8 V4.1 full layers 1.5-3.4x (T = 6-768, 200K context). Against aiter's Triton paged_mxfp4_mqa_logits: page 8 1.03-1.5x on every T; page 64 ahead or level on most shapes. Also: rope_rotate_activation (fp4) gains round_rope, rounding the rotated values to the input dtype before quantizing them, as a model whose RoPE writes its output back in that dtype (DeepSeek-V4.1's index query) does; with a bit-exact test (op_tests/test_dsv4_rotate_quant.py --round-rope). * feat(flydsl): ragged query rows in the row-group FP4 MQA logits kernel flydsl_pa_mqa_logits_fp4_rowgroup takes query_start_loc [B + 1] and max_query_len: sequence b's rows are the flattened rows query_start_loc[b] .. query_start_loc[b + 1] - 1, any count up to max_query_len (0 included, as cudagraph padding has). q_fp4 is [rows, 1, H, D / 2]; q_scale, weights, row_ends and out go by flattened row. Without row_ends, row n of a count-row sequence sees context - count + n + 1 keys. A sequence's rows still share each key load, so a batch whose requests carry different numbers of query rows no longer has to be scored one row a sequence. Rows past a sequence's count read row 0 and get bound 0: no keys walked, no stores. One build flag (ragged); counts and batch stay runtime values. The rectangular [B, NEXT_N] call is unchanged. V4.1 full layers (page 8, 200K context, T = 6..768, cold, MI355X): the ragged call runs 1.0-2.35x faster than scoring each row as its own sequence (T = 48: 24.8 vs 43.4 us, T = 768: 246 vs 580 us), and within -4.5..+1.6% of the rectangular call; uniform counts give the rectangular call's logits bit for bit. * feat(flydsl): ragged row-group MQA logits take their lengths from row_ends Ragged rows (query_start_loc) now require row_ends and take no context_lens: a sequence's length, which the kernel shares its workgroups out by, is its last row's bound, read in the kernel. A caller with per-row bounds no longer builds a per-sequence length. pages_per_block: a block-table entry may name that many consecutive pages (its first page's id / pages_per_block), expanded in the kernel where the page is used, so a caller whose pages are a multiple of the kernel's page size (DeepSeek-V4.1: 16 pages of 8 rows a 128-row PAGE) passes its own table as is. flydsl_pa_mqa_logits_fp4 forwards row_ends, query_start_loc, max_query_len and pages_per_block to the row-group kernel (page 8 / 64) and refuses them on its other pages. Uniform per-row bounds and pages_per_block = 1 build the same kernel as before; op_tests/test_flydsl_pa_mqa_logits_fp4.py passes at page 64 and 8. * feat(configs): DeepSeek-V4-Pro TP8 wo_b in the dsv4 mxscale bmm table N=7168, K=2048 (128x128 blocks), M 1..65536: the attention output projection at TP8, which had no row and took the heuristic kernel. Rows tuned on gfx950 as the table's others were; errRatio <= 0.0009 against the dequantized reference. Against the heuristic pick, same operands, cold: 1.01-1.03x at most M, 1.22x at M = 256. Only rows added. * feat(flydsl): M-tile FP4 MQA row-group kernel with a shape-only launch plan The row-group kernel scores MQA as a GEMM: M a sequence's query rows times the heads, cut into tiles of whole rows (tile_m, a row's head sum stays in its tile), N its keys. Prefill and decode are the same ragged rows (query_start_loc, row_ends); next_n is gone from the kernel. Work is shared out over M tiles by real length inside the kernel, so a long sequence's tiles get more workgroups. - make_fp4_mqa_plan / Fp4MqaPlan: the launch (tile from TILE_ROWS, rows a wave, warps, workgroups, depth, nontemporal loads) from the batch's shape alone, so a captured graph keeps it; lru_cached. - flydsl_pa_mqa_logits_fp4_rowgroup(plan, ...) takes ragged rows only. - Page 8 compiles a single depth-2 walk: the two walks in one kernel spilled 40 VGPRs at 6 rows of 32 heads; one walk is as fast or faster on every shape measured (B=64 6 rows at 131K keys 185 -> 157 us, per-row tables T=3072 954 -> 838 us). Page 64 keeps both walks. - Launch through a cached CompiledFunction (tensor_shim._run_compiled): host cost per call 87 -> 16 us (FlyDSL rebuilt the cache key from every argument on each call). - flydsl_pa_mqa_logits_fp4: ragged rows (query_start_loc + row_ends, an optional plan) run the row-group kernel; a rectangular call runs the varctx kernel as before rowgroup existed, no silent conversion. The page-8 op test passes its rows as ragged rows. Prefill (V4.1 geometry) 1.9-13x faster than scoring each row as its own sequence; decode on par with the previous row-group kernel; outputs bit-exact against it. * update * style: black-format the row-group FP4 MQA kernel
…ROCm#6163) The multi-GPU job has timed out at its 120-minute limit on every recent main run. Most multi-GPU tests start a new Pool and run init_distributed_environment for every case. Each init takes 3-30+ s, while the collective being tested takes microseconds. On MI350X, test_fused_ar_rms_per_group_quant (80 inits) and test_custom_allreduce (31 inits) alone took 60-90 % of the budget. Because of the timeouts, the tests that run last in alphabetical order never ran at all. Each worker now inits once, loops over every case, and tears down once. Per-case results come back to the parent as small scalars. Inputs are generated on the GPU from a per-case seed, and the reference is computed on the rank itself, so no large tensors are pickled or reduced on the CPU. Workers never raise inside the loop, so all ranks stay aligned on every collective; the parent does the assertions. CUDA graphs captured for each case are kept alive until teardown. Custom all-reduce caches a peer's IPC address under the local buffer pointer, so freeing a graph mid-sweep and getting the same address in a later capture would replay with stale peer pointers. CLI options, per-case log lines, and summary columns are unchanged. Separate launches remain only where a setting must be fixed before init: - tp size and topology - AITER_QUICK_REDUCE_QUANTIZATION - AITER_AR_1STAGE - the AITER_CUSTOM_AR_MAX_SIZE=0 fallback path Number of inits with CI default arguments: test_fused_ar_rms_per_group_quant 80 -> 1 test_custom_allreduce (sweep) 31 -> 1 gfx1250_poc/{allgather,p2p_bw,sync} 97 -> 3 gfx1250_poc/allreduce 20 -> 2 test_allgather 20 -> 1 test_fused_ar_mhc_post_only 18 -> 1 test_fused_ar_rms_mxfp4_quant 9 -> 3 test_reduce_scatter 8 -> 2 test_custom_group 8 -> 4 test_quick_all_reduce 7 -> 3 test_fused_qknorm_ar 6 -> 1 test_quick_all_reduce_rmsnorm had no __main__ block, so running it with `python3` in CI executed nothing. It now calls pytest.main from __main__. test_init_dist_env used a fixed env:// port (49374). In the first CI run that reached this test, the port was already in use (EADDRINUSE), so the test sat for 60 minutes while the peer ranks waited out the TCPStore connect timeout. It now uses a free port with a tcp:// init method. test_collective_profile's spawn path had the same problem with its fixed 29500 fallback; it now picks a free port before spawning.
* [Triton/Gluon] [gfx942] Enable sparse_mla_fwd on gfx942
sparse_mla_fwd was gated to gfx950. The Gluon kernel itself needs no change to
run on gfx942 -- the gl.amd.cdna4 intrinsics it uses lower fine there -- so the
gate was the only thing in the way, plus a tile size chosen against gfx950's
LDS budget.
gfx942 has 64 KB of LDS per workgroup against gfx950's 160 KB, so BLOCK_K=64
asks for 68608 B and fails to launch:
OutOfResources: shared memory, Required: 68608, Hardware limit: 65536
BLOCK_K=32 is the largest tile that fits. num_warps is pinned at 4 rather than
following block_k // 16, so the LDS cap does not halve it as a side effect
rather than as a tuning decision; that is worth 1.19x -> 1.55x on the prefill
shapes.
dot_precision="fp8" is rejected on gfx942. It feeds the cache's own code points
to the matrix core, which needs OCP e4m3 MFMA; CDNA3 has fp8 MFMA in the fnuz
encoding only, so those code points come out wrong. It was silently wrong
(rel-err 7.5e-1) rather than faulting, so this is now an explicit error naming
the alternative. dot_precision="bf16" dequantizes the tile into LDS and works
on every arch and every cache format, fp8 caches included.
The bench's fp8-dot series is built from FP8_DOT_ARCHS for the same reason:
with the arch gate widened it would otherwise reach that series on gfx942 and
raise on its first point, and triton's perf_report does not catch it, so the
run died before printing a table.
The test's cache builders go back to an explicit OCP e4m3 rather than the
arch-native fp8 dtype. These records are OCP by definition of the format, and
the kernel's dequant reads OCP code points; on gfx942 the native dtype is
float8_e4m3fnuz, so building with it decoded to garbage (rel-err 1.0) on every
fp8-cache case. _classify_flat already rejects fnuz when a uint8 view is not
hiding the dtype.
gfx950 is unchanged: _arch_block_k -> 64, _arch_num_warps -> 4, lds_limited
False, FP8_DOT_ARCHS contains gfx950, and FP8_DTYPE is OCP e4m3 there either
way, so every value and branch resolves exactly as before.
Validated on MI325X (gfx942), TP4, GLM-5.3-Flash at 131k context, where this
path serves the rope-free NoPE MLA that has no AITER kernel today:
op_tests test_sparse_mla 32 passed, 12 skipped (fp8 dots)
bench_sparse_mla runs clean, bf16 series only
sparse-MLA decode bucket 3.099 -> 0.400 ms/step 7.74x
sparse-MLA prefill 10988 -> 10330 us/call 1.06x
end-to-end sweep +7.14% mean output tok/s, 7/7 rows won
gsm8k 0.9719 vs 0.9712 strict-match (neutral)
* [Triton/Gluon] [gfx942] Bound the sparse_mla geometry against gfx942's LDS
_check_geometry has never had an upper bound on kv_lora_rank, and gfx942 already
takes the smaller of the two tiles this wrapper selects, so a latent too wide to
fit has nowhere left to go. It surfaced as Triton's OutOfResources at launch
rather than as something a caller can act on.
Measured on MI325X by reading the required byte count back out of
OutOfResources, the gfx942 footprint is exactly
BLOCK_K * (kv_lora_rank + 8) * 2
+ BLOCK_K * (qk_rope_head_dim + 8) * 2 (separate rope only)
+ 32 * BLOCK_K
on every point at BLOCK_K 32 and 64: 67072 B at kv_lora_rank=1024 rope-free,
71680 with a 64-wide rope, 132608 at 2048. Those constants hold only for bf16
tiles with the async path off, which is every gfx942 launch, since fp8 dots are
rejected there and lds_limited forces ASYNC_LDS off.
The check raises a ValueError naming the geometry, its footprint and the budget.
gfx950 is not in the table, so it is left to the launcher exactly as before.
test_sparse_mla 32 passed, 12 skipped (unchanged)
guard boundary 512 and 512+rope still launch; 1008, 1024 and 2048 are
rejected, and the computed figures match the measured
Required: byte counts exactly
* [Triton/Gluon] [gfx942] Cover the fp8-dot arch gate
The gate had no test on any arch, not just gfx942: the shape matrix skips its
fp8 cases off gfx950, so the branch that rejects fp8 dots was never reached and
the branch that accepts them only incidentally.
_resolve_dot_precision's arch parameter no longer defaults. That default was the
real hazard behind the missing coverage: there is one call site, and dropping
its third argument would have silently resolved every launch as gfx950 and
re-enabled fp8 dots on gfx942 -- which a unit test on the helper cannot see,
since the helper still behaves correctly when handed an explicit arch. Required,
that same edit is a TypeError caught by 32 of the 34 tests; verified by making
it and running the suite.
test_dot_precision_arch_gate parametrizes over SUPPORTED_ARCHS, so it needs no
GPU, never skips, and covers a newly added arch automatically.
test_sparse_mla 34 passed, 12 skipped (was 32 passed, 12 skipped)
* [Triton/Gluon] [gfx942] Reject the packed caches outside gfx950
fp8_dsv4_mla, fp8_g64 and the SWA+top-k two-loop return through _forward_paged
before the kernel's own arch gate, so widening that gate never reached them.
They route to pa_decode_sparse, whose packed driver is gated on
DEVICE_ARCH == "gfx950"; anywhere else they land in its fallback path, which
reads a plain grouped fp8 pool rather than these records and rejects them on
dtype:
fp8_g64 AssertionError: kv_scales supplied but unified_kv is
torch.uint8, expected torch.float8_e4m3fnuz
fp8_dsv4_mla RuntimeError: unified_kv dtype mismatch: kv=torch.uint8,
q=torch.bfloat16
Neither names the arch, so on gfx942 both read as a caller mistake. The
ordering is pre-existing, since the early return precedes the assert on main
too, but listing gfx942 in SUPPORTED_ARCHS advertises formats that have no
implementation there, so the wrapper now says so itself.
Reaching the fallback path instead is not an option: it wants the arch-native
fnuz dtype and these records carry OCP code points, so it would decode them
wrong in exactly the way dot_precision="fp8" does.
Verified on gfx942 through the public wrapper:
bf16, fp8_scalar, fp8_dsv32_mla launch, unchanged
fp8_g64, fp8_dsv4_mla ValueError naming the arch
test_sparse_mla 36 passed, 12 skipped (was 34 and 12)
gfx950 is in PACKED_ARCHS, so every branch resolves there as before.
* [Triton/Gluon] [gfx942] Cover the sparse_mla LDS budget guard
The GPU suite only launches the default 512/64 geometry, so a regression
in _check_lds_budget would restore OutOfResources while tests stayed
green. CPU-only cases pin the measured rope-free and separated-rope
boundaries, including the exact 64 KB point.
Signed-off-by: Jin Tao <jin.tao@amd.com>
* [Triton/Gluon] [gfx942] Reject fp8 q and caches on gfx942
The kernel reads every fp8 byte -- q, the cache, and the operands of the fp8
dots -- as OCP e4m3, which is gfx950's native fp8. On gfx942 the native fp8
is fnuz: vLLM's fp8 KV cache (current_platform.fp8_dtype()) and aiter's
quantizers (dtypes.fp8) both write float8_e4m3fnuz there. The encodings share
a bit layout with exponent bias 7 against 8, so a gfx942 cache read as OCP
comes out 2x too large, and fnuz's saturated +-240 (0x7F/0xFF) are NaN in OCP.
The fp8-cache tests passed on gfx942 only because their fixtures had been
switched to OCP to match the kernel, which no gfx942 producer writes. A cache
arrives as bytes, usually behind a uint8 view, so the wrapper cannot tell the
encoding from the tensor; _classify_flat catches a fnuz-typed flat pool and
nothing else.
So gfx942 now takes bf16 q and a bf16 cache only, and says why. That is what
the GLM-5.3-Flash path this PR enables uses: vLLM keeps that model's MLA
cache in bf16 and dispatches only bf16 q/kv here. FP8_DOT_ARCHS becomes
FP8_ARCHS and covers fp8 q and caches as well as the dots. The test fixtures
go back to the arch-native fp8, and the fp8 cases skip on gfx942 instead of
running on data nothing there produces. Reading fnuz natively, which would
also let the fp8 dots use CDNA3's fnuz MFMA directly, is left to a follow-up.
The >2 GB global-load test gains a bf16 case, so gfx942 keeps that path
covered now that the fp8 case skips there.
Verified on MI325X (gfx942):
test_sparse_mla 32 passed, 29 skipped (was 44 passed, 12 skipped;
every skip is an fp8 case)
native fnuz cache rejected through sparse_mla_fwd as fp8_scalar and
fp8_dsv32_mla; with the gate stubbed out the same call
returns rel-err 1.01 and the wrapper test fails
bench_sparse_mla runs, bf16 series only
gfx950 is in FP8_ARCHS, so every branch resolves there as before.
Signed-off-by: Jin Tao <jin.tao@amd.com>
* [Triton/Gluon] [gfx942] Take the sparse_mla LDS budget from arch_info
_check_lds_budget kept its own copy of gfx942's LDS capacity. It now reads
arch_info._LDS_CAP_BYTES, as pa_decode_sparse and the GEMM num_stages picker
do, so the capacity has one definition.
_LDS_CAP_BYTES lists gfx950 as well, so dict membership can no longer pick
the checked arch, and the gfx942-only gate is now explicit. The footprint
model was measured for gfx942's bf16, non-async tiles; gfx950 stays with the
launcher as before, which test_lds_budget_gfx950_is_unchecked pins.
test_sparse_mla 32 passed, 29 skipped (unchanged)
Signed-off-by: Jin Tao <jin.tao@amd.com>
* [Triton/Gluon] [gfx942] Pick the sparse_mla launch config by arch
_async_launch_config took an lds_limited flag that sparse_mla_fwd derived
from the arch. It now checks arch_info.get_arch() itself, so the arch
switch is visible where the config is chosen, and its signature is back to
main's.
No launch changes: lds_limited was exactly arch == "gfx942".
test_sparse_mla 32 passed, 29 skipped (unchanged)
Signed-off-by: Jin Tao <jin.tao@amd.com>
* [Triton/Gluon] [gfx942] Reword the fp8-dot arch error
"{arch}'s decodes fnuz" read as malformed. It now says the arch's native
fp8 is fnuz, the wording _check_fp8_arch already uses.
test_sparse_mla 32 passed, 29 skipped (unchanged)
Signed-off-by: Jin Tao <jin.tao@amd.com>
* [Triton/Gluon] [gfx942] Say which caches each dot_precision takes per arch
The dot_precision docstring still said bf16 dots work with every cache
format, and did not say fp8 dots are gfx950-only. On gfx942 the wrapper
takes a bf16 cache alone and rejects the rest, so the public docstring now
states the same scope the code enforces.
Docstring only.
Signed-off-by: Jin Tao <jin.tao@amd.com>
* [Triton/Gluon] [gfx942] Sort the sparse_mla test and bench imports
This PR's import edits in test_sparse_mla.py and bench_sparse_mla.py left
both blocks unsorted, which the pinned ruff 0.16.0 in pre-checks reports as
I001. Both files are clean at the merge base. Import order only.
ruff 0.16.0 check all three sparse_mla files pass
test_sparse_mla 32 passed, 29 skipped (unchanged)
bench_sparse_mla runs, bf16 series only
Signed-off-by: Jin Tao <jin.tao@amd.com>
* [Triton/Gluon] [gfx942] Check the sparse_mla LDS budget at the launch's KV pad
_check_lds_budget modelled the KV tile at the kernel's default 8-element row
pad, but a prefill launch (num_queries >= _PREFILL_MIN_ROWS, bf16 dots)
passes KV_LDS_PAD=16, so the guard undercounted every prefill launch by
BLOCK_K * 8 * 2 = 512 B. The pad is now computed once and handed to both
the guard and the launch, so the two cannot drift apart again.
No power-of-two geometry changes its accept/reject result, so no launch
that used to work is rejected. What changes is that the guard now reports
the footprint the launch actually needs, and it no longer depends on the
gap between pow-2 widths to stay safe. Widths in between never reach the
launcher: the KV tile's PaddedSharedLayout asserts a pow-2 shape.
Measured on MI325X with the guard stubbed out, reading metadata.shared when
the kernel fits and OutOfResources' Required: when it does not, at BLOCK_K
32, decode (8 queries) and prefill (2048):
kv_lora_rank / rope decode prefill
512 / 0 34304 34816 fit
512 / 64 38912 39424 fit
512 / 256 51200 51712 fit
512 / 512 67584 68096 rejected
1024 / 0 67072 67584 rejected
1024 / 64 71680 72192 rejected
2048 / 0 132608 133120 rejected
The guard now reproduces every one of those through sparse_mla_fwd. The
boundary test gains the prefill pad, with exact-fit points at 1000 (decode)
and 992 (prefill). Dropping the pad from the guard again fails the five
prefill rejects.
test_sparse_mla 43 passed, 29 skipped (was 32 and 29; the new ones are
boundary cases)
Signed-off-by: Jin Tao <jin.tao@amd.com>
* [Triton/Gluon] [gfx942] Read the sparse_mla tile and warps from configs/
_ARCH_BLOCK_K and _ARCH_NUM_WARPS were per-arch launch values in Python,
which the config rules in aiter-ops-triton.instructions.md and
configs/CLAUDE.md place in JSON. They now live in
configs/gfx942/gluon/attention/sparse_mla/DEFAULT.json BLOCK_K 32, num_warps 4
configs/gfx950/gluon/attention/sparse_mla/DEFAULT.json BLOCK_K 64, num_warps 4
read through resolve_config_dir() + load_config_json(), the single
DEFAULT.json read that configs/CLAUDE.md describes for attention wrappers.
The gfx950 file carries main's hardcoded block_k = 64 and
num_warps = block_k // 16, so nothing changes there. The entry is keyed by
the kernel, _sparse_mla, so the reduce can get its own entry later.
Only the values this PR made per-arch move. The rest of the launch policy
(the BLOCK_M rule, split-K, the async tile choice, the prefill KV pad) is
main's and stays as it was. Moving it all would rework gfx950's launch path,
which is its own change.
test_launch_config_published checks from any machine that every arch in
SUPPORTED_ARCHS ships the file, so adding an arch without one fails on
every runner rather than only on that arch's. The LDS boundary test takes
gfx942's BLOCK_K from the same file.
Verified on MI325X (gfx942): the compile-time launch arguments the wrapper
passes (BLOCK_K, num_warps, waves_per_eu, ASYNC_LDS, KV_LDS_PAD and the
rest) are identical before and after over 96 shapes spanning decode and
prefill, H 8/16/64, rope 0/64, has_invalid and forced split-K.
_get_config() costs 0.65 us per call.
test_sparse_mla 45 passed, 29 skipped (was 43 and 29; +2 config checks)
Signed-off-by: Jin Tao <jin.tao@amd.com>
* [Triton/Gluon] [gfx942] Leave gfx942's waves_per_eu to main's rule
_async_launch_config's gfx942 branch returned early with its own copy of
main's waves_per_eu rule, so gfx942 carried a second, arch-specific
statement of a launch value. All the branch has to do is keep the async
tiles off, which are sized for gfx950's LDS, so the arch check now sits in
the condition that enables them, and gfx942 falls through to main's rule
for everything else. The switch still reads arch_info.get_arch() inside
_async_launch_config, and the signature is still main's.
The async path also needs fp8 dots, which gfx942 rejects before it gets
here, so today the check only takes effect once fnuz fp8 dots exist there.
It stays for that case.
No launch changes. Checked on MI325X:
_async_launch_config old and new agree on all 10752 combinations of
its inputs under both gfx942 and gfx950, fp8_dots
included
launch arguments identical to the PR head over the same 96 shapes
test_sparse_mla 45 passed, 29 skipped (unchanged)
Signed-off-by: Jin Tao <jin.tao@amd.com>
…able (ROCm#6174) Running DeepSeek-V4.1-Flash at TP4 with DSpark (the ATOM accuracy job) logs "not found tuned config" for six bf16-in/bf16-out GEMM (N, K) pairs that have no gfx950 row anywhere in the bf16_tuned_gemm family: (32, 5120), (128, 512), (128, 5120), (384, 5120), (512, 5120) and the TP4 lm_head (32320, 5120). Each pair is tuned at the 25 M values the dsv41 table already uses (1..16384). That covers every power of two, which is where the lookup's gl=1 padding lands, so every M seen in CI resolves to a row. It also includes 6/12/24/48/96/192/288/384/768/1536, which are the DSpark verify batches (requests x 6) and now hit a row exactly. Backends tried: asm, opus, flydsl, triton, skinny, torch. err_ratio <= 0.031 for every row. Speedup over the torch default, median/max: (32, 5120) 2.13x / 2.29x (128, 512) 2.11x / 2.53x (128, 5120) 2.07x / 2.81x (384, 5120) 1.72x / 2.41x (512, 5120) 1.79x / 2.40x (32320, 5120) 1.02x / 1.19x (weight-bandwidth bound; torch stays best at 8 of the 25 M values) Validation: - test_config_shape_collision passes. - --run_config on all 150 rows: every row reports OK.
* gfx950 tuning bf16 tp8 shapes * Drop small-win Kimi K3 bf16 configs
…5909) * [Triton] Gluon PA decode: fix VGPR blow-up on Triton 3.8 Carry the prefetched KV tiles unconditionally in the PS decode kernels. Triton 3.8 if-converts the conditional carry into selects that keep both key tiles live (~260 VGPRs, 1 wave per SIMD). Add ps_decode_performance_test() to reproduce (not run in CI). * Apply review comments
* Run AITER full releases weekly at 1 a.m. Central * Update actionlint for timezone schedule support
* Fix MXFP4 and MXFP8 fallback for older Triton * Remove redundant quantization regression test * Cache scaled downcast capability and handle missing Gluon
* [Triton/Gluon] Add fp32 reference and tests for fused Qwen3-Next GDN decode Defines the op contract for a fused Qwen3-Next Gated DeltaNet decode kernel before any kernel lands: packed qkvz/ba split, depthwise causal conv1d with bias + SiLU, GDN gating, the fp32 recurrent state update, gated RMSNorm with a SiLU gate, and an optional per-head group FP8 quantization epilogue. This is a sibling of fused_kda_decode, not a variant of it. KDA uses a per-channel lower-bounded sigmoid gate, a sigmoid output gate, uniform head counts and no conv bias; Qwen3-Next GDN uses a per-head exp(-exp(A_log) * softplus(a + dt_bias)) decay, a SiLU output gate, num_v_heads == 2 * num_k_heads and a biased conv. The head topology alone (lp = H*K offsets) cannot express VH == 2*KH, so a mode flag on the existing kernel is not viable. The reference has been cross-checked on MI355X against all seven Artemis Gluon variants across all nine specialized batch sizes (M = 1,2,4,8,16,32,64,128,256). conv_state matches bit-exactly on every band; the other outputs agree within the ATOL/RTOL 0.05 bar this directory already uses for fused_kda_decode, with the residual being a few bf16 ULP from differing fp32 evaluation orders. The kernel-side tests are written against the pad_slot_id convention that aiter's other GDN ops already honour (fused_conv_recurrent_norm skips state_idx < 0; causal_conv1d_update_split_qkv takes an explicit pad_slot_id), rather than the reserved-slot-0 redirection the out-of-tree Artemis adapter maintains via a monkeypatched _replay_metadata. They skip until the op lands. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com> * [Triton/Gluon] Add fused Qwen3-Next GDN decode op for gfx950 One Gluon launch replaces the four-kernel Qwen3-Next decode chain (fused_qkvzba_split_reshape_cat -> causal_conv1d_update -> fused_recurrent_gated_delta_rule_packed_decode -> layer_norm_gated) and can optionally emit the per-head group-128 FP8 activations that a block-FP8 out_proj consumes directly, removing a separate quantization launch. A sibling of fused_kda_decode, not a variant of it: KDA uses a per-channel lower-bounded sigmoid gate, a sigmoid output gate, uniform head counts and no conv bias, and its lp = H*K offsets structurally cannot express VH == 2*KH. Verified bit-identical to the upstream Artemis kernel on every output (out, quantized, scales, ssm_state, conv_state) at M = 1, 32 and 256 on MI355X, and green against the fp32 reference across M = 1..128. Two changes on top of the upstream source: * pad_slot_id: rows whose state index is the sentinel are skipped without reading or writing state, matching the contract SGLang's replay_state_indices_validator documents and the one aiter's fused_conv_recurrent_norm already implements. This removes the upstream kernel's dependence on a caller-maintained reserved slot 0. * HAS_FP8: the quantization epilogue is optional, because an unquantized out_proj cannot accept an (fp8, scales) tuple. Support is reported by fused_gdn_decode_qkvz_supported() as (bool, reason) rather than asserted, so callers fall back and can log why. Coverage is deliberately narrow and stated: gfx950, VH == 2*KH, head dims 128, conv width 4. The Gluon body is generated by build_aiter_kernel.py, which pins the source sha256 and asserts every edit it applies; the source is ~600 lines of dense Gluon with explicit register layouts, where a transcription slip would surface only as a numerical mismatch. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com> * [Triton/Gluon] Drop the kernel generator from the aiter tree build_aiter_kernel.py reads the Artemis kernel pack, which is not part of this repository, so shipping it here would put a script in the tree that cannot run. Provenance stays where a reviewer needs it -- the generated kernel's header records the source file and its sha256 -- and the generator itself lives with the pack it reads. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com> * [Triton/Gluon] Add benchmark for fused_gdn_decode_qkvz The maintainer guide asks for kernel + wrapper + test + benchmark; this is the missing fourth. Reports achieved HBM bandwidth against arch peak. The kernel performs no MFMA and no gl.dot -- it is bound by the FP32 recurrent state, read and written once per token per value head, which dominates every other tensor by an order of magnitude. So bytes moved is the right roofline denominator. Two measurement choices the script documents inline, because getting either wrong silently invalidates the result: it times under CUDA graph replay, since timing the Python call measures host overhead that at small batches is ~2.5x the kernel itself; and it sizes the state pool independently of the batch, since sizing it to the batch puts small batches entirely in cache. Measured on MI355X: the kernel is latency-bound at small batch (1.4% of peak at batch 1) and becomes bandwidth-bound from batch 64, saturating at 5.0-5.2 TB/s, 62-65% of the 8 TB/s peak. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com> * [Triton/Gluon] Label the benchmark's timing as per-call The header said 'us'. The timed region is a CUDA graph enclosing the whole wrapper, so the number is per call. This op always issues a single launch, so the two coincide here -- but the label should not depend on that staying true, and the Artemis sources it derives from do use multi-launch paths at small batch where per-call and per-launch differ. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com> * [Triton/Gluon] Add the large-batch band to fused_gdn_decode_qkvz Above batch 128 the op now uses a second decomposition that keeps the whole [V, K] state tile in registers across both matrix-vector products, instead of the latency-hiding path that serves smaller batches. This is not a tuning of the existing kernel and could not be reached by one. The two move identical bytes -- one read and one write of the fp32 state tile, confirmed by source-level count -- and three candidate explanations for the gap were tested and rejected: launch-order locality (rewriting the grid to the other kernel's order gained 0.75%), store cache policy (.cs to .wt was slightly worse), and routing large batches through the kernel that does carry a batch ladder (about 5% of the gap). What remains is the register layout and reduction structure, which is a different algorithm. At batch 256 this takes the op from 53.5 to 46.5 us, and from 65% to 74.5% of peak HBM bandwidth with the same byte count. At 128 the two are within 2%; below 128 the existing path is better by up to 15%, so the threshold sits at the measured crossover. Selection is a range check rather than a shape table, so every batch size above the threshold is covered -- there are no gaps to fall into. Correctness is verified against the fp32 reference at 128, 192 and 256, and at the non-power-of-two sizes 136, 160, 200 and 248. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com> * [Triton/Gluon] Address review: predicate coverage, test gaps, conventions Review findings from ROCm#5520, verified against the code before fixing. The most serious was a hole in the tests: the correctness case computed ref_quant and never compared it, so the FP8 output was checked only through its scales. A broken epilogue could pass every case as long as the bf16 output, scales and state updates were right. Now asserted, in value space rather than code space -- a 1-ULP e4m3 code difference near the top of the range is enormous as a raw integer and meaningless as a value, so both sides are dequantized with their own scale. The support predicate claimed more than it checked: - conv_state and ssm_state are indexed with the same slot, but their slot counts were never compared, so a shorter conv pool passed and was then read out of bounds - only the two pools were required contiguous, while projected_qkvz, projected_ba, conv_weight, conv_bias, A_log, dt_bias and norm_weight are all addressed with flat pointer arithmetic -- a same-shaped strided view passed and was read from the wrong locations - conv_bias was missing from the device check, and A_log, dt_bias and norm_weight were not parameters at all, so the predicate could return (True, "") for a launch whose operands were on the CPU or wrongly shaped Also: the bf16-only path is now checked numerically rather than only for dtype and None, since HAS_FP8=False is a separate compiled specialization; a broken op module now fails instead of silently skipping every kernel test, with the skip narrowed to the optional Gluon dependency; the kernel module no longer imports torch, so it stays importable without PyTorch as the convention requires; and the three launchable kernels carry config-aware repr= callbacks. Dropped float8_e4m3fnuz rather than testing it: that is the gfx942 format and this op is gfx950-only, so advertising it promised a path that can never run. PAD_SLOT_ID is deliberately excluded from the repr keys. It is a protocol sentinel rather than a tuning key, and make_kernel_repr does not sanitize the sign of a negative value -- including it emits an invalid function identifier and fails at compile. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com> * [Triton/Gluon] Fix style CI: format generated output, modernize annotations The style checks were red on files I had not linted the same way CI does. Three distinct causes: The generated Gluon module was never formatted. It is copied from sources that are not black-formatted, so black rejected it. The generator now formats its own output rather than leaving it to be hand-patched. My local ruff runs used --select E4,E7,E9,F, which is narrower than the rule set the repository's pyproject actually enforces, so C408, ISC004, RUF059, UP006, UP035, UP045 and I001 all went unseen. Fixed at source: dict() calls rewritten as literals, implicit string concatenation inside return tuples parenthesised, typing.Optional/Tuple replaced with PEP 604/585 spellings, imports sorted, and the unread second element of the inline-asm pairs renamed to _scratch. Also removes build_aiter_kernel.py, which had been copied to the repository root by mistake and was never meant to ship here. Verified with the versions CI pins -- black 26.5.1 and ruff 0.16.0 -- run from a tree carrying this repository's pyproject, over every file the branch touches. 26 op tests still pass. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com> * [Triton/Gluon] Sort imports as the repository config does Ruff's I001 fired on three files. My earlier local run missed it because I linted copies flattened into a scratch directory, where the aiter package path no longer matched, so isort classified aiter imports as third-party and sorted them into a different, self-consistent order. Running ruff 0.16.0 against the real paths with the repository's own pyproject reproduces CI exactly. Worth noting for anyone reproducing: the ruff step uses reviewdog with filter-mode=diff_context and -fail-on-error=true, so any finding on a changed line fails the check -- whole-repo cleanliness is not the bar, the diff is. Verified with the pinned versions, black 26.5.1 and ruff 0.16.0, on the real tree. 26 op tests still pass. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com> * [Triton/Gluon] Move the op test into the gated_delta_net category Mirrors the public wrapper's category, as requested in review. Adds __init__.py to match the sibling test directories. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com> * [Triton/Gluon] Route the large band through the shared device helpers Addresses the review asks to stop duplicating operations inside the kernel file and to use the repository's activation helpers. _decode_group now calls _head_dynamics and _quantization_scale instead of re-implementing the gate, softplus, exp, sigmoid and FP8 scale inline, and _sigmoid delegates to _triton_kernels/activation._sigmoid_exp2. Outputs are bit-identical to the previous kernel on all five results -- bf16 output, quantized, scales, and both mutated state pools -- at batch 1, 2, 8, 32, 64, 128, 192 and 256. Kernel time is unchanged: 10.47 -> 10.53 us at batch 32, 27.03 -> 26.76 at 128, 46.58 -> 46.45 at 256. _normalize_qk is left in place. Unlike the others it is not a call-site swap: it takes a shared-memory handle and reads through .gather() with a hardcoded 128 offset, where the large band uses .slice().load() with an explicit gl.SliceLayout tied to its register layout. The arithmetic is the same either way, so routing through it would change layouts in a kernel that is bandwidth-bound above batch 128 without changing results. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com> * [Triton/Gluon] Make the decode op import-safe and gate on Triton >= 3.8 The wrapper imported the Gluon tile kernels at module scope, so the module raised ImportError on any build without triton.experimental.gluon -- before fused_gdn_decode_qkvz_supported() could report the reason and let the caller fall back, which is the one case the probe exists for. Import the three kernels at their launch sites instead, and defer arch_info for the same reason: detection can touch the driver at import time. Mirrors fused_gdn_prefill_qkvz.py. Add _gluon_supported(), an explicit Triton >= 3.8 floor checked before the import probe. Some ROCm 3.7 builds carry an earlier, incompatible Gluon where the import succeeds but the layout ops these kernels use miscompile, so probing the import alone is not enough. Trim the batch ladder from ten sizes to one per dispatch branch (1, 32, 64, 127, 128, 256), adding 127 to cover the previously untested 65-127 path. The branch mapping is recorded above the list. The Gluon kernel is unchanged and outputs stay bit-identical across batches 1-256, so the reported accuracy and end-to-end numbers are unaffected. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com> --------- Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
…ue head size (ROCm#3698) * [Triton] unified_attention: mask V load and output store by value head size unified_attention sizes everything off HEAD_SIZE = q.shape[-1] and reuses one dim_mask (offs_d < HEAD_SIZE) for the Q load, the K load, the V load and the output store. That holds when qk_head_dim == v_head_dim, but models with asymmetric heads (e.g. MiMo with qk=192, v=128) have a V cache and an output whose per-head width is the value head size, not HEAD_SIZE. With the single 192-wide mask the V load reads past the 128-wide head into the next kv-head (and past the buffer on the last head), and the output store writes 192 dims into a row whose per-head stride is 128, clobbering the next head and running off the end on the last head. On MiMo decode this shows up as garbled logits at low concurrency and a hipErrorIllegalAddress under load. Prefill is unaffected. Mask the V load by the V cache head stride (stride_v_cache_2) and the output store by the output per-head stride (output_stride_1), both already passed to the kernels. For symmetric models HEAD_SIZE, stride_v_cache_2 and output_stride_1 are all equal, so this is a no-op. No kernel signature or launcher change. Signed-off-by: adlashab <adlashab@amd.com> * [Triton] unified_attention: use explicit V_HEAD_SIZE for V load and output store The first version of this PR masked the V load and the output store with the V cache head stride (stride_v_cache_2) and the output per-head stride (output_stride_1). That works for the sglang layout but it leans on qk_head_size > v_head_size: when the value head is wider than the query head the V tile would only be loaded partially. Replace the stride trick with an explicit V_HEAD_SIZE / V_HEAD_SIZE_PADDED constexpr pair, mirroring HEAD_SIZE / HEAD_SIZE_PADDED. The V load, the acc tile, the segment buffer and the output store now size off the value head, while Q and K keep using HEAD_SIZE, so both orderings (v < qk and v > qk) are handled. This includes the V load in the SPLIT_UNMASKED_LOOP prefill loop of the 2d kernel. The triton subwrappers pass V_HEAD_SIZE from _UAParams.head_size_v, which already falls back to head_size for the shuffled layout, so for symmetric heads V_HEAD_SIZE == HEAD_SIZE and the generated kernel is identical to before. The gfx950 / gfx1250 gluon kernels are not touched, asymmetric heads are already routed away from them. Add asymmetric-head cases to test_unified_attention.py: the MiMo case (qk=192, v=128), a symmetric case, and the reverse ordering (qk=128, v=192), across decode and prefill shapes, bf16 and fp8, both backends. Signed-off-by: adlashab <adlashab@amd.com> * [Triton] unified_attention: keep asymmetric heads off the gfx1250 gluon 2d and reduce kernels The gfx1250 gluon 2d kernel and the gluon reducer still size V and the output off HEAD_SIZE, but is_2d_gluon_available and is_reduce_gluon_available only checked head_size_v == head_size on gfx950. Add the same check to both gfx1250 gates so asymmetric heads fall back to the triton kernels there. The gfx1250 3d path only runs with the shuffled cache, where head_size_v is head_size, so it needs no change. Also add V_HEAD_SIZE (and V_HEAD_SIZE_PADDED for the 2d kernel, which already lists HEAD_SIZE_PADDED) to the kernel reprs so specializations that differ only in the value head get different names. Signed-off-by: adlashab <adlashab@amd.com> * [Triton] unified_attention: guard asymmetric heads, pick config by the wider head Three follow-ups for asymmetric qk/v heads: - deny pre-shuffled KV cache with asymmetric heads for now. head_size_v falls back to head_size for the shuffled layout (its last dim is the vector width, not the head size) and the shuffled kernels lay V out with HEAD_SIZE, so an asymmetric shuffled cache would silently give wrong output. Compare the per-page element count of k and v instead. - assert out.shape[-1] == head_size_v. The kernels store with the V stride and mask, so an out sized by the qk head would get a partial write. - look the config up with max(head_size, head_size_v). There are no asymmetric configs yet, so pick by the wider head, the one that sets the accumulator and V tile size. Adds a small test for the two asserts. Signed-off-by: adlashab <adlashab@amd.com> --------- Signed-off-by: adlashab <adlashab@amd.com>
Signed-off-by: Martin Domanský <ragua@email.cz>
…gfx1250 (ROCm#4894) * Merge branch 'ROCm:main' into quant_mxfp4 * Update bench_quant_mxfp4.py * Resolve formatting issues * Resolve formatting issues * Resolve formatting issues * Resolve formatting issues * Resolve typing deprecation and import formatting * Fix Ruff import and __all__ ordering * Replace deprecated typing.Tuple in quant module * Use PEP 604 unions for optional tensor annotations * Include instruction for mxfp4 qunatization. * Change hueristic Added pre fetch slight modification in bench script * Consolidate comments and fix formatting/duplicate imports * Modifying shapes to be tested * Remove the inline asm part as the test triton patch provided by compiler team, now successfully uses the correct instruction * Removing unwanted files * Removing json file * Migrate gfx950 Gluon mxfp4/mxfp8 quant tuning to JSON configs Also update mxfp4/mxfp8 quant test files with representative shapes covering the benchmark shape sweep. * COnfig file format change * Add MXFP4/MXFP8 gfx950/gfx1250 gluon quant kernels (squashed) * style: black-format bench_quant_mxfp4_fp8.py * refactor(quant): switch gfx1250 mxfp4/mxfp8 block-config to JSON rule trees Replace the hardcoded if/else tuning chains and the flat bucket-key config lookup with configs/gfx1250/gluon/quant/{quant_mxfp4,quant_mxfp8}/DEFAULT.json default+rules trees, resolved via config_utils.select_tuned_config. Drops the now-unused quant_config_utils module. * fix(quant): pin mxfp4 gfx1250 config lookup to arch=gfx1250 _mxfp4_gfx1250_config's launch config is also reused, unchanged, by the plain-Triton fallback kernel on non-gfx1250 archs. Without an explicit arch override, resolve_config_dir defaulted to the running arch, so the config lookup 404'd on every non-gfx1250 arch (e.g. gfx950), breaking every dynamic_mxfp4_quant call there -- including the gfx950 stochastic-rounding tests and downcast_to_mxfp_pow2_scale_fp4. * fix(quant): remove dead/broken leftover code in _mxfp4_pack_op from main merge The merge of main introduced main's own refactor splitting _mxfp4_quant_op into scale computation + a new _mxfp4_pack_op(qx, ...) helper that takes already-scaled qx. The merge left dead leftover lines inside _mxfp4_pack_op that referenced an undefined 'x' (not a parameter of this function) and recomputed an unused scale, causing a NameError at trace time (flagged by CI as ruff F821/RUF059). * perf(quant): narrow gfx1250 mxfp8 small-M tile width for K > 1024 For M <= 32, the launch config previously picked BLOCK_SIZE_N purely from a fixed element budget over BLOCK_SIZE_M, independent of K. Benchmark sweeps (repeated trials, not single samples) show that for K > 1024 a narrower 128-wide tile robustly beats the wider budget-derived tile -- more independent CTAs hides TDM latency better than one bigger transfer per CTA at this size. ~11-13% bandwidth improvement at M=8, K=3072/7168 (closes the largest lag vs the Triton fallback at these shapes). * refactor(quant): table-based mxfp8 small-M BLOCK_SIZE_N + tighten comments Replace inline if/else BLOCK_SIZE_N arithmetic for mxfp8's M<=32 branch with an explicit _MXFP8_SMALL_M_BLOCK_SIZE_N lookup table + helper. Also shorten a duplicated multi-line comment in the gluon kernel and condense select_tuned_config's docstring. * Address PR review comments: bucket-key config lookup, kernel repr, lazy Gluon import, config dir rename, trim redundant test shapes - Replace the rule-tree select_tuned_config with lookup_tuned_config, a flattened bucket-key resolver ({schema, "M_GEQ_33.N_LEQ_16384": {...}, "any": {...}}) matching unified_attention_utils's existing pattern. - Add make_kernel_repr-based repr= to both gfx950 Gluon quant kernels. - Move the gfx950 Gluon kernel imports in quant.py out of the module header into the arch-gated dispatch branches that use them. - Rename configs/gfx950/gluon/quant/{mxfp4,mxfp8}_quant -> {mxfp4,mxfp8} to match the sibling PR's naming convention. - Drop one of three redundant large-non-square shapes from the mxfp4/ mxfp8 quant parametrize lists. * fix(quant): address gfx950-PR review feedback applicable to gfx1250 - Replace select_tuned_config's invented default+rules JSON schema with the repo's established composite bucket-key format (config_utils. lookup_config), matching the pattern used by unified_attention/gemm/ conv config families. Flattens mxfp4/mxfp8 DEFAULT.json into complete, mutually exclusive buckets instead of stacked rules. - Add repr= (make_kernel_repr) to the two launched gfx1250 gluon kernels so tuned specializations are distinguishable in traces/artifacts. - Make the gluon kernel import lazy, gated behind the gfx1250 arch check, matching gemm_a16w16.py's pattern -- importing aiter.ops.triton. quant.quant no longer unconditionally pulls in triton.experimental.gluon. - Rename configs/gfx1250/gluon/quant/{quant_mxfp4,quant_mxfp8} to {mxfp4_quant,mxfp8_quant} to align directory naming with the gfx950 sibling PR (ROCm#4383). * docs(quant): shorten config lookup docstrings and comments * docs(quant): shorten config lookup and gfx950 kernel comments * revert(quant): restore _triton_kernels/quant/quant.py to main (drop unneeded fp32-cast removal) * Update quant.py * Move gfx950 Gluon quant shape-derived block sizes into JSON and retune mid-M configs - Drop the Python BLOCK_SIZE_M/N overrides in dynamic_mxfp4_quant and dynamic_mxfp8_quant; launch configs now come only from the JSON tables. - Tune the 33<=M<=1024 buckets (more, smaller CTAs) for MXFP4 and MXFP8. - Compute e8m0 scale bias in int32 before narrowing to uint8. * Use lookup_config for gfx950 Gluon quant configs to match the sibling PR Replace lookup_tuned_config with lookup_config(table, axes, **values) and drop the "schema" key from the gfx950 mxfp4/mxfp8 JSON tables. * perf(quant): retune gfx1250 gluon mxfp4/mxfp8 configs; JSON is gluon-only - Re-tune per-bucket block/warp/buffer configs via the dispatch path with cold inputs (old config kept as a candidate); mxfp4 +10.6%, mxfp8 +5.4% geomean vs previous gluon configs. - Move the former in-code shape rules (small M / small N tile sizes) into the JSON buckets so the JSON is the single source of block sizes. - mxfp4: JSON is used only by the gfx1250 gluon launch; the Triton fallback is main's original heuristics again. - mxfp4: fix compile failures for M<=8 at N>1024 and M<=4 (wide BLOCK_SIZE_N at tiny BLOCK_SIZE_M) by using BLOCK_SIZE_N=128 tiles there. - Shorten comments/docstrings across the touched files. * bench_quant_mxfp4_fp8: match the gfx950 branch (drop duplicate line_vals, original comment) * gfx1250 quant: shared e8m0 helper, tile helpers, drop dead params, waves_per_eu in JSON, backend tests * gfx1250 quant: drop unreachable table entries, restore async_store note, add mxfp8 bucket shape * gfx1250 quant: int64 mxfp8 strides, overlap mxfp4 stores, infer table axes in get_quant_config * gfx1250 quant: retune mxfp4 bucket table (38 -> 30 keys) --------- Co-authored-by: Satya Nikhil Kodukula <nikhil.kodukula@gmail.com>
Actions/stale, daily: a pull request is labelled `stale` after 15 days without activity; any push or comment clears it. Nothing is closed -- -1, since leaving it unset would mean the action's 7-day default. Drafts and PRs labelled keep-open, blocked, do-not-close or dependencies are exempt. Issues are out of scope. Pinned to the v10.4.0 commit. Updated exempt PR labels to only include 'do-not-close'.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Kimi-K3 KDA decode currently uses separate operations. This PR adds a fused FlyDSL path on AMD gfx950 GPUs to reduce kernel launches and global-memory traffic.
Technical Details
flydsl_kda_decode,flydsl_kda_decode_with_f_b, andis_flydsl_kda_decode_supportedfromaiter.ops.flydsl.kda_decodeand the package.Test Plan
Covers reference correctness at batch sizes 1/8/16, cache behavior, input validation, and public API exports. Tests skip on unsupported hardware or when FlyDSL is unavailable.
Test Result
Reported on MI355X (gfx950):
Performance
On MI355X (gfx950), the new FlyDSL APIs were about 2× faster than the parent commit's fused Triton decode across batches 1–128:
Compared
80a3b0b09with301f09d83using the TP8 layout (12 heads, dimension 128). Results are medians of seven alternating A/B rounds with warmed GPU graph replay. The projection baseline uses a separate BF16 linear operation before Triton decode.All benchmark output/state checks passed (relative RMSE < 1%); convolution caches matched exactly. These are kernel timings, not serving throughput. Results were collected at
301f09d83, not the newer PR head; callers must opt into the new APIs.Submission Checklist