Back the FP8 rowwise grouped GEMM ops with FlyDSL on ROCm - #471

Open

aryaman-gupta wants to merge 39 commits into

meta-pytorch:mainfrom

aryaman-gupta:flydsl-rowwise-grouped

Commits on Aug 6, 2026

Add rowwise scaling to the grouped FP8 GEMM kernel

aryaman-gupta
and
claude
committed
3e685c7
Add opt-in tile autotuning for FlyDSL kernels

aryaman-gupta
and
claude
committed
ecbc5ff
Share the grouped GEMM host dispatch between scaling schemes

aryaman-gupta
and
claude
committed
6efcfbe
Back the FP8 rowwise grouped GEMM with FlyDSL on ROCm

aryaman-gupta
and
claude
committed
8c1381e
Cover the padded-layout rowwise grouped GEMM from gemm_test

aryaman-gupta
and
claude
committed
daf16ed
Support the padded per-group layout in the grouped FP8 GEMM

aryaman-gupta
and
claude
committed
85ee830
Mask the tail K tile so K need not divide the tile

aryaman-gupta
and
claude
committed
d66c45a
Mask the tail N block so N need not divide the tile

aryaman-gupta
and
claude
committed
a2cb94a
Keep pruning autotune candidates on both tile axes

aryaman-gupta
and
claude
committed
0151f08
Name the group-geometry encoding instead of flagging one of them

aryaman-gupta
and
claude
committed
1111ef8
Add a batched group layout for slabs that are entirely valid

aryaman-gupta
and
claude
committed
cae1b57
Accept cumulative offsets as the packed group geometry

aryaman-gupta
and
claude
committed
836ab15
Serve the 2D-3D and 3D-3D shapes of the rowwise grouped mm

aryaman-gupta
and
claude
committed
c9c6fc2
Drop the tile guard where the grid cannot produce a surplus tile

aryaman-gupta
and
claude
committed
28089ac
Extract group resolution into its own helper

aryaman-gupta
and
claude
committed
75f0da7
Let the groups divide N instead of M

aryaman-gupta
and
claude
committed
0f46390
Serve the 3D-2D shape of the rowwise grouped mm

aryaman-gupta
and
claude
committed
037028d
Allow the K loop to be rolled rather than unrolled

aryaman-gupta
and
claude
committed
aec313e
Roll the K loop by default and drop it from the tuning key

aryaman-gupta
and
claude
committed
7c2230a
Roll the ping-pong K loop too

aryaman-gupta
and
claude
committed
3eceb00
Let the groups divide the contraction

aryaman-gupta
and
claude
committed
58a187a
Tune the occupancy target alongside the tile

aryaman-gupta
and
claude
committed
f20a25c
Bound the K tail by the group's slice, not by the whole K

aryaman-gupta
and
claude
committed
b2d26a9
Cover group extents that do not land on a tile boundary

aryaman-gupta
and
claude
committed
f41db06
Serve the remaining rowwise grouped entry points from FlyDSL

aryaman-gupta
and
claude
committed
92d79f0
Run ROCm CI when the FlyDSL GEMM ops change
aryaman-gupta
committed
bf1ecc3
Make the group-K boundary case detectable

aryaman-gupta
and
claude
committed
377eaab
Hand the rowwise grouped ops to FlyDSL on ROCm

aryaman-gupta
and
claude
committed
c9ab9fd
Take preshuffled weights in the grouped mm where B stays whole

aryaman-gupta
and
claude
committed
9961710
Keep block scaling to whole scale blocks along K and N

aryaman-gupta
and
claude
committed
4012848
Name the preshuffle ops after the ops they mirror

aryaman-gupta
and
claude
committed
e09e2da
Correct the shared module's docstring

aryaman-gupta
and
claude
committed
df7552b
Name the kernel modules for what they compile

aryaman-gupta
and
claude
committed
ee4e4cc

Commits on Aug 7, 2026

Resolve group rows only where there are rows to resolve

aryaman-gupta
and
claude
committed
f37eb38
Zero the output where the grouped GEMM has nothing to launch

aryaman-gupta
and
claude
committed
4f77b45
Require a contiguous output where the groups divide M

aryaman-gupta
and
claude
committed
b323eca
Bound the epilogue store by its last column, not its first

aryaman-gupta
and
claude
committed
58c6e5b
Describe the registered ops to meta tracing

aryaman-gupta
and
claude
committed
36f5325
Key the tuning on one slab where the groups own one

aryaman-gupta
and
claude
committed
6177508