Back the FP8 rowwise grouped GEMM ops with FlyDSL on ROCm - #471
Open
aryaman-gupta wants to merge 39 commits into
Open
Back the FP8 rowwise grouped GEMM ops with FlyDSL on ROCm#471aryaman-gupta wants to merge 39 commits into
aryaman-gupta wants to merge 39 commits into
Conversation
The kernel only knew block scaling, which carries a scale per (scale_block_k x scale_block_n) block and so must scale each block's partial sum inside the K loop. Rowwise scaling instead carries one scale per row of A and per column of B, both constant along K, so they factor out of the reduction. Select the scheme with a compile-time `blockscale` flag. When it is off the K loop accumulates unscaled and a new epilogue pass applies the two scales once, which also frees tile_n of scale-block alignment and lets it go below 128. tile_k keeps its divisibility check under both schemes: scale_block_k sets the K-loop sub-block count, so a smaller tile_k yields zero sub-blocks and a compute loop that never runs. Co-Authored-By: Claude <noreply@anthropic.com>
Wraps FlyDSL's autotuner in the policy the other MSLK backends use: tune only when MSLK_AUTOTUNE_ENABLE is set, otherwise take a fixed default config with no benchmarking. Both halves matter -- benchmarking on a cache miss would make CI pay for a sweep, and it cannot happen inside a CUDA graph capture. Kernel modules supply their own launch function, candidate configs and cache key, so nothing here is tied to a kernel or an op. Also included are the pieces every caller would otherwise repeat: power-of-two bucketing for cache keys, a divisibility-based config pruner, and the workaround for FlyDSL's autotuner discarding the tuned function's return value. Co-Authored-By: Claude <noreply@anthropic.com>
One kernel serves every combination of B layout and scaling scheme, so the host-side work around it -- operand marshalling, grid extent, tile selection -- is the same for all of them. Move it into grouped_dispatch so each op module is left with only its own contract checks, and route tile selection through mslk.flydsl.autotune rather than a private copy of that plumbing. The candidate tiles are now per-scheme: rowwise scaling admits tile_n below the scale block, and both schemes gain tile_k=256, which the previous list omitted even though it wins on several shapes. Co-Authored-By: Claude <noreply@anthropic.com>
Registers mslk::f8f8bf16_rowwise_grouped_stacked against the shared grouped kernel compiled for rowwise scaling, plus a preshuffle sibling for callers that swizzle their weights into the MFMA B layout once at load time. CK served the stacked op on ROCm. Its GPU registration is now guarded to non-ROCm so the slot is free for the Python binding; the same registration serves CUTLASS on CUDA, hence a guard rather than a deletion. The CPU registration moves alongside f8f8bf16_groupwise_grouped, which likewise has no CPU fallback because it is dispatched from Python. Matches CK to within one or two bf16 ulps, and is faster on every shape swept across DS-V3, llama3_70b and llama3_405b except at total_M=1, where most groups are empty. Both the C++ and Python halves land together: either alone would leave the op unimplemented on ROCm. Co-Authored-By: Claude <noreply@anthropic.com>
f8f8bf16_rowwise_grouped_dynamic had no test. Add one alongside the stacked variant, covering empty, partially filled and completely filled groups, since the padded layout only defines the first zero_start_index_M rows of each slab. The reference is the existing bf16 loopover over the valid rows; the padded rows are asserted zero separately, which is the part of the contract zeroing_output_tensor controls. Co-Authored-By: Claude <noreply@anthropic.com>
The rowwise grouped GEMM has a second layout: instead of packing groups densely along M, each group owns a fixed slab of expected_m rows of which only the first zero_start_index_M[g] carry real tokens. Callers scatter into per-group slabs rather than compacting into one buffer, which keeps shapes static and the row counts on device. Because the slabs are contiguous the kernel can take the flattened [G * expected_m, ...] views, so every buffer size, loader and the epilogue are unchanged. Only two things differ, behind a compile-time flag: the group comes from a third grid axis rather than being resolved from row counts, and the row base is a fixed multiple of expected_m. The existing row-limit mask then keeps the padding rows unwritten, and the per-group scale_a addressing already in place for the packed layout turns out to be exactly what [G, scale_k, M] needs. Adds matmul_f8f8bf16_rowwise_grouped_dynamic and its preshuffle sibling. Neither is registered yet: CK still serves the op, and displacing it is a performance question settled separately. Co-Authored-By: Claude <noreply@anthropic.com>
The K loop is unrolled at compile time over whole tiles, so a K that stopped mid-tile was rejected. That narrowed the ops well below what they replace: CK accepts any K, and requires only that N divide its store width. Compile a tail-masked variant when K stops mid-tile and select it on the host, mirroring how CK chooses between its KPadding and Default specialisations. The mask is emitted only for the final tile, so every other tile and every shape where K divides keeps exactly the loads it had before. Masked buffer loads read out of range and return zero, which is what the tail needs; letting the offset run past K instead would read the next row, since that is still inside the buffer. K now only has to reach a 16-byte load boundary rather than a whole tile. Preshuffled B keeps the stricter rule, its swizzled layout interleaving K so a partial tile cannot be masked one load at a time. Co-Authored-By: Claude <noreply@anthropic.com>
N had to divide tile_n, and the MFMA tiling only admits tile_n in multiples of 64, so the ops rejected shapes CK accepts at a granularity of 8. Mask the tail the same way the K tail is handled: round the N grid up, drop B rows past N on load, and suppress epilogue stores whose columns fall past N. Both are predicated buffer accesses, so a masked load reads zero and a masked store does nothing, and the predicates are only emitted for the padded variant. N now only has to reach the widest vectorised store rather than a whole tile. As with K, preshuffled B keeps the stricter rule, its swizzled layout interleaving N so a partial block cannot be masked one access at a time. Co-Authored-By: Claude <noreply@anthropic.com>
Tail masking made a tile that overruns N or K valid rather than an error, so the prune had been relaxed to stop rejecting them. That was too broad: such a tile spends part of every iteration on padding, so it is not going to win, and benchmarking it only lengthens tuning. Prune on both axes again. The case that needs padding is the one where nothing divides, and there prune_by_divisibility falls back to the full list, so those shapes are still tuned over every candidate. Co-Authored-By: Claude <noreply@anthropic.com>
The kernel learns where each group's rows live in more than one way, and a boolean named for a single alternative cannot express a third. Replace the masked flag with a layout enum, so the packed and per-group-slab encodings are peers and adding another does not multiply booleans that cannot be combined. Behaviour is unchanged: masked=False becomes layout="sizes" and masked=True becomes layout="padded". Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
A grouped GEMM whose groups all carry their full complement of rows is an ordinary batched GEMM, which is the padded layout with the row counts implied rather than supplied. Reading them from a [G] tensor would make the caller materialise and fill one on every call, so the layout derives the count from the slab height instead and never touches the metadata operand. Rows still have to be masked. A full slab does not mean a tile fits inside it: tile_m need not divide the slab height, and the overrun lands on the next group's rows, racing with the group that owns them. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The torch-native grouped GEMM API hands the group boundaries over as an inclusive prefix sum in int32 rather than as int64 per-group row counts. The resolution loop already carries that prefix, so a group's row count is the step between successive offsets and decoding costs one subtract. Differencing the offsets host-side would need its own kernel launch on every call, which is the wrong trade for an op whose weak point is decode latency. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
f8f8bf16_rowwise_grouped_mm picks its grouped axis from the operand ranks. Two of the four combinations are the layouts the kernel now has: 2D-3D is the packed layout addressed by offsets, and 3D-3D is the batched one. The op stays unregistered. It is a single entry point covering all four combinations, so claiming it before 3D-2D and 2D-2D exist would break the shapes they cover, and registering at all puts CK out of reach as the oracle these were checked against. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The packed layouts launch a host-known upper bound of M-tiles and the padded one can have fewer valid rows than its slab, so both need a tile to be able to retire itself. The batched layout has neither property: its grid is exactly ceil(slab / tile_m) tiles and every row is valid, so the guard is always taken and can be left out of the kernel rather than emitted and folded. Worth 0.4-1.2% on four of five shapes measured under graph replay; a fifth is too noisy run-to-run to call. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Working out which group owns a tile is the only part of the kernel the layout changes; everything after it consumes the same coordinates. Having all four encodings inline left the kernel body opening with sixty lines of branching before any of the work it actually describes. No behaviour change: the final ISA is byte-identical across all eight configurations that share this kernel -- both scaling schemes, both B layouts and all four layouts. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The torch-native grouped GEMM also comes in a form where the weights are one [total_N, K] matrix that the groups partition by row, and the results are packed side by side into [M, total_N]. The work per group is unchanged; only which axis is ragged, and therefore where the group has to be resolved from. Resolution moves to the N-block index, mirroring the M-side prefix scan. Every N-side address was already a runtime base plus a local offset, so the rest is substitution: B loses its per-group stride and is indexed by the global column, scale_b likewise, and the tail masks bound against the group's column end rather than N. The output's leading dimension stops doubling as that bound, since a row now spans every group. A and scale_a keep addressing the flattened per-group slabs while the epilogue indexes rows within one slab, which the two taking bx_m separately allows. The tail mask is always emitted here: a group's column end is a runtime value, so it cannot be elided the way a compile-time N remainder can. Rowwise plain-B only. Ragged N would need per-group scale blocks, and the preshuffled layout swizzles N across the whole matrix. The final ISA is unchanged for all eight pre-existing configurations. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Three of the four rank combinations are now the kernel's own layouts. The weights arrive 2D here, so the group count comes from the offsets rather than from a leading weight dimension, and tuning keys on one group's shape rather than the concatenation -- the same normalisation CK applies before its heuristic. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The K loop is unrolled over every tile, which folds the tile index into each address and leaves the pipeline free of loop overhead, but makes the traced IR, the compile time and the emitted code all grow linearly with K. At K=8192 that is 5271 lines of ISA and 13s to compile, against 990 lines and 1.5s rolled -- and the rolled figures do not move with K at all. The rolled form keeps the pipeline identical, issuing the next tile's global loads before the current tile computes, so the prefetched registers ride the loop beside the accumulators; the final tile is peeled, having no successor to prefetch. Output is bitwise identical to the unrolled loop on every shape checked, so the reduction order is unchanged. It lives in the kernel body rather than beside the other K-loop helpers because only the kernel function's own source is AST-rewritten, and the loop construct is a rewritten one. The flag is off everywhere: the eight configurations that share this kernel still emit byte-identical ISA. Nothing yet decides when rolling should be preferred; that needs its own analysis. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Rolling is the framework's own default and what both shipped grouped-GEMM libraries do, so it is the better resting position for a kernel that JITs per shape: compile time and code size stop tracking K. It stays out of the autotune key. A tuning space holding a fully unrolled candidate would compile one per tile config at a cost that grows with K, which is why no surveyed kernel offers one -- the unroll factors that are tuned elsewhere are small fixed constants. Revisit when the factor itself becomes a tuning dimension. The preshuffled path still runs the two-deep ping-pong loop, which is unrolled, so the dispatch resolves the flag away there rather than the caller having to. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The preshuffled path runs its own K loop, staging A through a ping-pong pair because B needs no LDS, and it was the one left unrolled. It rolls a pair of tiles at a time, which keeps the buffer alternation compile-time inside the body while the loop itself rolls, and peels the pairs that have no successor to prefetch so the rolled body carries no tail tests. The B tile and the first A pack ride the loop beside the accumulators, being produced a pair ahead of use. At K=8192 that is 1199 lines of ISA and 1.8s to compile, against 4635 lines and 13.9s unrolled, and the rolled figures are flat in K as on the plain path. Output is bitwise identical to the unrolled loop for K-tile counts 1 through 8, covering both parities of the pair count. The per-pair body is now shared rather than written once per loop form, since only the loop construct has to live in the kernel's own source. With the flag off all sixteen configurations still emit byte-identical ISA. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The last of the four rank combinations groups along K: both operands are one matrix that the groups slice by column, and every group contracts over its own slice to produce a whole output, so nothing is packed and the result is a slab per group. Which group owns a tile is structural here rather than ragged -- every group produces a full output, so the group is simply a grid axis and only the ends of its K slice have to be read. The raggedness moves instead into the trip count, which is why the loop had to roll first: a group's K length is only known on the device, so this layout always rolls. A is shared by every group while scale_a and the output are per group, so the three take different row bases; B is one matrix with no per-group row base, while its scales still have one. Each group's K must be a multiple of the vectorised load width, which CK checks nowhere and is silently wrong below. A group contracting over nothing is skipped and its output left to the caller, as CK leaves it. Matches CK across uniform and ragged group lengths, non-tile lengths, all-16 lengths, and empty groups. On [512, 16] at N=512 CK leaves the second group entirely unwritten; this does not. The eight pre-existing configurations still emit byte-identical ISA. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Preshuffled B is held in registers across a whole pair of K tiles to cover HBM latency, where the plain path stages it through LDS and can drop it after the ds_write. That longer live range costs about 25 registers at the same tile, and at a wide tile_n it lands the kernel on exactly 256 -- the boundary between two waves per SIMD and one. A rolled K loop needs a few registers more than straight-line code, since its carried values are all live at the back-edge, and that was enough to fall off the edge: 264 to 297 registers, one wave, and a runtime 1.4 to 1.6x worse than unrolled at the shapes where it bit. waves_per_eu was already a parameter of the kernel and never used. Setting it to 2 fits the allocator back under 256 and closes that gap: at 16384x7168x2304 and a fixed 64x256x128 tile, rolled goes from 1.61x unrolled to 1.03x. It is a per-shape choice rather than a constant -- autotune picks 2 on some shapes and leaves it to the compiler on others -- so it joins the tile in the config space. Only two values are worth sweeping: three waves would need 170 registers and four would need 128, which no configuration of this kernel approaches. The hint is in the module name, since two configs that differ only in it are different kernels. Unlike an unroll factor this costs nothing extra to compile -- same trace, same instruction count, only a different allocator target -- which is affordable precisely because the loop now rolls. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Where the groups divide K, the tail predicate was gated on whether the total K divides the tile. That says nothing about whether a group's slice does: with a total of 2304 and eight groups of 288, the total is tile-aligned and the slice is not, so no predicate was emitted at all and every group read past its own K end into the next one's data. Wrong on two of six shape families measured and on every one of 34 tile configurations, by roughly twenty times the FP8 residual -- large enough to be plainly wrong, small enough to pass for quantisation error in a column of its own. It showed up as a divergence from CK rather than as an implausible number. Inert for the other layouts: they pass no bound, so the expression short-circuits to what it was. An earlier ragged-K sweep missed this because every case in it happened to have a total K that was not tile-aligned, so the predicate was emitted by accident. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The existing grouped mm tests size every group to a multiple of 64, or hold it uniform, so a group's extent always divided the tile in practice. That is the one property the kernel is most likely to get wrong, and two bugs reached the benchmarks through the gap. 2D-2D keeps the total K tile-aligned while each group's slice is not, which is what breaks a tail predicate keyed on the total rather than on the slice. 3D-2D does the same on N, including a group narrower than one tile. 3D-3D sweeps slab heights that do not divide the M tile, where the last tile overhangs into the next group's rows. Group extents stay within what both implementations support -- a multiple of 16 along K, of 8 along N -- so these say nothing about the shapes below that, where CK is silently wrong. Adds 12 cases and 14 seconds. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The dynamic and mm ops now have a FlyDSL implementation for every shape they accept, so bind them alongside the two that were already bound. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
mslk/gemm/flydsl/ holds the op wrappers for the FlyDSL GEMM kernels but was not in the ROCm trigger paths, so a change confined to it would skip ROCm CI.
The comparison tolerance is absolute at these magnitudes, and a group that reads a little past its own K end sums extra terms that cancel, leaving an error of the same order as the quantisation noise. The check passed either way, so it said nothing about the boundary. Give every other group constant, single-signed data. Extra terms then accumulate coherently and the error lands orders of magnitude above the tolerance, which is what makes the case worth running. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The dynamic and mm ops kept a CK implementation on the CUDA key, so the Python binding took effect by overriding it, which torch warns about. Leave those slots free on ROCm the way the stacked op already does, so the binding is the only implementation rather than the winning one. Declare a schema for the padded layout's preshuffle sibling as well; it had none, so it was reachable only as a plain function. CK still serves the Tensor[] variants, which have no FlyDSL counterpart, and every one of these ops on CUDA. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The kernel already accepts the MFMA B layout for the packed and batched layouts; only an entry point was missing. Add one for the 2D-3D and 3D-3D operand ranks, which group along M or not at all and so leave each group an entire [N, K] for the swizzle to apply to. The other two ranks group along N and K, the axes the swizzle interleaves across the whole matrix, so a group boundary falls inside it. They raise rather than fall back to the plain path, which would quietly ignore weights the caller had already shuffled. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Block scaling counts whole scale blocks, flooring k and n by the block size. That count used to be exact because k had to divide the K tile and n the N tile, and a tile covers a whole number of scale blocks. Masking the tail tile relaxed those bounds, and with them the property block scaling was relying on: a K or N that covers a partial block now leaves it out of the count and misindexes the scales from there on, dropping the tail of the contraction. Require the block bound directly under block scaling. It admits exactly what the tile bound used to admit, so nothing that worked stops working, and the rowwise path keeps the relaxation, carrying a scale per row and per column rather than per block. The resulting error is inside the tolerance this suite compares outputs with, so the test asserts the input is refused rather than checking numerics. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
f8f8bf16_rowwise_grouped_preshuffle reads as the sibling of f8f8bf16_rowwise_grouped, which is a different op taking tensor lists. It is the sibling of _stacked and carries that op's schema, so name it _stacked_preshuffle, matching the _dynamic and _mm siblings. The Python entry point had the same gap and gets the same treatment. Register the ops FlyDSL implements in one place per dispatch key. Excluding them from the CUDA key took a second conditional on USE_ROCM in the CUDA block, while the CPU block expressed the same thing through the arm it already had; put both in that arm and say once, against the list, which module binds each. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
It pointed at a masked kernel that has no counterpart here and described the module in terms of block scaling alone, though it serves both scaling schemes. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The pair was named for one scaling scheme and one group layout. The kernel serves both schemes and six layouts, including the masked one the "contiguous" half was there to distinguish it from, and "contiguous" also names the preshuffle pipeline variant inside the same module. Only the module and symbol names change. The kernel name that keys the compile cache is built from its own strings and is untouched. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Where the groups divide K each one spans the whole output, so the resolver's answers were all overwritten immediately after the call. It was reached with a layout it does not classify, which had it read the offsets as though they were int64 row counts and bound the descriptor to twice the operand's length. The backend dropped the loads, so the code it produced was already correct; take the call out so that stays true of the source. Fold the two blocks that qualify the K-grouped layout into one, and document the tile-padding flags and the occupancy hint alongside the rest. Name the CShuffle lane count the epilogue store width is derived from, and drop "contiguous" from the cache key, which named a distinction the kernel no longer draws. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
A contraction of length zero sums to zero, but the shapes that skip the launch returned the buffer as allocated. torch.empty leaves it holding whatever the caching allocator last had there, which is real data in any workload that recycles allocations, so the op returned that instead of a product. Zeroing covers the shapes that hold no elements too, where it does nothing. Say in the docstring that an allocated output has the one shape, so the layouts whose output is shaped differently keep supplying their own. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The kernel addresses the output as a flat row-major buffer and writes it in place, so a strided one is written to the wrong addresses. Its three sibling ranks already say so; this one did not, and a caller passing a slice of a wider tensor got a quietly wrong result rather than an error. Drop an attribute of the tuning decorator that nothing reads, and name the preshuffle siblings the gemm package registers now that all three have one. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Where the groups divide N, a store that began inside a group could finish outside it: the predicate compared the first of the columns it writes against the group's end, so a group whose column count is not a multiple of the store width overran into the next group's columns, which that group also writes. The result depended on which block ran last. Predicate on the last column instead. The straddling store is then dropped and those columns are left unwritten, which is what the op already documented and is at least the same every run. For a group whose columns do divide the store width the two tests agree, so nothing that was supported changes; the epilogue grows by sixteen instructions, which measures as no change on the shapes it applies to. Make the group metadata contiguous like the other operands, and describe the store bound as what it is: four columns, held at eight to match CK. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
aryaman-gupta
changed the title
aryaman-gupta
marked this pull request as ready for review
August 7, 2026 13:17
Only the stacked op and its preshuffle sibling had a fake, so tracing any of the other four raised UnsupportedOperatorException and torch.compile could not see through them. The padded layout returns a slab per group; the mm ops write the caller's buffer and return it, so the fake is that argument. State two contracts the code now keeps but did not say: a shape that reduces over nothing comes back zeroed, whichever layout asked for it, and the output where the groups divide K is the flattened per-group one. Assert that the row resolver is not reached with the layout whose metadata it would misread. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The comment said the key describes the work a block does, but M went in as the concatenation for the layouts that give each group a slab of it, so a shape with sixty-four groups of a hundred and twenty-eight rows keyed as eight thousand for blocks that only ever see a hundred and twenty-eight. Normalise it the way the N-grouped layout already did. The key selects which shapes share a tuned result, not what is measured, so a shape tuned on its own picks the same config either way. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
@jwfromm has imported this pull request. If you are a Meta employee, you can view this in D115222099. |
Contributor
|
Impressive speedups! Ill take a look at this today so we can get it merged :) |
Contributor
|
It would be great if you could also add wrappers for these new variants to |
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.
Summary
Implements the FP8 rowwise-scaled grouped GEMM entry points with FlyDSL and binds them on
ROCm, extending the grouped kernel added for the groupwise op in #470. On ROCm these ops are now
served by FlyDSL rather than CK; on CUDA nothing changes.
f8f8bf16_rowwise_grouped_stacked[G]int64 row countsf8f8bf16_rowwise_grouped_dynamic[G, M, K]slabs, runtime valid-row countf8f8bf16_rowwise_grouped_mm..._stacked_preshuffle,..._dynamic_preshuffle,..._mm_preshufflef8f8bf16_rowwise_groupedand_catare left on CK: both are markedUNSUPPORTED AND DEPRECATEDin CK's own source, have no in-repo callers and no tests.Design notes
One kernel, not six. The group geometry is a compile-time
layoutenum —sizes,offsets,padded,batched,n_offsets,k_offsets. Only the resolution step differs between them; theloaders, K loop and epilogue are shared.
n_offsets(groups divide N) andk_offsets(groupsdivide K) resolve the group from the N-block index and from the grid z axis respectively, rather
than from a row-count scan.
Scaling scheme is a compile-time flag. Rowwise carries one scale per row of A and per column
of B, both constant along K, so they factor out of the reduction and apply once in the epilogue
instead of per scale block inside the K loop. This also frees the tile from scale-block alignment,
making
tile_n = 64available, which block scaling cannot express.The K loop is rolled by default. Unrolling folds the tile index into every address, but traced
IR, compile time and code size then grow linearly with K: at K=8192 the plain loop is 5271 lines of
ISA and 12.9 s to compile unrolled, against 990 lines and 1.5 s rolled, with bitwise-identical
output. Both loop forms share one body.
k_offsetsis always rolled, since a group's K length isa runtime value.
Preshuffled B is offered only where a group owns a whole
[N, K]. The MFMA swizzle interleavesN and K across the whole matrix, so a group boundary inside either axis falls inside the swizzle.
_mm_preshuffletherefore serves the 2D-3D and 3D-3D ranks and rejects 3D-2D and 2D-2D rather thansilently ignoring weights the caller has already shuffled.
Tile selection is opt-in. Tuning runs only under
MSLK_AUTOTUNE_ENABLE, matching the CK andCUTLASS paths; otherwise a fixed default tile is used with no benchmarking.
mslk/flydsl/autotune.pyis generic and carries nothing specific to this kernel.
Performance
FlyDSL vs CK, CUDA-graph replay, G=8, gfx950, idle GPU. Figures are the reduction in kernel time
against CK,
so a positive number is the share of CK's time saved and a negative one is the share added.
59 of 72 shapes are faster.
_mm2D-3D_mm3D-3D_mm3D-2D_mm2D-2D_dynamicTwo caveats belong with these numbers:
MSLK_AUTOTUNE_ENABLEis read by both, andCK's tuning path raises
Kernels to tune over is emptyfor these entry points, so the benchmarktoggles it per side. Each is at its best available setting, not the same setting.
1280x8192, wherethere are too few N-blocks to fill the machine; the same weakness is reported in Flydsl groupwise grouped (#470) #470 against
Triton, so it reproduces across two kernels and two baselines. Most of the remainder are
total_M=1 at large N, where launch overhead dominates.
Both sides are timed after a 0.5 s clock ramp. Without it, a short window on an idle GPU measures
the clock ramp as much as the kernel — about 12% on these shapes — and whichever side is timed
second inherits the other's warm-up.
Per-shape detail
_mm2D-3D — M is total rows across all groups_dynamic— M is the per-group slab height, 75% of rows valid_mm3D-3D, 3D-2D and 2D-2D follow the same pattern; medians are in the summary table.Correctness
test/gemm/gemm_test.py -k "grouped or batched_gemm": 92 passed, 67 skipped, 0 failed ongfx950. The existing upstream contract tests for
_stacked,_dynamicand all four_mmrankcombinations now exercise FlyDSL rather than CK, so they are the primary coverage.
Added tests cover the boundaries the existing cases did not reach:
(3D-3D)
mm, and rejection of the two ranks that cannot take preshuffled weightsThe group-K boundary test gives alternate groups constant, single-signed data. A read that runs
past a group's K end then accumulates coherently rather than cancelling; with random data on both
sides the resulting error is the same order as the FP8 quantisation noise and sits inside this
file's comparison tolerance, so the test would pass whether or not the boundary is respected.
Across the benchmark,
max|fly − ck|is at most 0.0039 on every row.Unit tests pass on MI350 (gfx950) and MI300 (gfx942).
Notes for reviewers
gemm_ops.cppfrees the CK slots on ROCm for_stacked,_dynamicand_mm, so the Pythonbinding is the only implementation there rather than an override of one. A consequence is that a
ROCm build without FlyDSL has no implementation for these ops; this already held for
_stacked.CK continues to serve every one of them on CUDA.
once
gemm_ops.cppis recompiled.grouped_gemm_blockscale_{common,contiguous}.pybecomefp8_grouped_gemm{_common,}.py: the kernel serves both scaling schemes and six layouts, soneither "blockscale" nor "contiguous" described it. No behaviour change — the compile-cache key
is built from its own strings and is unaffected.
like CK these are documented rather than asserted: group K a multiple of 16 (2D-2D) and group N a
multiple of 8 (3D-2D). Measured against a dequantised reference, CK and FlyDSL are equally
inaccurate when group K is not a multiple of 16, so this is not a change in what works. For
ragged N, the epilogue drops the store that would cross a group boundary, leaving those columns
unwritten, rather than overrunning into the next group.
mslk/gemm/flydsl/**, which the path filter did not previously cover.