gfx1250 opus gemm splitk fuse - #4246
Merged
Merged
Conversation
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
demonsan
changed the title
demonsan
force-pushed
the
binglue/gfx1250-opus-gemm-splitk-fuse
branch
7 times, most recently
from
July 27, 2026 02:26
8b15272 to
b7a7bac
Compare
Append the gfx1250 K=7168 shape block to dsv4_bf16_tuned_gemm.csv: tuned shapes use the opus results from dsv4_gfx1250_tuned_gemm.csv, and the remaining untuned shapes fall back to triton, matching the baseline scheme. Co-authored-by: Cursor <cursoragent@cursor.com>
Add the fused split-k pipeline and split-k reduce launch for gfx1250 opus gemm, extend split-k reduce to support the bf16 workspace + fp32-bias path, and update codegen/gen_instances to emit the fused instances. Co-authored-by: Cursor <cursoragent@cursor.com>
The aiter imports must run after _patch_flaky_hip_device_count(); mark them with noqa: E402 to satisfy the ruff code-style CI. Co-authored-by: Cursor <cursoragent@cursor.com>
…sults Overwrite the 59 tuned gfx1250 K=7168 shapes with the latest tuning run; some large-M shapes now select triton over opus. Co-authored-by: Cursor <cursoragent@cursor.com>
Add a golden-checked CUDA-graph opus_gemm sweep (default entry) with split-K workspace pre-warm on the capture stream, and print a latency-vs-tuned-CSV table for gfx1250. Co-authored-by: Cursor <cursoragent@cursor.com>
…1250 dispatch - Add opus_gemm_workspace_release() (current stream) and opus_gemm_workspace_release_all() (all streams) to free the per-stream split-K workspace buffer, host/device handles and registry entry, with eager-mode guards and pre-free synchronization; wire up pybind + Python. - Make the gfx1250 dispatch self-contained: move the shared flat-array dispatch POD types into opus_gemm_heuristic_dispatch_gfx1250.cuh (opus_gfx1250_detail) and drop the gfx950 header include so a gfx1250-only build compiles without OPUS_BUILD_HAS_GFX950. - Allocate the split-K workspace handle from plain host memory on the gfx942/gfx1250 (device-mirror) path instead of hipHostMallocCoherent, so those builds use no pinned/coherent host memory (reclaimed by the OS at process exit); gfx950 keeps the coherent handle it dereferences on device. Co-authored-by: Cursor <cursoragent@cursor.com>
…try) Replace the SplitkWsRegistry/hipMalloc-based workspace management for gfx1250 split-K kernels with direct pointer passing from Python-allocated torch.empty buffers (bf16 by default, matching OPUS_WS_BF16=1). Changes: - kargs struct: ws_handle* -> void* ptr_ws (direct pointer, no indirection) - main kernels: kargs.ws_handle->ptr -> kargs.ptr_ws - reduce kernel: const opus_splitk_ws_handle* -> const void* ws_ptr - codegen launcher: remove registry/grow logic, add workspace tensor param - C++ dispatch: opus_gemm_a16w16_tune gets optional workspace param; gfx1250 split-K kids pass workspace, others unchanged - Python layer: allocate workspace via torch.empty(dtype=bf16) for kid [20000,30000) - OpusA16W16NoscaleKernel for gfx1250 moved to opus_gfx1250_detail namespace (6-arg signature with workspace, separate from gfx950/gfx942 5-arg type) - gen_instances.py manifest: emit 6-arg forward decl for gfx1250 split-K launchers - gfx950/gfx942 paths unchanged (still use SplitkWsRegistry) Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: binglue.lyu <binglue.lyu@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
…rewarm
The gfx1250 split-K partial-sum workspace was allocated as a bf16 buffer sized
from an over-estimated raw element count. Kids whose kernel writes an fp32
workspace then overran the buffer (fp32 is 2x the bytes), corrupting adjacent
device memory and hanging the machine.
Size the workspace from the selected kid's OWN kernel definition instead:
* allocate a [batch, split_k, padded_M, padded_N] tensor by element count in
the kid's workspace dtype (fuse kids use compile-time fuse_split_k /
fuse_ws_dtype; ws-variant kids use fp32 with runtime split_k), never a raw
byte size;
* fall back to a safe fp32 over-estimate when the kid table can't be loaded.
Also drop the eager graph-capture prewarm in tuned_gemm: the workspace is a
plain torch.empty passed to the launcher, and torch's HIP graph-capture-aware
caching allocator makes a single torch.empty path valid for both eager and
capture (replay-stable address, no pre-warm needed). The gfx942/gfx950 hipMalloc
registry path is unaffected and still fails loudly if used under cudagraphs
without a warm.
Co-authored-by: Cursor <cursoragent@cursor.com>
Replaces the positional tdm_desc constructor with tdm_traits::resolve<> over an unordered pack of policy tags, so a call site names only the concerns it varies and every default and validation lives in one place. Adds the dim1..dim3 stride programming that 3D-5D windows were missing, so a tile past 2D walks planes by the D#'s own pitch instead of landing on top of itself, and rejects gather with ndim>2 at compile time since both claim D# groups 2/3. Moves the LDS write point out of window state into a make_descriptor() parameter, leaving slot cycling to the pipeline that already owns the s_wait_tensorcnt depth. Makes every offset the public API takes an element count, matching the extents, strides and origins it already took. The LDS byte scaling folds into the address add, so the 4wave GEMM comes out 10 instructions shorter.
…g 24 The two gather bits in D# group 0 dword 0 were transposed. The reference lowering in llvm-project (AMDGPUMakeDmaBaseLowering in AMDGPUToROCDL.cpp, with the offsets pinned by name in the gfx1250 conversion test) puts gather_mode at bit 30 and index_size at bit 31; opus had them the other way round. gather<32> sets both bits either way, which is why the encoding was accidentally correct everywhere it had been exercised, but gather<16> emitted 0x80000001 where the reference emits 0x40000001 -- index_size of 32 bits with gather disabled, so the window silently degraded to a plain contiguous copy instead of failing. Since async_store shares make_descriptor(), scatter carried the same bug and is fixed with it. The rest of the descriptor was checked field by field against the same lowering and agrees, as does the 16-bit index packing and the reuse of tile_dim1 as the valid-index count. Also drops the runtime rewrite of those bits in make_descriptor(). Both are compile-time and already sit in group0_dword0_const, which nothing else patches, so it was dead -- and holding one encoding in two places is what let the transposition survive. Three clang-24 builtin changes broke the header outright. aac296b87e8c moved the AMDGPU builtins from GNU vectors to ExtVector, so the lax conversions that used to paper over spelling differences are gone, and 7318bc734a5b retyped D# group 0 as unsigned while leaving groups 1-3 signed. fp16_t is therefore registered as _Float16 from clang 24 on, which is what the fp16 wmma/mfma builtins now demand and what opus itself used before clang 20; the older compilers keep __fp16 rather than shifting their numerics, since respelling turns arithmetic on fp16_t from the fp32-intermediate v_fma_mixlo_f16 into a native v_fma_f16. Group 0 gets a bit_cast behind OPUS_HAS_UNSIGNED_TENSOR_D0, generating nothing since the lanes are unchanged, and the b64/b128 buffer store vdata goes unconditionally unsigned because it was declared that way in the .def era too. Verified to produce identical instructions under clang 23 and clang 24. Tightens the sub-byte element check. Only array<fp4_t, 2> and its vector form are accepted now: the bare type reports sizeof 1 for a 4-bit value so the descriptor strides twice as far as the data, and an array or vector of any other width makes the D# element several values at once, rescaling every extent, stride and pad. The old sizeof test caught only the bare type. The remainder is presentation. tdm_traits becomes a namespace rather than a struct of statics, tdm_desc reaches through one Traits alias instead of five, default_cache_policy comes from make_cache_policy() now that it is not stuck in an incomplete-class context, and the longer comment blocks are condensed to the file's single-line style with the measurements kept. Co-authored-by: Cursor <cursoragent@cursor.com>
The gather mode and index size bits go back where they were. The previous commit moved gather_mode to bit 30 and index_size to 31 to agree with MLIR's AMDGPUToROCDL lowering; on gfx1250 the lowering is the one that is wrong. Issuing one descriptor twice with nothing but that dword changed, 0x80000001 fetches the listed rows and 0x40000001 returns consecutive rows from the origin. Only gather<16> can tell the two orderings apart -- gather<32> is 0xC0000001 either way -- which is how the transposition passed review and why no 32-bit workload noticed it. tdm_groups23 now receives tile_dim3 and tile_dim4 as template arguments. tdm_desc computed them into group2_dword3_const and group3_dword2_const, but the groups23 aggregate zeroed both operand groups and neither constant was ever read, so a rank-4 tile collapsed onto its dim3 = 0 slice and a rank-5 tile onto dim3 = dim4 = 0. Ranks 2 and 3 were unaffected because tile_dim2 rides in group1, which is why nothing before this noticed. test_tdm_feature_gfx1250.cu is the suite that found both: one self-contained TU covering element sizes 1/2/4/8 bytes and array<fp4_t,2>, tile ranks 2..5 each with a move along its outermost dimension, gather and scatter at both index widths, runtime descriptor state, padding, multicast and clamping. Every check is byte-exact against a position hash, and untouched memory is checked as hard as copied memory -- clamping, short gather lists and the tile_dim1 patch are only visible that way. Both halves of the group0 encoding are asserted, so a future transposition goes red instead of silently costing gather<16> its rows. test_tdm_gfx1250.cu is removed: it is written against tdm_window, which the refactor took out of opus.hpp, so it has not compiled for some time. Its coverage -- a pipelined multi-wave GEMM with cluster multicast -- is not what the new suite checks and wants porting rather than deleting. Around the harness: setup.py gains a positive per-arch list so the suite is built only for gfx1250, an arch fallback through torch for hosts without rocminfo, and named entries for two sources that do not currently compile there, so one stale source no longer takes the whole .so with it. The wmma_scale tests gain the missing-symbol guard the rest of the harness already uses, since they look their launchers up by name at call time. Verified on gfx1250 under both the ROCm 7.x clang and the clang 24 / HIP 7.15 toolchain in the opus_rocm_compiler container: 26 runtime checks pass, the 17 compile-time rejections all fire, opus_tdm_example is unaffected, and the bf16 GEMM still runs at 147 TFlops.
The "feature" qualifier distinguished nothing -- it is the only TDM suite in the directory. setup.py's source list and arch-only map follow the new name.
The TDM refactor replaced the positional tdm_window<> constructor with an unordered pack of policy tags and moved every offset to element units, which left the three gfx1250 a16w16 split-K pipelines and their traits header on an API that no longer exists. Rebuild WindowA/WindowB on opus::tdm<> with a padding_auto tag in place of the kLdsPadEn / kPadInterval / kPadAmount constants the old descriptor carried, and retire the byte-denominated slot sizes in favour of kSlotElemsA/B. The producer loops now take the slot index as a template parameter and hand make_descriptor() an element offset, so the double-buffer half is chosen at the issue site instead of living in window state. The clusterlaunch variants read their multicast masks from tdm_traits::peers_along_y/x<> rather than deriving them in a loop, and the fuse variant's workspace staging follows the same element-offset form. Retune the gfx1250 blocks of the dsv4 and kimik3 bf16 GEMM configs (180 + 345 shapes) against opus and triton measured back to back in one session, keeping the faster backend per shape: 212 shapes select opus and none regress against the triton fallback they replace. The sweep searched the 496 plain + clusterlaunch kids; OPUS_TUNE_NO_FUSE=1 is the switch that held the ~1.4k fuse kids out of the compile set, so those remain unmeasured. Co-authored-by: Cursor <cursoragent@cursor.com>
tdm_window<> and its positional make() no longer exist. The 2D x-load becomes an opus::tdm<> window whose double-buffer half is handed to async_load() as an element count, rather than being folded into the LDS base address in bytes. The 3D residual load has no make() overload to move to, so it programs the descriptor through make_from_layout() with explicit shape, pitch and origin arrays in D# order. That also retires the placeholder barrier argument the old eight-parameter call had to thread through. Co-authored-by: Cursor <cursoragent@cursor.com>
The style CI has been red on this branch. Three separate findings, all on lines the branch adds: ruff flags the five `# noqa: E402` directives on the imports that have to run after _patch_flaky_hip_device_count(). Under the 0.16.0 the workflow pins, those imports no longer trip E402 at all, so the suppressions are dead and RUF100 reports them. Drop them; the comment above the block still records why the imports sit where they do. ruff also flags two idioms in gemm_op_a16w16.py that reviewdog already commented on: lru_cache(maxsize=None) becomes functools.cache, and the splitK>1 ternary becomes max(1, ...). black wants the _alloc_splitk_workspace call collapsed onto one line, which it now fits after those edits. Co-authored-by: Cursor <cursoragent@cursor.com>
demonsan
force-pushed
the
binglue/gfx1250-opus-gemm-splitk-fuse
branch
from
August 10, 2026 06:18
86d1eba to
25dd628
Compare
… an exact cluster fill A clusterlaunch kid used to require the tile grid to be a whole number of clusters, enforced by an assert at launch and mirrored by a divisibility filter in the tuner. That left most shapes with no cluster candidate at all. The launcher now rounds the tile grid up to whole clusters and the workgroups the round-up adds leave at `tile_oob` in the pipeline prologue -- after paying their one cluster-barrier arrival, which -3 counts per workgroup, and before any TDM is issued. Their peers therefore never wait on them and never see a request whose extents disagree, so any (CWGM, CWGN) is legal for any (M, N). The workspace strides stay on the unrounded tile counts, so the reduce kernel is unaffected, and the epilogue store loses its now-dead OOB guard. This also retires the 2D-cluster lockout: what those configurations hung on was the tile-less workgroup streaming the whole K loop as zero-extent multicast loads at peers whose extents were real, not a deeper strided-A multicast defect. The tuner consequently sweeps (cwm, cwn) over [1,4]x[1,4] minus (1,1), bounded by a round-up budget rather than by divisibility: drop a cluster side wider than the grid it rides on, drop anything leaving more than GFX1250_MAX_CLUSTER_WASTE of the launched workgroups tile-less, then rank by (waste bucket, widest multicast group, grid-aspect match). Ranking by group size rather than by proximity to 8/16 follows the measurement -- a 12x12 grid runs 10.7 us at 3x3 against 13.5 us at 1x2, and on a lopsided 4x32 grid aspect is worth as much as a size step (2x4 at 11.2 us against 4x2 at 13.3 us). dsv4 gfx1250 is retuned against this, opus only, on 179 of its 180 shapes (M=16384 N=129280 has a 3.95 GiB output and every kid is rejected by the 4 GiB buffer-resource filter). The tuned times over those shapes sum 1.52x lower, and opus now holds 157 of the 180 rows against 101 before.
…e is fixed The fused single-kernel split-K pipeline still misbehaves, so keeping its 1378 kids in the table only costs tuning time and risks a bad pick. Gate the whole sweep on GFX1250_SPLITK_FUSE_ENABLED (False): the kid list comes out empty, so nothing looks the family up, the tuner has nothing to select and the codegen has nothing to emit -- and the [21000, ...) band it used to claim is free for another kernel family. The factory, the codegen emitter and the device pipeline all stay in place, so re-enabling is a one-line flip. This replaces the OPUS_TUNE_NO_FUSE escape hatch that every recent tune run had to remember to set. Co-authored-by: Cursor <cursoragent@cursor.com>
Conflicts came from main's new a8w8_mxscale BMM family (#4320) landing in the same two files this branch edits, and both sides were purely additive: * opus_gemm_common.py OpusGemmInstance: keep this branch's fuse_* axes next to main's mxscale BMM axes; both groups are defaulted trailing fields. * gen_instances.py manifest emit: keep both signature templates and fold the gfx1250 6-arg workspace choice into main's A16W16_TUNE_TAGS branch, after the BMM branch (the tag sets are disjoint, so the order is free). Verified after the merge: 496 gfx1250 kids with none in the fuse band, tuner candidate counts unchanged, and codegen over gfx950+gfx942+gfx1250 emits both the 44 BMM and the 63 workspace launcher declarations. Co-authored-by: Cursor <cursoragent@cursor.com>
…mpile opus_gemm_lookup.h and opus_gemm_a16w16_tune_lookup.h each defined one macro per output dtype, listing every kid in the compile subset regardless of arch. Each arch's dispatch header then expanded those macros into a table of its own launcher function-pointer type -- but gfx1250's a16w16 launchers take a workspace tensor that gfx950/gfx942's do not, so any build targeting gfx1250 alongside another arch failed to type-check: gfx950's table held gfx1250 function pointers and vice versa. GPU_ARCHS filtering of |S| hid this for single-arch builds only. A gfx950+gfx1250 host compile produced 439 such errors across all six tables. Emit one macro per (dtype, arch) instead and have each arch header expand only its own, which is also what the per-arch host TU split already does for the split-K reduce ABI. Table contents are unchanged for every single-arch build -- verified entry-for-entry against the base commit for gfx950, gfx942 and gfx1250, with the rest of the generated blob byte identical. Winners in get_tune_dict are now keyed on arch too. The (M,N,K) tables are per arch, so without arch in the key a shape tuned on two arches kept only whichever CSV row was read last and the other arch silently fell back to its heuristic: an all-arch build dropped 7 of 308 tuned rows (gfx1250 266/268, gfx950 35/40) before this change and drops none after. Verified by host-compiling opus_gemm.cu for gfx950;gfx1250, gfx942;gfx950;gfx1250, gfx942;gfx1250 and each arch alone. Note that a full mixed-arch device build is still blocked by an unrelated issue: the gfx1250 clusterlaunch pipeline applies __cluster_dims__ to its kernel unguarded, which the gfx950 device pass rejects. That needs an arch-guarded device stub and is not addressed here. Co-authored-by: Cursor <cursoragent@cursor.com>
…the cluster kernels Both clusterlaunch kernels already compile their body away to an empty stub outside gfx1250, but the __cluster_dims__ attribute sat outside that guard. A mixed-arch build compiles every device TU once per offload arch, and the gfx950/gfx942 pass rejects the attribute before it ever reaches the body: opus_gemm_pipeline_a16w16_clusterlaunch_tdm_splitk_ws_gfx1250.cuh:75:1: error: 'cluster_dims' is not supported for this GPU architecture 1 warning and 1 error generated when compiling for gfx950. Guard it the same way, keeping the attribute on the host pass, which is where the launch reads the cluster geometry from. Extended to the fuse kernel too -- that family is unregistered today, so it has no TU to fail in, but the bug is identical and would come back with it. gfx1250 codegen is unchanged by construction (__gfx1250__ is defined on that pass). Verified: gfx1250 device TUs now compile clean with --offload-arch=gfx1250 --offload-arch=gfx950, a gfx1250-only build still produces a correct GEMM (rel_err 2.7e-03 bf16), and four compiled clusterlaunch kids dispatched by id are numerically clean (rel_err 5.4e-05), so cluster multicast is intact. A full mixed-arch build now gets past the gfx1250 TUs and hits the mirror image of this problem, which is not addressed here: gfx950's kid TUs are also compiled for gfx1250, where get_warp_size() is 32 instead of 64, so wave-size assertions in the gfx950 traits fail (117 errors, e.g. "splitk BLOCK_SIZE must cover exactly 4 waves"). Those TUs have no business being compiled for another arch at all -- the fix belongs in the build flags (per-source --no-offload-arch), not in more source guards. Co-authored-by: Cursor <cursoragent@cursor.com>
…s compile hipcc runs every TU once per --offload-arch, so a GPU_ARCHS=gfx950;gfx1250 build also instantiated the gfx950 kids for gfx1250: the gfx950 traits then derived their layouts from a 32-wide wave and tripped `BLOCK_SIZE == 4 * get_warp_size()`, and the reverse direction reached gfx1250-only kernel attributes. Wrap the generated device TUs (per-kid instances and the per-arch splitk reduce) in `#if !defined(__HIP_DEVICE_COMPILE__) || defined(__<arch>__)`, covering the #include as well, so a foreign device pass is an empty TU that never parses another arch's headers. The host pass stays inside the guard since it emits the __device_stub__ symbols the fused host TU links against. Each object now carries its kernel only in its own arch's code object (verified with llvm-nm on the unbundled images for gfx942/gfx950/gfx1250). Builds clean for gfx950, gfx942, gfx1250, gfx950;gfx1250 and all three at once; op_tests/test_opus_a16w16_gemm.py passes 15/15 on gfx1250 against both the gfx1250-only and the mixed module. Co-authored-by: Cursor <cursoragent@cursor.com>
The GLM 5.x tuning commit added a block_m grid divisor but only assigned it in the gfx950 branch, so gfx1250/gfx942 raised UnboundLocalError at launch. Those kernels take no BLOCK_M and run one query row per program, so 1 restores the original (seq_len,) grid. Co-authored-by: Cursor <cursoragent@cursor.com>
junhaha666
approved these changes
Aug 14, 2026
Fyzyukk
added a commit
to Fyzyukk/aiter
that referenced
this pull request
Aug 17, 2026
Reapply upstream commit 1b741c0 on top of the unified public API, canonical exact-kid registry, and caller-owned Torch workspace design. Adopt the policy-tag N-D TDM API, clusterlaunch grid round-up, compile-time gfx1250 reducer dispatch, mixed-arch device guards, tuned configs, and the final decision to leave the fused family unregistered.
xiaobochen-amd
pushed a commit
to xiaobochen-amd/aiter
that referenced
this pull request
Aug 28, 2026
… offset, not tensor bytes (#4) * [TRITON][GLUON] Prefill MQA Logits kernel tuning for GLM 5.x (ROCm#4563) (cherry picked from commit aded0f8) Also carries the two-line `block_m = 1` for the non-gfx950 branch from ROCm#4246. This commit routes the launch grid through block_m but only defines it under `arch == "gfx950"`, so without those lines every other arch raises NameError at the kernel launch. Upstream ran with that hole from ROCm#4563 until ROCm#4246. * [Bugfix][Triton][gfx950] fp8_mqa_logits: gate buffer ops on the int32 offset, not tensor bytes Prefill shapes whose fp32 logits pass 2 GiB abort the AMDGCN backend at JIT time (Sequence.h:275 "Begin must be less or equal to End"). BLOCK_M=2 is selected for seq_len > 4096 but only compiles with buffer stores, and the buffer-store gate switches off at 2 GiB, so the two combine into a hard crash. For GLM-5.x that is any chunked prefill of 8192 tokens against a context past 65,536. The gate had the wrong unit. Buffer ops address through a 32-bit offset, but the kernel re-bases the pointer per row and per KV tile before each access, so that offset never has to span the tensor. What must fit in int32 is the largest element offset the kernel forms, because the row strides stay 32-bit on the buffer path and are only widened to int64 on the fallback path. Counting bytes rather than elements made the limit 4x too tight for an fp32 output. Measured on MI355X, 32 heads x 128 head_dim, against the plain-store path. Output is bit-identical in every case (max rel err 0.00e+00): s_q s_k logits plain store buffer store speedup 8192 65536 2.00 GiB 3.557 ms 2.941 ms 1.21x 8192 95457 2.91 GiB 6.105 ms 4.577 ms 1.33x 8192 131072 4.00 GiB 9.812 ms 6.489 ms 1.51x 16384 131072 8.00 GiB 21.738 ms 12.879 ms 1.69x 8192 262144 8.00 GiB 25.710 ms 14.184 ms 1.81x The new boundary is exact rather than approximate: 8192x262144 and 16384x131072 both place the largest offset at exactly INT32_MAX and are bit-correct, while 16384x139264 (1.06 x 2^31) and 32768x131072 (2^32) fault. End to end, a 21-layer GLM-5.x indexer prefill of 9,695 new tokens against a 96,960-token context at chunk 8192 drops from 150.2 ms to 118.4 ms. The existing cases top out at s_q=1024, s_k=1560, four orders of magnitude below the gate, which is why nothing caught this. Added 8192x65664 and 8192x98304, which crash the process on current main. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: Mehmet Cagri <mehmet.kaymak@amd.com> Co-authored-by: Jiejing Zhang <jiejizha@smci355-ccs-aus-n02-29.prov.aus.ccs.cpe.ice.amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
JohnQinAMD
pushed a commit
to xiaobochen-amd/aiter
that referenced
this pull request
Aug 28, 2026
… offset, not tensor bytes (#4) * [TRITON][GLUON] Prefill MQA Logits kernel tuning for GLM 5.x (ROCm#4563) (cherry picked from commit aded0f8) Also carries the two-line `block_m = 1` for the non-gfx950 branch from ROCm#4246. This commit routes the launch grid through block_m but only defines it under `arch == "gfx950"`, so without those lines every other arch raises NameError at the kernel launch. Upstream ran with that hole from ROCm#4563 until ROCm#4246. * [Bugfix][Triton][gfx950] fp8_mqa_logits: gate buffer ops on the int32 offset, not tensor bytes Prefill shapes whose fp32 logits pass 2 GiB abort the AMDGCN backend at JIT time (Sequence.h:275 "Begin must be less or equal to End"). BLOCK_M=2 is selected for seq_len > 4096 but only compiles with buffer stores, and the buffer-store gate switches off at 2 GiB, so the two combine into a hard crash. For GLM-5.x that is any chunked prefill of 8192 tokens against a context past 65,536. The gate had the wrong unit. Buffer ops address through a 32-bit offset, but the kernel re-bases the pointer per row and per KV tile before each access, so that offset never has to span the tensor. What must fit in int32 is the largest element offset the kernel forms, because the row strides stay 32-bit on the buffer path and are only widened to int64 on the fallback path. Counting bytes rather than elements made the limit 4x too tight for an fp32 output. Measured on MI355X, 32 heads x 128 head_dim, against the plain-store path. Output is bit-identical in every case (max rel err 0.00e+00): s_q s_k logits plain store buffer store speedup 8192 65536 2.00 GiB 3.557 ms 2.941 ms 1.21x 8192 95457 2.91 GiB 6.105 ms 4.577 ms 1.33x 8192 131072 4.00 GiB 9.812 ms 6.489 ms 1.51x 16384 131072 8.00 GiB 21.738 ms 12.879 ms 1.69x 8192 262144 8.00 GiB 25.710 ms 14.184 ms 1.81x The new boundary is exact rather than approximate: 8192x262144 and 16384x131072 both place the largest offset at exactly INT32_MAX and are bit-correct, while 16384x139264 (1.06 x 2^31) and 32768x131072 (2^32) fault. End to end, a 21-layer GLM-5.x indexer prefill of 9,695 new tokens against a 96,960-token context at chunk 8192 drops from 150.2 ms to 118.4 ms. The existing cases top out at s_q=1024, s_k=1560, four orders of magnitude below the gate, which is why nothing caught this. Added 8192x65664 and 8192x98304, which crash the process on current main. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: Mehmet Cagri <mehmet.kaymak@amd.com> Co-authored-by: Jiejing Zhang <jiejizha@smci355-ccs-aus-n02-29.prov.aus.ccs.cpe.ice.amd.com> Co-authored-by: Cursor <cursoragent@cursor.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
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Add a fused split-K opus GEMM pipeline for gfx1250 (a16w16 / bf16) and wire up the
corresponding DSv4 / Kimi-K3 bf16 GEMM tuned configs. The existing split-K path relies on a
separate reduce kernel; for small-M / large-K DSv4 shapes this incurs an extra
kernel launch and an HBM round-trip for the partial sums. This PR introduces a
single-kernel ("splitk_fuse") variant that keeps the partials resident in cache
and reduces them in-cluster, and merges freshly tuned gfx1250 configs so
these shapes pick the fastest available kernel (with a triton fallback for the untuned ones).
Technical Details
opus_gemm_pipeline_a16w16_clusterlaunch_tdm_splitk_fuse_gfx1250.cuh:performs the split-K partial reduction inside a single kernel using a cluster
barrier (instead of a separate reduce kernel + semaphore). Non-last split WGs
cast their fp32 partial to the workspace dtype and store it dirty-resident in
GL2 (CPOL
TH_WB|SCOPE_DEV); the last WG TDM-stages the published partials intoLDS and folds in bias before writing the bf16/fp32 C tile. Workspace defaults to
bf16, with an fp32 option when bf16 reduce precision is insufficient.
splitk_reduce_launch_gfx1250.cuh.splitk_reduce_gfx1250.cuh,opus_gemm.cu,opus_gemm_common.py): support the bf16-workspace + fp32-bias path (fp32 biasfolded in the reduce before the cast to bf16), same kernel name/ABI.
codegen/common.py,codegen/gen_instances_gfx1250.py,gen_instances.py,opus_gemm_tune.py): emit the fused instances and theadditional bf16-out + fp32-bias reduce instantiations across all compile-time
split_k values and HAS_OOB.
opus_gemm_pipeline_a16w16_clusterlaunch_tdm_splitk_ws_gfx1250.cuh,codegen/gen_instances_gfx1250.py,opus_gemm_tune.py): a clusterlaunch kid used torequire the tile grid to be a whole number of clusters, which left most shapes with no
cluster candidate at all. The launcher now rounds the grid up and the surplus workgroups
leave at
tile_oob-- after paying their one cluster-barrier arrival, which-3countsper workgroup, and before issuing any TDM -- so any (CWGM, CWGN) is legal for any (M, N).
This also retires the 2D-cluster lockout: what those hung on was the tile-less workgroup
streaming the whole K loop as zero-extent multicast loads at peers whose extents were
real. The tuner consequently sweeps (cwm, cwn) over [1,4]x[1,4] minus (1,1), bounded by a
round-up budget instead of divisibility.
aiter/configs/model_configs/dsv4_bf16_tuned_gemm.csv,aiter/configs/model_configs/kimik3_bf16_tuned_gemm.csv): retune the full gfx1250block in both files (180 + 345 = 525 shapes) against opus and triton in the same
session and keep the faster backend per shape. dsv4 is then retuned a second time
against the round-up (opus only) -- see the performance section.
Performance
Setup
524 measured (see limitation below).
csrc/gemm_a16w16/gemm_a16w16_tune.py. opus swept 496 AOT-compiled non-fuse kids(28 plain
cluster_tdm_ws+ 468clusterlaunch_tdm_ws) with split-K search; triton usedits own autotune.
shapes already tuned on this branch. The remaining 225 shapes (M>512 or N>10000) are
triton-only and are unchanged by this comparison.
opus vs triton, head-to-head (299 shapes)
opus is faster on 212 / 299 shapes (71%), median speedup 1.37x, mean 1.83x.
The win is largest in the decode regime (M<=32), where the shapes are launch- and
memory-bound and opus lands at a ~5.6-5.9 us floor against triton's
14 us. The gap narrows2048 triton is ahead, which is whyas M grows and both backends become compute-bound; past M
the merged config keeps triton there.
Where triton still wins
triton wins on 87 of the 299 shapes. These are not scattered: they cluster into whole
(N, K) families, hitting nearly every M in the family, so this is a systematic tile/split-K
selection gap rather than measurement noise.
Since the merged config picks the faster backend per shape, these shapes keep the triton
kernel and no shape regresses.
dsv4 retune on the clusterlaunch round-up (180 shapes)
The round-up removes the exact-cluster-fill requirement, so the tuner can offer any
(cwm, cwn) in [1,4]x[1,4] on any shape. dsv4's gfx1250 block was retuned against it,
opus only, over all 180 shapes (179 measured, see limitations) -- including the
N=32320 / N=129280 families that the sweep above had left to triton. Same 496-kid
non-fuse compile set, split-K search, max err_ratio 0.018.
Against the previous config -- which already picked the faster of opus/triton per
shape -- the tuned latencies over those 179 rows sum 1.52x lower (51.6 ms ->
33.9 ms). 112 rows are replaced and 67 keep their old entry; the merge takes the new
measurement only when it is faster, so no row regresses.
The large end is where it pays, because those shapes previously had no usable cluster
candidate and fell back to triton: M=16384 N=32320 K=7168 goes 11.82 ms -> 6.94 ms
(1.70x), M=4096 N=129280 K=4096 5.99 ms -> 3.79 ms (1.58x), M=1024 N=2048 K=7168
55.7 us -> 30.9 us (1.80x). Below M=64 at N>=32320 triton still wins (0.80-0.93x) and
keeps those rows.
Where both configs already chose opus (101 rows) the result is a wash: median 1.016x,
47 faster and 36 slower, all within 10% and all sitting on the 4-5 us launch-latency
floor. The gain comes from the 78 rows that were triton's, where opus now wins 56
(median 1.30x).
152 of the 157 opus rows are clusterlaunch kids, and 99 of them use a (tile, cluster)
pair the previous filter could not select -- a 2D cluster, cwm>2, or (24 rows) a grid
that does not divide the cluster at all. Those 99 rows are a median 1.26x faster than
what the previous config had for them. The selected dims skew wide: 4x4 on 55 rows, 1x4
on 42, 4x2 on 14, 2x4 on 12, the remaining 29 spread over 1x2 / 1x3 / 4x3 / 2x2 / 4x1 / 2x1.
That width is also why the candidate ranking prefers the widest cluster over a product
near 8/16. Measured on grids the cluster fills exactly, which takes the round-up out of
the picture (tile 128x128x128):
At equal size the cluster wants the grid's aspect: on a 4x32 grid 2x4 runs 11.2 us
against 13.3 us for 4x2, and 1x4 (11.4 us) beats 4x1 (20.1 us).
Resulting config
dsv4_bf16_tuned_gemm.csvkimik3_bf16_tuned_gemm.csvdsv4 reflects the round-up retune; kimik3 is from the earlier opus-vs-triton sweep and
has not been retuned against the round-up yet.
All 59 gfx1250 shapes that already carried a tuned entry on this branch improved
(median 1.83x, no regressions) and kept their backend choice. That gain is mostly a
machine-configuration difference between the two tuning sessions rather than a code change:
the five triton-selected rows among them sped up by a comparable amount, and triton is
untouched by this PR.
Limitations
the 496 plain + clusterlaunch kids; the 1378
splitk_fusekids this PR adds were skipped tokeep the build tractable. The numbers therefore validate the retuned configs, not the fused
pipeline itself, which is covered by the functional test below. A fuse-inclusive sweep is
follow-up work.
with an identical build selected the same kernel for 290/290 shapes but shifted latencies by
a median of 4-6%, with 32% of shapes moving more than 10%. Kernel selection is reproducible;
individual speedup figures should be read with that error bar. The aggregate results are well
outside it.
the correctness check (err_ratio 0.49) -- its a16w16 kernel computes output byte offsets in
int32 and this output is 4.2 GB, so addressing wraps past 2^31 bytes, with corruption
starting exactly at row 8305 = 2^30 / 129280. opus has no candidate either: the output is
3.95 GiB and every kid is rejected by the 4 GiB buffer-resource filter, since these kernels
build 32-bit AMDGPU buffer resources. The row is left as an untuned
autofallback.config's entry rather than against a same-session triton run. The 78 rows that were
triton's carry that caveat; the triton numbers they are compared with come from the sweep
above, same machine and build.
Test Plan
op_tests/opus/test_splitk_fuse.py(kid 21000) against aCPU fp32 reference (torch GPU bf16 matmul is unreliable on the current gfx1250
ROCm build), covering bf16/fp32 output and bias/no-bias:
run(64, 512, 2048, bf16, bias=False/True)run(64, 512, 2048, fp32, bias=False)run(128, 1024, 4096, bf16, bias=False)every tuned row passes the tuner's correctness check (max err_ratio 0.042 for opus).
scenarios each (full / M-partial / N-partial / 1-of-16 alive / ragged / split-K),
plus 20 consecutive sweeps at 4x4 -- no hang, no mismatch, no writes outside the
tile from workgroups that leave at
tile_oob.tuner's correctness check, max err_ratio 0.018.
Test Result
bad=0) against the CPU reference for bothbf16/fp32 output and with/without bias.
against the previous triton-fallback baseline.