[gfx1201] add tuned ck_gemm_a8w8_blockscale configs for various qwen3 models and default case - #1
Closed
big-yellow-duck wants to merge 1 commit into
Closed
[gfx1201] add tuned ck_gemm_a8w8_blockscale configs for various qwen3 models and default case#1big-yellow-duck wants to merge 1 commit into
big-yellow-duck wants to merge 1 commit into
Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
big-yellow-duck
pushed a commit
that referenced
this pull request
Jul 7, 2026
* [fused_qk_norm_rope] add 1way kernel; replace ds_bpermute with DPP+ds_swizzle in rms-reduce/NEOX
New kernel `fused_qk_norm_rope_1way` mirrors the existing 2way kernel for
single-token-stream models: per-head RMSNorm followed by RoPE on q/k,
supporting NEOX (half-head split) and interleaved (adjacent-pair) styles,
head_size ∈ {64, 128, 256}, BF16/FP16. No partial rotary, no KV cache,
no quantization — minimal scope of 2way.
Shared rope_common.h shuffle helpers reworked to take advantage of DPP and
ds_swizzle on gfx9xx. The previous implementation went through __shfl_xor
with width=32 on a 64-lane wave, which lowers to ds_bpermute_b32 (~10 cyc)
even for compile-time-constant XOR offsets:
block_utils::warp_reduce_sum<float>:
before: 5x ds_bpermute_b32 (offsets 16,8,4,2,1) + 1x __shfl broadcast
after: 3x ds_swizzle_b32 (offsets 16,8,4 via XOR mask) +
2x v_mov_b32_dpp (offsets 2,1 via quad_perm 0x4e/0xb1)
The XOR butterfly is symmetric so the post-reduce broadcast is a no-op
and is removed. Order kept as 16->1 to make the FP32 accumulation
bitwise-identical to the previous bpermute-based path.
warp_shfl_xor_sync_vec<T, vec_size, XorOffset>: new helper for vectorised
constant XOR shuffles, lowers to ds_swizzle. Used at the two NEOX
neighbour-swap call sites in fused_rope_rms_1way_kernel — replaces the
runtime `lane + neighbor_offset` arithmetic that lowered to ds_bpermute.
Verification on production config B=2 T=4096 D=128 NEOX, Hq=24 Hk=25 (BF16):
- Bitwise identical output vs the old ds_bpermute helper across NEOX +
interleaved, BF16 + FP16, D ∈ {64, 128, 256}, T ∈ {32, 1024, 4096}.
- Latency: ~8% reduction (652 GB/s achieved, up from 600 GB/s).
- PMC: SQ_WAIT_INST_ANY -25%, SQ_BUSY_CYCLES -11%, SQ_INSTS_VALU -10%.
- Disasm: 0 ds_bpermute_b32 (was 10), 7 ds_swizzle_b32, 2 v_mov_b32_dpp.
Also fixes a pre-existing OOB read on the interleaved RoPE path: cos/sin
are only VEC_SIZE/2 elements per lane but the vec_t::load issued a full
VEC_SIZE read, racing past the cos_sin buffer tail on the last token. The
1way kernel now uses scalar loads of exactly VEC_SIZE/2; the 2way kernel
still has the original load (separate fix recommended).
Tests: op_tests/test_fused_qk_norm_rope_cache_quant.py adds a
`test_qk_norm_rope_1way` + sweep across dtype/D/T/interleaved, all pass
checkAllclose(rtol=1e-2, atol=0.05) vs the torch reference.
* [fused_qk_norm_rope] add quad fast path to 1way kernel
Adds a 4-head-group ("quad") fast path to the 1way fused QK-Norm+RoPE kernel
for shapes where num_heads_q % 4 == 0 && num_heads_k % 4 == 0. Each physical
wave (64 lanes) maps to 4 heads x 16 lanes-per-head, packing 4 adjacent heads
into one wave instead of the 1-head-per-wave default path. Renames the internal
"pair-x2" variant to "quad" for clarity; the 2-head "pair" variant did not
pay off vs quad and is dropped.
* [fused_qk_norm_rope] quad kernel: kill NEOX divergent branch, use packed bf16 cvt
Two bit-exact-equivalent optimizations to fused_rope_rms_1way_quad_kernel
identified from rocprofv3 ATT trace + disasm inspection.
1. NEOX rope: replace the divergent if(is_lower_half){...}else{...} with a
per-lane cndmask select. The divergent branch forced the compiler to
emit two copies of the rope math AND the bf16 cvt sequence, with
s_and_saveexec / s_xor / s_or EXEC mask switches between them. The new
form computes both x*c - nx*nc and x*nc + nx*c in FP32 on every lane
(same op order as the original divergent code) and selects via cndmask.
2. Add f32x2_to_bf16x2_rne + pack_f32_to_vec_t<T,N> helpers in
rope_common.h, adapted from the gfx94 RNE branch of float_2_bf16_pair<0>
in aiter/csrc/kernels/mla/hk/hk_mla_buffer_managers.cuh. The helper
replaces the compiler default static_cast<bf16>(float) expansion (13
instructions plus EXEC mask switches per output) with a 10-instruction
VALU sequence with no EXEC mask manipulation. Used for both RMSNorm
output and rope output in the NEOX and INTR paths. RNE rounding is
bit-identical to the ctor for non-NaN inputs (NaN is replaced with
canonical 0x7FFF, which is unreachable from finite RMSNorm / rope
inputs).
Verified bit-exact equivalent to pre-change kernel: stashed both files,
built HEAD (53f3f65), dumped raw output bytes for 6 (seed, mode, shape)
configs (NEOX/INTR x T={127,1024,8192} x (Hq,Hk) in {(16,16),(24,24),
(32,32)}); restored optimizations, rebuilt, dumped again; md5sum compare
matched on all 12 binary files (~150M bf16 elements total).
Disasm impact for quad<bf16, 128, NEOX, 6, 6> on gfx942:
total disasm lines: 1212 -> 591 (-51%)
default static_cast<bf16> seqs: 48 -> 0
s_and_saveexec EXEC mask switches: 47 -> 2 (-96%)
Wall-clock impact at B=1, T=8192, Hq=Hk=24, D=128 (Qwen-Image-2 shape):
NEOX: 224.77 -> 141.05 us (1.59x speedup, -83.7 us)
INTR: 152.07 -> 123.07 us (1.24x speedup, -29.0 us)
* [fused_qk_norm_rope] fallback 1way kernel: inline RMSNorm + packed bf16 cvt
Apply the same kill-divergent-branch + packed-bf16-cvt optimizations that
landed for the quad fast path (e66f419) to the single-head-per-warp
fused_rope_rms_1way_kernel fallback (used when num_heads_q or num_heads_k
is not divisible by 4 — e.g. Qwen-Image-2 prefill, Hq=24, Hk=25).
Three changes, all bit-exact-equivalent:
1. Inline RMSNorm instead of the shared mrope_utils::warp_rms_norm_<T,N>
helper. The shared helper does
acc = sum( static_cast<float>(input[i])^2 ) ; reduce
input[i] = static_cast<T>(static_cast<float>(input[i]) * s_val
* static_cast<float>(gamma[i]))
which forces the compiler to re-cast bf16->f32 in the writeback
(redundant v_lshlrev_b32 conversions) and emit the default 13-instr
static_cast<bf16>(float) sequence per output. Inlining lets us cache
the f32 reads in a stack array v[VEC_SIZE], do the writeback in f32,
and pack via pack_f32_to_vec_t (10 instr per bf16x2 pair). RNE
rounding bit-identical for finite inputs (NaN -> canonical 0x7fff).
Per the user's directive, the shared helper is left untouched (it is
used by 4+ unrelated kernels) — the inlining is local to the 1way
fallback.
2. NEOX rope: replace the divergent if(is_lower_half){...}else{...} with
a per-lane cndmask select, same as e66f419 for the quad kernel.
Both expressions are evaluated in the SAME FP32 op order as the
original divergent code (mul + mul + sub for lower, mul + mul + add
for upper) — bit-exact equivalent.
3. Stage both NEOX and INTR rope outputs in float[VEC_SIZE] then pack
via pack_f32_to_vec_t for the same bf16-cvt-instr-count win.
Verified numerically against the pure-FP32 PyTorch reference at rtol=1e-2
atol=0.05 (matches aiter checkAllclose) on shapes that exercise the
fallback path (Hq,Hk in {(24,25),(25,25),(18,18),(20,21)}, T in
{256,1024,4096}, NEOX+INTR): all 24 cells pass with fail=0/N.
Wall-clock impact (rocprofv3 kernel-trace, B=1, D=128, median of 30 iters,
gfx942 / MI300X). Geomean over 32 (T, Hq, Hk, layout) cells: 1.22x.
Hot shape (Qwen-Image-2 prefill, Hq=24, Hk=25):
T=256 NEOX: 13.02 -> 10.15 us (1.28x)
T=256 INTR: 10.38 -> 9.33 us (1.11x)
T=4096 NEOX: 156.11 -> 116.17 us (1.34x)
T=4096 INTR: 119.98 -> 107.40 us (1.12x)
T=8192 NEOX: 305.22 -> 225.58 us (1.35x)
T=8192 INTR: 231.51 -> 206.33 us (1.12x)
NEOX consistently wins more than INTR (~1.30x vs ~1.10x): NEOX gets all
three optimizations, INTR gets only inline-rmsnorm + packed-rope (no
divergent branch to refactor). VGPR (NEOX/INTR) goes 24/20 -> 32/24 — no
spill, occupancy unchanged.
* [fused_qk_norm_rope] fix host-pass build of rope_common.h includers
bb22fd3 added DPP-based warp_reduce_sum / half_warp_reduce_sum /
warp_shfl_xor_sync_vec helpers in mrope_utils::block_utils that referenced
opus::mov_dpp / opus::number / opus::bool_constant directly. But
rope_common.h gates `#include "opus/opus.hpp"` with #ifdef
__HIP_DEVICE_COMPILE__ — so during the HIP host pass, `opus` is not a
declared namespace. clang HIP parses template bodies in BOTH passes for
non-dependent-name lookup, so any TU that includes rope_common.h without
otherwise pulling in opus.hpp transitively (e.g. via quant_utils.cuh)
fails the host pass with "error: use of undeclared identifier 'opus'".
This caused CI build failure on
csrc/kernels/rope/general_2c_cached_positions_offsets_fwd_kernels.cu,
which is why the failure was masked locally — the fused_qk_norm_rope JIT
target pulls quant_utils.cuh (-> opus.hpp) and was building fine.
Two changes, both bit-for-bit equivalent at the GPU instruction level:
1. warp_reduce_sum / half_warp_reduce_sum bodies are now wrapped in
#ifdef __HIP_DEVICE_COMPILE__ (matching the existing pattern used
throughout rope_common.h, e.g. line 564). The host pass now skips the
opus::mov_dpp calls entirely and the function returns `val` unchanged
— fine because these helpers are __device__-only and never called
from host code. Device-pass body is unchanged.
2. warp_shfl_xor_sync_vec's tag-dispatch parameter went from
`opus::number<XorOffset> = {}` to `std::integral_constant<int,
XorOffset> = {}`. The signature is parsed in both passes regardless
of #ifdef wrapping the body, so this one CAN'T be hidden the way #1
is. opus::number<I> is publicly derived from
std::integral_constant<index_t, I> (csrc/include/opus/opus.hpp:57), so
existing callers passing opus::number<X>{} continue to work
unchanged via pass-by-value slicing of the empty derived type.
Verified:
- general_2c_cached_positions_offsets_fwd_kernels.cu now builds (full JIT
re-run succeeds; previously failed in host pass on line 7260).
- fused_rope_rms_1way_quad and fused_rope_rms_1way (fallback) both
numerically match pure-FP32 PyTorch reference at rtol=1e-2 atol=0.05
for shapes that exercise both warp_reduce_sum and warp_shfl_xor_sync_vec
(NEOX path) and the fallback inline RMSNorm.
- rocprofv3 fallback sweep over 32 (T, Hq, Hk, NEOX/INTR) cells: geomean
speedup vs pre-388f737ba baseline = 1.22x (identical to the measurement
before this fix — opus::mov_dpp lowers to the same v_*_dpp instruction
as __builtin_amdgcn_mov_dpp; opus.hpp:1559-1562 is a thin wrapper).
* [fused_qk_norm_rope] quad kernel: scalarize Q/K split branch when QUAD_*_CT is even
When both QUAD_Q_CT and QUAD_K_CT are even, `is_q = global_warp_id <
T*QUAD_Q_CT` is uniform across the full 64-lane physical wave. Readfirstlane
the warp_id so the compiler emits s_cmp + s_cbranch instead of the divergent
v_cmp + s_and_saveexec + s_xor + s_cbranch_execz sequence, saving a few cycles
and EXEC-mask thrash per wave. Falls back to the original per-lane path for
odd QUAD_*_CT (e.g. H=12 -> QUAD=3) where the boundary can cut a wave.
* fix: align 1way RMSNorm cast order with diffusers
* fix: keep cos_sin in fp32 in 1way fused QK norm + RoPE
* style: black reformat assert in 1way op_tests
---------
Co-authored-by: LiuYinfeng01 <yinfeliu@amd.com>
BadrBasowid
pushed a commit
that referenced
this pull request
Aug 19, 2026
* [FLYDSL] add MLA decode reduce kernel for gfx942 Add FlyDSL MLA reduce kernel and production opt-in fallback, with multi-token decode support (decode_qlen > 1), num_kv_splits dispatch compat, and fp8 partial/output dtype options. Include correctness tests against the HIP kernel and standalone bench/profile harness scripts. Co-authored-by: Cursor <cursoragent@cursor.com> * Update op_tests/prof_mla_reduce.py Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * [FLYDSL] drop redundant Claude-generated comments from MLA reduce Co-Authored-By: Claude <noreply@anthropic.com> * style: black format MLA reduce FlyDSL files Fix Check Code Style with Black CI failure on PR ROCm#3901. Co-authored-by: Cursor <cursoragent@cursor.com> * [FLYDSL] replace tier string literals with Tier enum in mla_reduce.py Replace bare string tier values ("simple", "m64", "m256", "mlds") with a proper Tier(enum.Enum) class. select_tier() now returns Tier; the compile_mla_reduce() tier param is typed Tier = Tier.SIMPLE; all string comparisons use enum members; the LDS global_sym_name f-string uses tier.value to preserve the original naming. Co-Authored-By: Claude <noreply@anthropic.com> * style: black 26 tuple unpack in bench_mla_reduce_standalone Match psf/black@stable (Black 26) used by CI Checks workflow. Co-authored-by: Cursor <cursoragent@cursor.com> * [FLYDSL] use flydsl.utils.env OptBool for MLA reduce opt-in gate Replace the raw os.environ.get("AITER_MLA_REDUCE_FLYDSL") read with FlyDSL's typed env helper (OptBool on an EnvManager subclass). The helper import is kept inside the existing try/except alongside the FlyDSL availability check, so the gate still falls back silently to HIP when FlyDSL is not installed. Drops the now-unused `import os`. Co-Authored-By: Claude <noreply@anthropic.com> * [FLYDSL] consolidate mla_reduce expr imports and relocate HIP bench Use fx.* qualified flydsl.expr access per review feedback and move the standalone MLA reduce benchmark under op_benchmarks/hip/. Co-authored-by: Cursor <cursoragent@cursor.com> * chore: drop dev-only prof_mla_reduce.py harness Remove the standalone rocprofv3 driver; it was not used in CI or production testing and is superseded by the bench/correctness harnesses. Co-authored-by: Cursor <cursoragent@cursor.com> * [FLYDSL] black-format MLA reduce and trim review docstrings Remove unverified HBM-bound and dev-doc references from MLA reduce docstrings and apply Black wrapping in mla_reduce.py. Co-authored-by: Cursor <cursoragent@cursor.com> * Refactor MLA reduce test harness into pytest and FlyDSL bench script. Split shared helpers into flydsl_mla_reduce_common, convert correctness coverage to pytest (HIP + torch ref for GLM Dv=256), and add a dedicated FlyDSL bandwidth benchmark alongside the existing HIP bench. Co-authored-by: Cursor <cursoragent@cursor.com> * Add MLA reduce serving guards and discriminating differential tests. Harden the FlyDSL kernel with gather/store bounds checks and a store q-range guard, plus in-process guards-on/off tests that prove the guards matter via mapped-allocation slack fixtures rather than only asserting passes-with-fix. Co-authored-by: Cursor <cursoragent@cursor.com> * [FLYDSL] opt1: persistent grid-stride launch for MLA reduce Wrap the per-work-item body in process_work_item() and add a kn_mla_reduce_v1_ps -style 1-D persistent grid (num_cu*OCC*2 blocks) that grid-strides over the flat work index with a CSR-sentinel early-out. Host (mla_reduce_kernels.py) and the bench/test harness auto-select persistent via should_use_persistent_launch when H*NTG*num_reduce_tile exceeds the HIP threshold. Guard semantics (bounds-checked gather, q-range clamp, disable_guards) preserved. Bench (MI300X, H=16, Dv=512): sparse 16384-tile/8-active 127.8us->12.3us (~10x); uniform tiles=8/splits=32 7.3->6.8us; tiles=256/splits=8 flat (below threshold). pytest: 40 passed + 8 slow/serving (sparse grid) passed. Co-authored-by: Cursor <cursoragent@cursor.com> * [FLYDSL] opt2: vectorize MLA reduce epilogue store Truncate each accumulator element to out_t before packing into the out_t vector so the epilogue lowers to buffer_store_dwordx2/4 instead of a scalarized per-element store. Drops the unused f32 acc_vt staging vector. pytest: 40 passed (matrix, not-slow). Bench within noise of opt1 at the prod replay shapes (store path already saturated post-opt1). Co-authored-by: Cursor <cursoragent@cursor.com> * [FLYDSL] opt3: replace _as_list with tuple loop-carried state Use a tuple init and range_constexpr (compile-time) indexing in the massive-path accumulator loop, removing _as_list and its runtime range(n) index workaround that could lower to i64 index arithmetic. VEC==1 scalar unwrap after the yield is handled explicitly. pytest: 40 passed (matrix, not-slow). Bench flat at prod shapes (bandwidth -bound kernel; this is a codegen cleanup). Co-authored-by: Cursor <cursoragent@cursor.com> * [FLYDSL] opt4: wire MLA reduce waves_per_eu (default 4) Stamp rocdl.waves_per_eu on the emitted gpu.func at compile time and expose AITER_MLA_REDUCE_WAVES_PER_EU for sweeps. Default lowered 8->4 (opt4 sweep: on H=16 Dv=512 tiles=8 splits=32, wpe=4 best at 6.6us; H=128 graph time flat across 2/4/6/8). Host wrapper and bench/test harness pass the env-resolved value. pytest: 40 passed (matrix, not-slow). Co-authored-by: Cursor <cursoragent@cursor.com> * [FLYDSL] opt5: convert MLA reduce seq loop to FlyDSL range Replace the raw scf.ForOp seq loop with FlyDSL range(seq0, ub_seq, ntg, init=None) so the AST rewriter emits scf_range without iter_args, letting hot_loop_scheduler interleave the inner split-loop VMEM loads with compute. Body unchanged; seq -> seq_i32 induction value. pytest: 40 passed (matrix, not-slow). Perf-neutral on the prod shapes by CUDA-graph replay (the bandwidth-bound reduce is already at 24-72% BW); kept for codegen hygiene / scheduler enablement. Co-authored-by: Cursor <cursoragent@cursor.com> * [FLYDSL] fix massive-tier os prefetch with deferred bounds guard Defer the pmap bounds select to point-of-use in emit_massive_body so the prefetched os[s+1] load stays in flight (vmcnt(1)) instead of draining every iteration. Add load_split_o_raw and carry a float OOB mask folded into the LDS scale at the FMA. Wire Tier.ALL as the production compile path (device-side runtime tier selection per tile, mirrors HIP) and update tests/benchmarks accordingly. * [FLYDSL] stage reduce_partial_map to LDS once per work item Cooperative-copy pmap[t0:t1] into lds_pmap before the split loop so gather_row reads LDS instead of repeated global pmap loads (mirrors reduce.cu:431-438). Adds pmap to the shared allocator for all tiers. Co-authored-by: Cursor <cursoragent@cursor.com> * mla_reduce: depth-2 double-rate pipeline in emit_massive_body Process 2 splits/iter with two output buffer_loads in flight (mirror HIP reduce_output_massive oaccu_0/oaccu_1). Carrying two distinct loaded vectors (os0/os1) plus the next pair's prefetch as separate SSA yields lets the compiler allocate separate VGPRs, so the accumulate loop reaches s_waitcnt vmcnt(1) overlap instead of the depth-1 vmcnt(0) drain (register aliasing). Split index is clamped for OOB prefetch (lds_pmap[0]/lds_scale[0] always written) with a deferred float mask -> no NaN*0 pollution. Graph mode: b8_s32 9.0->7.0us (ratio 1.50->1.17x), b1_s128 25.8->16.9us (ratio 1.86->1.22x). SIMPLE unchanged. vgpr 42->56. * mla_reduce: generalize accumulate to GRP=8 double-rate pipeline Generalize the depth-2 double-rate loop to process GRP splits/iter with GRP output buffer_loads in flight (loop-carried grouped state). GRP=8 is the gfx942 sweet spot: vgpr ~125 keeps 2 waves/SIMD while the accumulate loop reaches s_waitcnt vmcnt(7) overlap (was vmcnt(1) at depth-2). GRP=16 pushes vgpr ~199 (1 wave/SIMD) for no b8_s32 gain and only marginal b1_s128 gain. Graph mode: b8_s32 7.0->6.5 (ratio 1.17->1.08x), b1_s128 16.9->15.3 (ratio 1.22->1.10x) vs depth-2. SIMPLE unchanged. buffer_load_dwordx4 14->50, vmcnt_max 3->7. * mla_reduce: vectorize group LDS reads (ds_read_b128) in emit_massive_body Collapse the per-split scalar lds_pmap/lds_scale gathers into one wide STensor vector load per GRP group. The group's pmap indices and lse scales live at contiguous LDS slots (base = i*GRP, 16B-aligned), so read them with ds_read_b128 instead of GRP scalar ds_read_b32. ATT showed LDS/SMEM-wait (lgkmcnt) as the #1 stall category once the GRP=8 pipeline hid VMEM; this cuts ds_read_b32 66->10 (ds_read_b128 10->24), vgpr 125->123. Refactor gather_row -> row_from_pmap(pmap_value, local_seq) so scalar and vectorized paths share the identical bounds clamp. OOB tail lanes of the vector read hit stale slots: substitute slot-0's pmap value (always staged; massive body only runs when n_splits>1) so the row computation never sees a stale value even with guards disabled, and select-force the scale to 0 for invalid splits (no stale NaN reaches the FMA, and selecting on the LDS scale does not touch the VMEM os load so the deferred-guard vmcnt overlap holds). Graph-mode vs HIP (gfx942, GLM-5.2 serving): b1_s128 15.3->13.1us (1.09x -> 0.95x, BEATS HIP); b8_s32 6.5->6.0us (1.08x -> 1.00x, parity). SIMPLE tier unchanged. Correctness 44/44 (matrix + differential + graph). * mla_reduce: tier-dependent GRP=16 for M256/MLDS (-8% b1_s128 graph, 0.95x->0.87x) The long-loop M256/MLDS accumulate path (nlse>=4) uses a GRP=16 double-rate software pipeline (16 output buffer_loads in flight, deeper vmcnt overlap over the many-iteration loop): b1_s128 @128 splits drops 13.1->12.0us graph (0.87x HIP), reproduced across two runs. The M64 path (nlse=1) keeps GRP=8: a blanket GRP=16 regressed the low-split tail (b8_s5/s6 @5-6 splits) +11% because each group wastes 10 masked lanes instead of 2, with no compensating gain on b8_s32 (within run noise). Scoping GRP to the long-loop tiers captures the b1_s128 win with zero tail regression. 44/44 correctness matrix + differential + cudagraph replay pass. * mla_reduce: hoist group-0 os prefetch ahead of the LSE scale barrier (lever #6) emit_massive_body's group-0 output loads depend only on lds_pmap (staged + barriered at the top of the work item), not on lds_scale (written by the warp0 LSE reduce). Splitting load_group into a pmap/os phase (load_os_group) and a scale phase (load_scales) lets the GRP os buffer_loads issue *before* the scale barrier, overlapping ~GRP VMEM loads with the warp0 LSE reduce + barrier wait. ISA confirms it: the group-0 buffer_load_dwordx4 batch is now emitted before s_barrier, with the lds_scale ds_read_b128 after it; vmcnt_max=14 overlap and the uniform SGPR gather descriptor (no waterfall) are preserved. Graph mode: b8_s32 6.1->6.0, b8_s26 6.1->6.0, b1_s128 12.0->11.9 (reproduced across two runs, zero regressions). 44/44 correctness matrix + differential + cudagraph replay pass. * mla_reduce: scalar indptr sentinel via raw pointer deref (-0.2us all active shapes graph, b8_s32 1.00x->0.98x) Load the CSR traversal sentinels (`last`, per-work-item `tile_start`) with a raw uniform `llvm.load` + `rocdl.readfirstlane` instead of the GTensor `buffer_load` path. A `buffer_load` is inherently a vector memory op (voffset addressing) and never lowers to `s_load_dword`, so simply wrapping the GTensor load in readfirstlane (tried, inert) leaves the vector load + `s_waitcnt lgkmcnt(0)` traversal-floor stall in place. Dereferencing the uniform address raw makes the load scalarizable, mirroring HIP `__builtin_amdgcn_readfirstlane(p_reduce_indptr[tile])` (reduce.cu:688). Graph-mode (clean back-to-back A/B): every M64/SIMPLE shape drops ~0.2us (b8_s32 6.1->5.9 = beats HIP 6.0; b8_s6 4.5->4.25; b8_s3 4.4->4.2) with b1_s128 held at 11.9 (0.86x). 44/44 correctness (matrix + differential + graph replay). Kept t0/t1 inside process_work_item on the GTensor path: scalarizing those regressed b1_s128 +1.2us (they feed the M256 accumulate-loop bounds). * mla_reduce: invariant sentinel load -> scalar s_load_dword (-0.1..0.2us all shapes graph) Mark the traversal indptr sentinel llvm.load invariant. reduce_indptr is read-only in the kernel and the tile index is block-uniform, so the AMDGPU backend now scalarizes the uniform-address load into s_load_dword (SMEM) instead of a per-lane global_load_dword + s_waitcnt vmcnt(0). This matches HIP kn_mla_reduce_v1_ps's scalar sentinel and removes the #1 traversal stall (re-profile: line 316 was 42.5K vmcnt, 35% of total). ISA: global_load 1->0, s_load_dword 2->3, vmcnt(0) 15->14. Graph A/B (same session, fresh JIT): b8_s32 5.95->5.9 (0.98x), b8_s13 4.8->4.7, b8_s6/s5 4.25->4.1, b8_s3 4.3->4.1, b8_s2 4.2->4.0; b1_s128 held 11.9 (0.86x). 44/44 correctness (matrix + differential + graph replay). * mla_reduce: persistent grid = num_cu (T2, SIMPLE 1.08x->0.97x, beats HIP; dense unchanged) The persistent launch used grid = num_cu*OCC*2 (=4864), mirroring HIP `num_cu*kOccupancy*2`. But the FlyDSL Tier.ALL kernel runs at occupancy 1 wave/SIMD (193 VGPR from the shared massive accumulate path), so that 16x grid is ~8x oversubscribed on the sparse serving profile: thousands of blocks each do a single sentinel `s_load_dword` then terminate, and at occupancy 1 that latency cannot be hidden. Dropping the grid to num_cu (mult=1) trims the wasted blocks. The grid-stride loop still covers any input (correctness unchanged); genuine dense MLDS work is bandwidth-bound and unaffected. Graph-mode (two runs each, clean A/B vs the mult=16 baseline this session): - b8_s3 4.1->3.7 (1.08x->0.97x) and b8_s2 4.1->3.4 (1.11x->0.92x) -- SIMPLE now BEATS HIP (3.8/3.7), the skill's gated SIMPLE beat target. - b8_s26 5.9->5.7, b8_s32 5.9->5.8, b1_s128 11.9->11.8 -- held/marginally better. - b8_s13/s6/s5 flat at 4.7/4.1/4.1 -- the M64 mid-tail occupancy floor (emit_massive VGPR), out of this skill's scope. - Dense uniform (256 tiles x 304 splits, MLDS): 722.9us -> 722.9us, no regression. - mult=4 regressed b8_s5 (4.1->4.4); mult=1 is the sweep optimum. Overridable via MLA_PS_GRID_MULT. 44/44 correctness (matrix + differential + graph replay). * mla_reduce: joint-config NUM_THREADS=256 (VEC 4->2, VGPR 193->133) Joint-search knob NUM_THREADS (MLA_NUM_THREADS) set to 256 as the new default. Halving VEC = Dv/NUM_THREADS (4->2) cuts the per-thread output accumulator live-set: whole-kernel VGPR 193->133 and the M256/MLDS accumulate throughput improves. Sweep (graph, optionb vs 39dd385): b1_s128 11.8->11.0 (0.86x->0.80x), mid-tail b8_s32/s26/s13/s6/s5 and SIMPLE b8_s3/s2 unchanged, dense uniform 722.9->712.8us. 44/44 correctness (matrix + differential + CUDA-graph replay). No regressions; hard gates (b1_s128<=0.90x, SIMPLE, dense) all satisfied. Also plumbs env-overridable joint-search knobs for the search: MLA_GRP_M256 / MLA_GRP_M64 (accumulate GRP per tier) and MLA_M64_HI_THR / MLA_M64_HI_GRP (runtime M64 sub-split), defaults unchanged. * mla_reduce: capture-safe host per-tier dispatch via num_kv_splits (opt-in) Realizes the per-tier occupancy win (opt5: M64 alone occ-3, ~0.2-0.3us faster on the mid-tail; b8_s6/s5 cross below HIP) capture-safely, WITHOUT CUDA-graph conditional nodes (infeasible on ROCm 7.2.4). The wrapper now picks the tier on the HOST from num_kv_splits (a pure host scalar upper bound on per-tile n_splits; no device read/sync) when AITER_MLA_REDUCE_HOST_TIER=1. Default stays Tier.ALL (unchanged). select_tier is monotonic and each per-tier body reduces a tile's actual n_splits (tier only caps LSE-register width), so select_tier(num_kv_splits) is correct for all tiles. Because num_kv_splits is constant for a fixed CUDA-graph capture config, PyTorch's per-config capture bakes the correct per-tier kernel and every replay reuses it -- capture-safe, no extra launch. Adds 3 wrapper-level graph capture/replay tests (M64, over-provisioned M256, and the heterogeneous [8,304]->MLDS upper-bound safety case). 47/47 pass. * mla_reduce: split-K for low-tile/high-split (b1_s128 11.0->7.4us) Cooperative multi-block split-K reduction for the latency-bound low-tile / high-split decode case (b1_s128 = 1 active tile x H=16 = 16 active blocks / 304 CUs, each serially reducing 128 splits). Opt-in via AITER_MLA_REDUCE_SPLITK (default OFF); default path byte-for-byte untouched. Two-kernel scheme (kernel boundary = free cross-block fence): - sk_partial_kernel (grid active_tiles*H*K): each block online-softmax partial-reduces a contiguous split subset of one (tile,head) into a pre-allocated scratch buffer (weighted acc + running max + sum-exp). - sk_combine_kernel (grid active_tiles*H): merges the K partials by global max renormalization; lse = ln(sum l)+M, matching the baseline exactly. plan_splitk engages only when profitable (max_seqlen_q==1, splits>=64, active_tiles*H<num_cu) from host-visible metadata. Scratch is pre-allocated once (lru_cache) and reused every CUDA-graph replay: no alloc / device sync / .item() in the launch path (capture-safe). Measured (GPU4 MI300X gfx942, graph us, x2): b1_s128 11.0->7.4 (0.67x baseline, 0.52x HIP 14.0), K=16 sweet spot (K=8 8.4, K=32 8.6). No regression on any b8_* shape (they don't engage). Correctness 53/53 (47 default + 6 new split-K: vs torch-ref K in {4,8,16}, vs HIP, cudagraph-replay, default-OFF). * mla_reduce: guard host-tier dispatch against num_kv_splits under-baking (opt-in) The num_kv_splits reaching flydsl_mla_reduce_v1 on the real dispatch is max_split_per_batch (a per-BATCH split budget), NOT a per-tile upper bound. The metadata (csrc/kernels/mla/metadata/v1_2_device.cuh:858, num_splits = min(num_clusters, max_split_per_batch * num_batches)) uses it only as a global payload divisor; the greedy CU load balancer can then concentrate a skewed batch's reduce tile up to num_cu = num_clusters splits (measured per-tile n_splits 171 at max_split_per_batch=32). So select_tier(num_kv_splits) = M64/M256 could bake a body whose fixed LSE cap (64/256) is below a real tile's split count and SILENTLY drop the overflow (smoking gun: abs err 1.83 on a 128-split tile). Add _safe_host_tier(): under AITER_MLA_REDUCE_HOST_TIER=1 trust only Tier.MLDS (nlse=5 covers 320 >= LDS_MAX_SPLITS=304 = num_clusters, the hard per-tile ceiling) and fall back to the always-correct device-side Tier.ALL for every smaller selection, so the opt-in flag can no longer under-bake. Default (flag off = Tier.ALL) is byte-identical. The true per-tile bound get_mla_decode_fwd_max_splits (= cu_num * occupancy = 304) selects MLDS, so the num_kv_splits=None dispatch path stays correct with no behavior change. Correctness (both flags on): test_flydsl_mla_reduce.py 47/47; test_mla_persistent.py (script) uniform + varlen + msb=304 pass. Also fix docstrings that incorrectly claimed num_kv_splits is a per-tile upper bound. * mla_reduce: device-adaptive capture-safe split-K, default-on Fold the b1_s128 cooperative split-K win into the production wrapper (flydsl_mla_reduce_v1) as a capture-safe, default-on path. plan_splitk_capture_safe takes its ENTIRE plan from host-only values: final_output.size(0) (= decode batch = active tiles), num_kv_splits (the max_split_per_batch budget, a true upper bound on actual per-tile splits), and num_cu -- no device read/sync -- so it engages under CUDA-graph capture, unlike the opt-in plan_splitk which reads the CSR via .item(). The per-tile K allocation stays device-adaptive, so ONE capture is correct across replays whose per-tile split counts vary. - b1_s128 graph 11.0 -> 7.4us (0.53x HIP); byte-identical single-kernel path on every other shape (heuristic declines when num_kv_splits<64 or grid saturated). - Correctness: full matrix 57/57 + new test_da_splitk_capture_safe_varying_splits (one bs=1 capture stays correct across per-tile splits [128,304,64,200,8,96]). - Off switch: AITER_MLA_REDUCE_DA_SPLITK=0. Decode-only (max_seqlen_q==1); assumes active tiles = CSR prefix. Existing opt-in plan_splitk untouched. * mla_reduce: gate DA split-K on actual_max_splits (phase-1 prototype) Add optional actual_max_splits to plan_splitk_capture_safe and flydsl_mla_reduce_v1 so engagement uses the true max per-tile split width instead of the loose num_kv_splits budget (~304 on persistent decode). - derive_actual_max_splits(reduce_indptr): planning-time CSR max (phase 2 replaces with metadata-emitted scalar). - When actual_max_splits is set, engage_splits uses it for min_splits gate and K sizing; None preserves legacy behavior. - Crossover probe: short-context over-engage edge closed (splits 2-32 delta ~0us); b1_s128 win preserved (128 splits: -3.6us vs DA-off). - Correctness: 61/61 incl. 4 new tests. Opt-in via actual_max_splits= arg; not folded to decode until phase-2 C++ metadata emit. * mla_reduce: emit actual_max_splits from metadata; gate split-K on it (phase-2) Phase-1 gated device-adaptive split-K on actual_max_splits derived host-side from the reduce CSR. Phase 2 has get_mla_metadata_v1 emit that scalar natively so no host CSR reduction is needed. C++ emit (opt-in, back-compatible): - New optional trailing output reduce_max_split (1-elem int32) on get_mla_metadata_v1. v1_2_device.cuh fills it via atomicMax at the two reduce_indptr write sites (parallel: num_frags, serial: num_splits) = max_t(reduce_indptr[t+1]-reduce_indptr[t]). Launcher zeroes it with hipMemsetAsync on the stream (capture-safe). - Threaded through metadata.cu, mla.h, pybind (rocm_ops.hpp), attention.py. - Defaults to nullopt; nullptr guard => non-emitting/HIP callers pay nothing. Plumbing: - mla_decode_fwd(actual_max_splits=None) -> both _mla_reduce_v1_dispatch sites -> flydsl_mla_reduce_v1. Pure pass-through; no hot-path sync. Tests (74/74 FlyDSL reduce matrix; +13 phase-2): - emitted scalar == derive_actual_max_splits for both planners x 10 shapes - metadata-sourced crossover + over-provisioned-budget edge closure - dispatch forwards actual_max_splits - persistent decode e2e (bs1 ctx8192 bf16+fp8): decode err = 0 Overhead is opt-in and on the once-per-shape planning path only (+2.7-7.4us metadata), cheaper than the phase-1 device .max() reduction. Co-authored-by: Cursor <cursoragent@cursor.com> * mla_reduce: lv3 depth lever - default-on M64 deep sub-split (GRP16, -5% b8_s32/s26 graph) The accumulate loop is already a GRP-wide double-rate software pipeline (vmcnt overlap up to 15, well beyond the HIP depth-2 oaccu_0/oaccu_1), so the literal "depth-1 -> depth-2" conversion would regress. The remaining pipeline-depth win is in the M64 path: enable the existing device-side sub-split by default (M64_HI_THR 0->8) so high-split M64 tiles take the deeper GRP=16 accumulate (more os buffer_loads in flight) while the low-split tail keeps GRP=8. Measured graph us (optionb, same session): b8_s32 5.9->5.6 (-5%, 0.97->0.92x HIP), b8_s26 5.9->5.6, b8_s13 4.7->4.5; b1_s128 (M256) and the SIMPLE tail (b8_s3) unchanged; b8_s6/s5 keep GRP=8 (byte-identical). ISA: vmcnt max 15->24, buffer_load 97->122 (deeper overlap confirmed). Correctness: 74/74 pytest matrix pass. Capture-safe (device branch). * mla_reduce: adaptive active_tiles×H launch (default-on, multi-tile only) Launch one block per active (tile, head) on sparse multi-tile decode instead of the persistent grid-stride kernel, eliminating the traversal WhileOp floor. Gated: num_final_rows > 1 (split-K owns bs=1), num_final_rows < num_reduce_tile. Opt-out: AITER_MLA_REDUCE_ADAPTIVE_LAUNCH=0. b8_s32 graph 5.6→5.2µs (1.15× HIP). Co-authored-by: Cursor <cursoragent@cursor.com> * mla_reduce: drop redundant accumulate mask from GRP pipeline state load_scales already zeroes invalid split scales and OOB os reads are pmap0-substituted, so carrying mask_g through the loop only added dead VALU. Removes ~230 instructions and cuts b8_s32 graph 5.2→4.9µs on the production adaptive path without changing SIMPLE or split-K behavior. Co-authored-by: Cursor <cursoragent@cursor.com> * mla_reduce: low-split direct pmap path for adaptive decode When actual_max_splits is known and <= 8, compile a separate adaptive kernel that reads reduce_partial_map directly instead of staging to LDS first. This removes the fixed pmap-staging barrier on low-split serving shapes while high-split captures keep the vectorized LDS pmap path via separate JIT entries. Co-authored-by: Cursor <cursoragent@cursor.com> * bench_mla_reduce: add GLM-5.2 serving scoreboard as default path Make PR numbers reproducible from the aiter tree: default bench runs the production wrapper (trimmed final_output, actual_max_splits, adaptive + DA split-K) against hip/wrapper-daoff/wrapper-daon. Keep uniform/irregular/replay under --mode for synthetic and metadata replay sweeps. Co-authored-by: Cursor <cursoragent@cursor.com> * Drop actual_max_splits from mla_decode_fwd; resolve via warmup cache. Auto-derive the split-K gate input inside flydsl_mla_reduce_v1 from reduce_indptr using a capture-safe warmup-populated cache, so callers need not thread actual_max_splits through mla.py. Removes the stale dispatch docstring and fixes mla_prefill_ps_fwd forwarding a missing arg. Co-authored-by: Cursor <cursoragent@cursor.com> * bench_mla_reduce: use production actual_max_splits auto-resolve Stop passing actual_max_splits explicitly in the serving harness; let flydsl_mla_reduce_v1 resolve it from reduce_indptr via the warmup cache so graph replay matches mla_decode_fwd. Plan annotation still uses derive_actual_max_splits at host planning time. Co-authored-by: Cursor <cursoragent@cursor.com> * mla_reduce: always-on DA split-K; simplify bench to hip/wrapper Remove AITER_MLA_REDUCE_DA_SPLITK and da_splitk_enabled(); split-K engage logic in plan_splitk_capture_safe is always active. Bench backends are now hip and wrapper only (production flydsl_mla_reduce_v1 path). Co-authored-by: Cursor <cursoragent@cursor.com> * mla_reduce: drop AI-narrative comments, simplify wrapper/bench/tests Trim docstrings to production style, remove opt-in host-tier/adaptive-launch env gates (Tier.ALL + adaptive-on are now unconditional), dedupe wrapper setup, rewrite the benchmark to the standard @benchmark()/pandas pattern, and drop dead helpers left over from the cleanup. Co-authored-by: Cursor <cursoragent@cursor.com> * mla_reduce: restore uniform/irregular benches and fix test/dispatch issues Restore uniform and irregular benchmark sweeps under the @benchmark interface. Re-read AITER_MLA_REDUCE_FLYDSL on each dispatch (cache only FlyDSL availability) and work around module_mla_metadata check_args failures in reduce_max_split tests. Co-authored-by: Cursor <cursoragent@cursor.com> * mla_reduce: fix black/ruff style on kernel and hip bench Remove an unused import, rename ambiguous loop variables, and apply black formatting so branch Python files pass lint checks. Co-authored-by: Cursor <cursoragent@cursor.com> * mla_reduce: fix HIP reference LDS sizing on GPUs with <304 CUs hip_ref/hip_ref_like_fout passed num_kv_splits=0, so the HIP kernel sized its LDS from the GPU's CU count. On gfx950 (256 CUs) that undersizes the buffer for fixtures using up to LDS_MAX_SPLITS (304) splits, causing failures on MI35X CI while gfx942 (304 CUs) passed. Co-authored-by: Cursor <cursoragent@cursor.com> * mla_reduce: fix CI failures on metadata tests and MI35X Forward reduce_max_split through get_mla_metadata_v1 wrapper so metadata gate tests can call the public API, and scope FlyDSL MLA reduce tests to gfx942 only until MI35X is a supported target. Co-authored-by: Cursor <cursoragent@cursor.com> * remove unnecessary pytest markers section * mla_reduce: remove superseded reduce_max_split metadata emission The phase-2 reduce_max_split output on get_mla_metadata_v1 was replaced by the capture-safe warmup cache (_resolve_actual_max_splits), which resolves actual_max_splits host-side without threading a metadata scalar. It was dead in production (only tests used it), so drop it entirely from the C++ metadata kernel/binding, the Python op + wrapper, and its tests. Co-authored-by: Cursor <cursoragent@cursor.com> * Restore CP round-robin comment dropped by unrelated cleanup The comment documenting round-robin context-parallel semantics was removed as collateral damage in a Claude-comment cleanup pass; it doesn't belong to this branch's work, so restore it. Co-authored-by: Cursor <cursoragent@cursor.com> * space * mla_reduce: fold HIP baseline bench into FlyDSL script, drop replay mode Add an opt-in --include-hip flag to the FlyDSL benchmark so it can run the production HIP kernel as a comparison candidate across all sweeps, and remove the now-redundant standalone HIP bench (and its unused replay mode). Co-authored-by: Cursor <cursoragent@cursor.com> * Fix MLA reduce test imports for CI. CI runs op tests with python3 directly, so fix the import path for flydsl_mla_reduce_common. Co-authored-by: Cursor <cursoragent@cursor.com> * test(mla_reduce): merge FlyDSL perf bench into the correctness suite Drop pytest in favor of a plain python3 entrypoint so aiter CI's exit-code-based test runner actually exercises this suite, then fold the standalone bench script's serving/uniform/irregular perf sweeps in as the tail of the same run (gate first, sweep only if it passes). Co-authored-by: Cursor <cursoragent@cursor.com> * style(mla_reduce): black-format test file Co-authored-by: Cursor <cursoragent@cursor.com> * Inline flydsl_mla_reduce_common into test_flydsl_mla_reduce Merge the single-use helper module into its only consumer so the op test is self-contained, per review feedback on PR 3901. Co-authored-by: Cursor <cursoragent@cursor.com> * refactor FlyDSL MLA reduce APIs Co-authored-by: Cursor <cursoragent@cursor.com> * fix FlyDSL MLA runtime and cache isolation Co-authored-by: Cursor <cursoragent@cursor.com> * fix: harden FlyDSL MLA reducer safety Keep the validated reducer changes and CI corrections while excluding internal PR notes from the source tree. Co-authored-by: Cursor <cursoragent@cursor.com> * refactor(flydsl): remove env-var tuning knobs from mla_reduce Replace all os.environ reads in the kernel with function parameters (defaults unchanged), and drop a dead math-wrapper branch. Co-authored-by: Cursor <cursoragent@cursor.com> * test(flydsl): consolidate MLA reduce coverage Merge eager/CUDA-graph replay test pairs behind a replay flag, fold the small_split graph case into the shared case table, extract shared helpers for the guard-differential and split-K b1_s128 tests, replace the hand-written run_checks() call list with a data-driven registry, dedupe benchmark roofline/candidate/table logic across the three perf sweeps, and drop redundant local imports/duplicated constants. No scenario, assertion, or CLI behavior changes. Co-authored-by: Cursor <cursoragent@cursor.com> * refactor(flydsl): move MLA scratch and output stores off buffer_ops Two of the four raw buffer_ops paths in mla_reduce now go through the public layout API, cutting direct call sites from 13 to 5. The f32 scratch load and store use a buffer tensor with a copy atom over a row view; VEC=8 f32 is wider than one atom, so it composes two atoms over adjacent chunks -- the same shape fused_compress_attn already uses, expressed through the layout API. The final-output store moves into shared helpers used by both the normal and split-K combine kernels, where a single atom always covers the fragment. A probe confirmed the layout copy still emits one packed store, refuting an in-code comment that had claimed it would scalarize the write. Invariant checks pass, ISA opcode counts are identical across all four emitted kernels, VGPR counts are unchanged, and the perf sweep shows no regression. The two paths still on buffer_ops are documented in the code. Co-authored-by: Cursor <cursoragent@cursor.com> * refactor(flydsl): drop remaining buffer_ops from mla_reduce mla_reduce no longer references buffer_ops. The indptr reads use a plain layout view, which the backend already lowers to a scalar load, so the raw descriptor and the explicit readfirstlane broadcast were both unnecessary. The partial-output load now uses the same slice/copy form as the output store; the regression seen in the earlier attempt came from a redundant predicate, not from the load itself. The f32 atom/chunk helpers move to module level so the 2-D and 3-D paths share one composition. All bounds guards are unchanged. Invariant checks pass, the split-K ISA is byte-identical, VGPR counts are unchanged, and the perf sweep shows no regression. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(flydsl): satisfy Ruff 0.16 style checks Remove stale noqa and simplify mechanical expressions flagged by the pinned CI linter. Co-authored-by: Cursor <cursoragent@cursor.com> * refactor(flydsl): clean up MLA reduce dispatch and tests Co-authored-by: Cursor <cursoragent@cursor.com> * Update aiter/mla.py Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * refactor(flydsl): use empty split-K scratch buffers The partial kernel fully overwrites scratch rows each launch, so zero-init is unnecessary. Co-authored-by: Cursor <cursoragent@cursor.com> * Keep FlyDSL MLA reduction decode-only Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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
Add tuned configs for gfx1201
ck_gemm_a8w8_blockscalekernel, targeting various Qwen3 model variants. gfx1201 supports the FP8 dtype so these tuned configs speed up the gemm_a8w8_blockscale for inference in vLLM.Technical Details
Tuning Process
The tuning was performed using the CK GEMM tuner:
The tuned configurations are added to the existing GEMM configuration files and are automatically selected based on the input tensor dimensions and the target architecture (gfx1201).
Test Plan
The tuned kernels were validated using the GEMM test suite:
Tests cover various matrix dimensions (M: 1-10240, N: 24576, K: 1536) that are representative of Qwen3 inference workloads.
Test Result
All tests pass with zero error, confirming the correctness of the tuned configurations.
Submission Checklist