feat(tensilelite): StreamK WG-cluster reduction + cooperative loads for gfx1250#9420
Draft
jaopaulolc wants to merge 3 commits into
Draft
feat(tensilelite): StreamK WG-cluster reduction + cooperative loads for gfx1250#9420jaopaulolc wants to merge 3 commits into
jaopaulolc wants to merge 3 commits into
Conversation
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
Codecov Report✅ All modified and coverable lines are covered by tests. ❌ Your project status has failed because the head coverage (76.84%) is below the target coverage (80.00%). You can increase the head coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## develop #9420 +/- ##
===========================================
+ Coverage 64.60% 64.61% +0.01%
===========================================
Files 2688 2688
Lines 422872 422802 -70
Branches 62766 62750 -16
===========================================
- Hits 273189 273188 -1
+ Misses 128995 128927 -68
+ Partials 20688 20687 -1
*This pull request uses carry forward flags. Click here to find out more.
🚀 New features to boost your workflow:
|
jaopaulolc
added
the
rocm:gemm algos
jaopaulolc
force-pushed
the
users/jaopaulolc/aihpblas-3929-streamk-wg-clusters
branch
from
July 15, 2026 19:24
d20dcde to
bbfc19a
Compare
jaopaulolc
changed the base branch from
develop
to
users/nhenders/cluster-streamk-gfx1250
July 15, 2026 19:24
jaopaulolc
changed the title
jaopaulolc
changed the base branch from
users/nhenders/cluster-streamk-gfx1250
to
develop
July 15, 2026 19:39
jaopaulolc
changed the title
jaopaulolc
changed the title
jaopaulolc
force-pushed
the
users/jaopaulolc/aihpblas-3929-streamk-wg-clusters
branch
5 times, most recently
from
July 15, 2026 23:59
6f3ede9 to
c6be5b7
Compare
jaopaulolc
changed the title
jaopaulolc
force-pushed
the
users/jaopaulolc/aihpblas-3929-streamk-wg-clusters
branch
9 times, most recently
from
July 17, 2026 19:54
7edc2d5 to
009e5e4
Compare
JIRA ID: AIHPBLAS-3929 StreamK GEMM tiles whose partial results are reduced across workgroups (WGs) synchronize that reduction with a cross-CU global-workspace flag spin-wait. On gfx1250, the WGs that reduce one tile can instead be co-located into a hardware cluster and synchronized with a cluster split barrier, removing the spin-wait from the fast path. Adds a new opt-in solution parameter StreamKClusterReduction. When enabled, a StreamK tile's fixup peers are placed in a single 1-D workgroup cluster (ClusterDim = [C, 1]); the HW WG-id remap (WorkGroup0 = cluster_x*C + wg_x) makes each cluster own the contiguous StreamK index range [c*C, c*C+C), i.e. exactly one tile's peers. - Host: getSKGridImpl rounds skGrid to C*tiles (fixed even split, SKItersPerWG = itersPerTile/C) and keeps the launch grid a multiple of C; the new field is threaded through SizeMapping / ContractionSolution (de)serialization. - Kernel: new clusterReduceSignal / clusterReduceWait / clusterReduceIntraCheck helpers in StreamK.py emit the wave-0-elected s_barrier_signal/s_barrier_wait -3 handshake in the fixup epilogue. A uniform intra-cluster predicate selects the barrier fast path or the retained global-flag path so a cluster never mixes the two (deadlock-safe). The ttmp9 WG-id workaround is skipped under clustering so StreamKIdx stays the global index. - Solution.py keeps Multicast and the mainloop ClusterBarrier off for a StreamK cluster: its peers iterate different K-splits and never march the mainloop in lockstep, so splicing the cooperative subtile ClusterBarrier there would mismatch the cross-WG barrier counts and deadlock. MVP scope: StreamK variant 3 (two-tile) only; barrier-only (no Multicast / cooperative-load reuse); linear-reduction fast path only (tree reduction keeps the global-flag path); fixed even split (one tile per cluster); SCOPE_DEV fences retained. The global-flag reduction is retained as a runtime/compile fallback. Design rationale in docs/design/streamk-wg-clusters.md. - Unit (CPU asm-string): test_streamk_cluster_reduction.py. - Snapshot codegen characterization: test_streamk_cluster_reduction_gfx1250_char.py with a designed gfx1250 config and syrupy golden. - On-FFM gfx1250 GPU roundtrip: test_streamk_cluster_reduction_gpu.py (@requires_gpu_gfx1250, SIGALRM deadlock watchdog). - C++ client GEMM: sk_mxf8gemm_cluster_reduction.yaml (real multi-WG cluster StreamK launch). Validated on the FFM gfx1250 simulator: 60 pytest passed, 26 C++ gtests passed, baseline StreamK client GEMM 31/31 validated, and a real multi-WG cluster StreamK GEMM validated with cluster_size=4x1x1 (no deadlock). Hardware sign-off pending. - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. Low. The feature is off by default and strictly opt-in via StreamKClusterReduction with a narrow validation predicate (SK3 + gfx1250 + [C,1]); the global-flag reduction path is retained unchanged as the runtime fallback.
…ulticast
JIRA ID: AIHPBLAS-3929
The TDM multicast ("cluster load") mask machinery was duplicated inline across
KernelWriter, KernelWriterAssembly, and Subtile/SubtileGREmit, and Multicast was
implicitly coupled to ClusterDim != [1, 1]. That coupling made it impossible to
compose barrier-only clustering and cooperative-load clustering independently, or
to reuse the mask attach at a new load site.
- New ClusterLoad component (Components/ClusterLoad.py, registered via
Component.py + Components/__init__.py). ClusterLoadTDM centralizes the mask
value computation (computeMasks), the MulticastMask* SGPR declare/undeclare
(declareSgprs / undeclareSgprs), the topology decision (combined vs split
A/B masks), and the descriptor attach at each load site (applyToDescriptor).
It is a behavior-preserving extraction: every method emits byte-identical
assembly to the previous inline code, and computeMasks receives the exact
SGPR operands the caller already holds so register indices are unchanged.
Selection is capability-based (HasTDM + TDMInst == 3), like TensorDataMoverLoad.
- KernelWriter / KernelWriterAssembly / SubtileGREmit now route mask declare,
undeclare, compute, and descriptor attach through ClusterLoadTDM.find(...).
- Tri-state Multicast parameter (ValidParameters + GlobalParameters default -1):
-1 = auto (legacy: ClusterDim != [1,1] implies Multicast, minus the StreamK
cluster paths), 0 = force off, 1 = force on. Default -1 reproduces the historic
derivation exactly, so every existing YAML (which omits Multicast) derives
byte-identically.
- Solution.py derives Multicast/ClusterBarrier from the tri-state, keying the
legacy auto coupling and the ClusterBarrier gate on StreamK == 0.
- Common/Utilities.clusterEnabled() helper for the ClusterDim != [1,1] test.
- Unit: test_cluster_load_component.py (component find/declare/compute/attach),
test_multicast_tristate.py (-1/0/1 derivation, legacy equivalence).
Low. Behavior-preserving refactor; the tri-state defaults to -1 (legacy auto),
so emitted assembly is unchanged for all existing configs.
…icast) JIRA ID: AIHPBLAS-3929 StreamK data-parallel (DP) workgroups that process M-adjacent tiles share the same B (N-block) tile over the full K range, yet each reloads B from global memory. On gfx1250 those consecutive-WG [C,1] cluster peers can load B once and TDM-multicast it to all C members, cutting redundant B traffic in the DP region. Adds the opt-in StreamKMulticast solution parameter (built on the ClusterLoad component + tri-state Multicast). - Kernel (StreamK.py, StreamKTwoTileDPFirst): the B-broadcast mask is gated at runtime by clusterMulticastValid (streamKMulticastMaskPredicate) -- the cluster must be M-aligned (nWG0 % C == 0) and fully populated (clusterBase + C <= totalTiles); otherwise B falls back to a self-only (normal) load, so results are correct for all sizes and a masked target is never left without a matching load. At the DP->SK boundary streamKMulticastBoundaryClear rewrites MulticastMaskB to the self-only mask (SK partial-tile peers no longer co-issue the identical full-K B load). The SK3 preLoop ttmp9/ttmp7 reread is skipped under clusterEnabled(ClusterDim), keeping StreamKIdx the global index. - Solution.py: explicit StreamKMulticast branch turns Multicast on (superseding the legacy auto-off for StreamK), plus _validateStreamKMulticast (SK3, mutually exclusive with StreamKClusterReduction, no StreamKAtomic, XCCMapping=0, ClusterDim=[C,1] with C pow2 in [2,16], gfx1250 HasTDM/TDMInst!=0). - Host: streamKMulticast threaded through SizeMapping / ContractionSolution (de)serialization; getSKGridImpl rounds the launch grid up to a multiple of C (spatial single-round DP), and the cluster grid multiple-of-C guard now covers StreamKMulticast too. - Unit (CPU asm-string): test_streamk_multicast.py. - Snapshot codegen characterization: test_streamk_multicast_gfx1250_char.py with a designed gfx1250 config (streamk_multicast.yaml) and syrupy golden. - C++ client GEMM: sk_mxf8gemm_multicast.yaml. Design rationale in docs/design/cluster-load-component-and-streamk-multicast.md. Low. Off by default and strictly opt-in via StreamKMulticast with a narrow validation predicate; the cooperative B load is additionally gated at runtime by clusterMulticastValid, falling back to normal per-workgroup B loads otherwise.
jaopaulolc
force-pushed
the
users/jaopaulolc/aihpblas-3929-streamk-wg-clusters
branch
from
July 17, 2026 20:07
009e5e4 to
0eed89b
Compare
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
2 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.
JIRA ID : AIHPBLAS-3929
Motivation
StreamK GEMM tiles whose partial results must be reduced across workgroups (WGs) synchronize that reduction with a cross-CU global-workspace flag spin-wait. On gfx1250 the WGs that reduce one tile can instead be co-located into a hardware WG cluster and synchronized with a cluster split barrier, removing the spin-wait from the fast path while keeping the global-flag path as a fallback. The same cluster co-location also enables cooperative in-cluster data movement: one WG loads a shared operand tile and TDM-multicasts it to its cluster peers, cutting redundant global loads of the B operand across the data-parallel (DP) region. This PR builds on the 1-D WG-cluster support for Stream-K on gfx1250 (#9464), which is now merged into
develop.Technical Details
Three commits, each building on the previous:
(a) StreamKClusterReduction — cluster-barrier reduction. Adds the opt-in solution parameter
StreamKClusterReduction. A StreamK tile's K-split reduction peers are placed in a single 1-D WG cluster (ClusterDim = [C, 1]); the HW WG-id remap (WorkGroup0 = cluster_x*C + wg_x) makes each cluster own the contiguous StreamK index range[c*C, c*C+C)— exactly one tile's peers — which are reduced via a cluster split barrier (s_barrier_signal -3/s_barrier_wait -3). The kernel (Components/StreamK.py) emits the wave-0-elected handshake in the fixup epilogue, gated by a single uniform intra-cluster predicate (a pure function of the cluster's top StreamK index and the SK grid size) so a cluster can never mix barrier-signalers with flag-setters — the deadlock-safety invariant. The global-flag path is retained as the runtime fallback. Host accounting uses a fixed even split (skGrid = C * tiles,SKItersPerWG = itersPerTile / C) and keeps the launch grid a multiple ofC; host selection rejects the config via theClusterReductionIterCheckpredicate (emitted fromContractions.py) whenceil(K / DepthU) % C != 0(the peers do not split evenly). A workspace-fallback guard inContractionSolution.cppsets the cluster-reduction kernargs only whenskGrid == C * tiles, so a fallback grid reverts cleanly to the global-flag path.(b) Reusable ClusterLoad component + tri-state Multicast. Extracts the cooperative cluster-load (TDM multicast-mask) machinery that was duplicated inline across
KernelWriter,KernelWriterAssembly, andSubtile/SubtileGREmitinto a reusable tensilelite component (Components/ClusterLoad.py,ClusterLoadTDM), centralizing mask computation,MulticastMask*SGPR declare/undeclare, the combined-vs-split A/B topology decision, and the per-load-site descriptor attach.Multicastis generalized from a bool into a tri-state parameter[-1, 0, 1]decoupled fromClusterDim(-1= auto/legacy derivation,0= force off,1= force on). Derivation now emits an INTMulticast(0/1) on every path (post-derivation coercion removed); a pre-loadcoerceLegacyMulticastTypenormalizes the legacy boolMulticastcarried by shipped library YAMLs to the tri-state int, since msgpack serializes bool and int as distinct wire types. This is a behavior-preserving, snapshot-gated refactor: default-1reproduces the historic derivation byte-for-byte (zero golden deltas) for every existing config.(c) StreamKMulticast — StreamK cooperative TDM B-multicast. Adds cooperative TDM B-multicast for the DP region built on the ClusterLoad component.
StreamKMulticastis a DERIVED-ONLY internal state key (likeClusterBarrier), NOT a user-settable/benchmark parameter: onStreamK == 3, a non-[1,1]ClusterDimwithoutStreamKClusterReductionauto-enables it (the "bare StreamK cluster" collapse — "ClusterDim without cluster loads or reduction" is no longer a supported state, so this is the sole enable site). The DP cluster is spatial[C,1]peers withskGrid = ceil(tiles/C) * C; within a cluster one WG performs the TDM B load and multicasts it to its peers, and at the DP→SK boundary the multicast mask is rewritten to a self-only mask so the StreamK region reverts to independent per-WG loads. The path is conservative — a runtimeclusterMulticastValidpredicate (M-alignednWG0 % C == 0and fully populatedclusterBase + C <= totalTiles) disables the cooperative load whenever correctness is not guaranteed, falling back to independent loads._validateStreamKMulticasthard-rejects any derived config it cannot satisfy (SK3 only, mutually exclusive withStreamKClusterReduction, noStreamKAtomic,StreamKXCCMapping = 0,ClusterDim = [C, 1]withCa power of two in[2, 16], gfx1250HasTDM/TDMInst == 3), rejects the config where it would be auto-enabled butMulticast = 0, and SK4/SK5 (dynamic/hybrid) with a non-[1,1]ClusterDimare rejected outright — cluster support is SK3-only.MVP / scope boundaries. StreamK variant 3 (two-tile DP-first) focus; linear-reduction fast path (tree reduction keeps the global-flag path); fixed even split, one tile per cluster for the reduction;
SCOPE_DEVfences retained (no cluster-scope relaxation); every new path is opt-in or derived behind a narrow validation predicate with the pre-existing behavior retained as the runtime fallback. All validation to date is on the FFM gfx1250 simulator — there is no real-silicon validation yet and hardware sign-off is pending.Test Plan
test_streamk_cluster_reduction.py); ClusterLoad component find/declare/compute/attach (test_cluster_load_component.py); tri-state Multicast derivation and legacy equivalence (test_multicast_tristate.py,test_multicast_legacy_coercion.py); StreamKMulticast DP multicast + DP→SK boundary mask-clear (test_streamk_multicast.py); SK4/SK5 + ClusterDim reject (test_streamk_cluster_sk45_reject.py)..ambr): reduction-3handshake and global-flag fallback (test_streamk_cluster_reduction_gfx1250_char.py); cooperative-load extraction byte-exactness (test_streamk_cluster_coop_load_gfx1250_char.py); StreamKMulticast codegen (test_streamk_cluster_multicast_gfx1250_char.py).test_streamk_cluster_reduction_gpu.py) and C++ client GEMM underTensile/Tests/common/streamk/gfx1250/core/—sk_mxf8gemm_cluster_reduction.yaml/sk_mxf4gemm_cluster_reduction.yaml(F8+F4 reduction) andsk_mxf8gemm_cluster_multicast.yaml/sk_mxf4gemm_cluster_multicast.yaml(F8+F4 multicast). The obsolete bare-cluster test fixtures were removed.Test Result
Validated on the FFM gfx1250 simulator only (no real silicon; hardware sign-off pending):
.ambr) deltas confirming the (b) extraction is byte-exact for existing configs.Submission Checklist
Risk level
Low. Every layer is off by default or derived only under a narrow validation predicate (SK3 + gfx1250 cluster-barrier +
[C,1]for the reduction; the SK3ClusterDimcollapse plusclusterMulticastValidfor the cooperative multicast). The (b) refactor is snapshot-proven byte-exact, and the pre-existing global-flag reduction and independent-load paths are retained unchanged as runtime fallbacks. Validation is on the gfx1250 simulator only; hardware sign-off is still pending.