Add gfx1250 a8w8 mxscale BMM scaffold with preshuffled B. - #5068
Draft
yzhou103 wants to merge 4 commits into
Draft
Add gfx1250 a8w8 mxscale BMM scaffold with preshuffled B.#5068yzhou103 wants to merge 4 commits into
yzhou103 wants to merge 4 commits into
Conversation
Introduce a TDM 2-producer/2-consumer batched GEMM path for MI450 that consumes shuffle_weight(16,16) weights, uses DS V4 1x128 e8m0 blockscale (pack_e8m0x4 broadcast), and wires host/pybind/JIT plumbing plus a probe test for frag layout validation on gfx1250 hardware. Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags: |
Co-authored-by: Cursor <cursoragent@cursor.com>
…sweep found
The raw binding takes a kernelId and defaults it to 0, the 128x128 prefill tile.
On the DSV4 wo_a decode shapes that default costs 1.44x-2.24x. Add
_heuristic_bpreshuffle_kid and a bmm_a8w8_mxscale_bpreshuffle_opus entry point
shaped like its bmm_a8w8_mxscale_opus sibling, so a caller gets the tile the
shape wants without knowing the kid table.
The rule is a workgroup count, not an M threshold, because that is what both
sides of the trade key on. Below the CU count time FALLS as the grid shrinks --
64 -> 32 -> 16 workgroups all get faster -- since a narrower B_N means more
workgroups each re-reading the same B_M=16 rows of A and the same per-WMMA
scales, and that duplication, not occupancy, is what binds. Above it the grid
stops buying anything and kid0's B_M=128 takes over. Fitted on batch 1..16 x
m 1..256 at n=1024 k=4096: names the measured winner in 30 of 35 cells, the
five misses all near-ties of 0.5%-2.3% against a sweep whose own agreement with
rocprofv3 is 2.7%. Decode tiles are confined to m<=256, the region actually
swept -- extrapolating the rule past it put batch=2 m=512 on a tile 1.52x off.
This inverts the premise the decode tiles were added under. They exist because
kid0 "leaves 94% of the CUs idle", i.e. to raise the workgroup count; kid1, the
16x32 tile that premise produced, wins none of the 35 cells and is 6x off the
best at batch=16 m=256. Fewer and fatter, not more and narrower.
None of this is visible in wall time. A host dispatch costs ~8 us and these
kernels are 9-16, so an event-timed loop reads every decode tile at 19.2 us to
three digits; on kernel time they span 1.53x. Every earlier decode conclusion
taken from wall time is blind, including the recorded B_N ladder.
The header changes are comment-only (the .so is byte-identical) and record two
further results:
* the cluster-launch variant, previously unmeasured. mClusterWg=2 -- the B
multicast across M-tile peers -- is worth 3.4%-6.0% on prefill at SplitK=1.
SplitK itself does not pay: on prefill it costs ~2.5x per doubling because
the SplitK=1 grid already fills the machine while kid0's partial tile is a
full B_M x B_N fp32, and on decode entering the cluster path costs a flat
~21% that nothing wins back.
* a re-measurement that CONTRADICTS the scale-panel notes in the traits
header. kid13 is neutral and kid14 costs 11%-12% where the file records
+8.1%; and the 38.60 us kid0 that file discards as unreproducible is what
reproduces now, against the ~48 us it says both compilers give. Since the
panel's claimed win is a ratio against that baseline the two are probably
one thing. Recorded, not resolved: nothing above it has been rewritten and
no conclusion has been deleted, because the cause has not been found.
Gates: op_tests/test_opus_a8w8_bmm_bpreshuffle_gfx1250.py 21/21,
op_tests/test_opus_a8w8_bmm_bpreshuffle_cc_gfx1250.py 21/21, and the dispatched
path bit-exact against kid0 at ten shapes spanning all three tiles.
Co-Authored-By: Claude Opus 5 (1M context) <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.
Introduce a TDM 2-producer/2-consumer batched GEMM path for MI450 that consumes shuffle_weight(16,16) weights, uses DS V4 1x128 e8m0 blockscale (pack_e8m0x4 broadcast), and wires host/pybind/JIT plumbing plus a probe test for frag layout validation on gfx1250 hardware.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist