Port GEMM/MoE/conv kernels to the layout API by coderfeli · Pull Request #913 · ROCm/FlyDSL

coderfeli · 2026-07-27T14:49:30Z

Summary

Ports four GEMM / MoE / conv kernel families from manual buffer/byte-offset addressing to the FlyDSL layout API (the examples/04-preshuffle_gemm.py idiom: make_buffer_tensor + flat_divide + partition_S/D, make_composed_layout(SwizzleType) for LDS swizzle, make_tiled_mma/fx.gemm, crd2idx for swizzled offsets). Each port is cold-correct (FLYDSL_RUNTIME_ENABLE_CACHE=0) and perf-neutral (median-of-7 on gfx950, within clock noise).

moe_gemm_2stage is intentionally not in this PR (separate WIP).

Ports

gemm/preshuffle_gemm.py — buffer_ops removed; epilogue scale addressing on the layout API. Cold: fp8/int8/fp16/bf16/fp4/fp6/a8w8 + fused-epilogue modes + ragged-M. Perf parity (8192³ fp8 ~2237 TFLOPS, independently re-measured).
gemm/fp8_gemm_8wave.py + fp8_gemm_4wave.py — MMA via in-kernel make_tiled_mma; g2s + s2r LDS swizzle via prologue-hoisted crd2idx over make_composed_layout(SwizzleType(3,4,4)) (hot-loop ISA identical). Direct-DMA BufferCopyLDS + split-16@64 pack kept (repo idiom). fp8_gemm_utils.py untouched.
moe/mxfp_moe/gemm2.py — A-LDS read on real layout algebra (make_composed_layout(SwizzleType) + crd2idx, bit-for-bit vs the manual XOR), all gemm2 variants incl. a8w4 e2e stage-2. Cold 6/6 + 228/0 e2e; perf parity (fp4 1280.7, a8w4 828.4 TFLOPS).
conv/conv3d_implicit.py + conv3d_implicit_fp8.py — im2col gather via make_buffer_tensor/logical_divide/BufferCopyLDS128b (flat element index, OOB→sentinel); BIG_IN (>2GB) via make_buffer_ptr(rebased, num_records_bytes); fp8 MMA via in-kernel make_mma_atom(MFMA_Scale). Cold: bf16 26/26.

Documented "manual by necessity" (kept, with ISA evidence)

A few hot-loop sites are left manual on purpose — the layout-API form is op-count-identical but perturbs the tuned instruction schedule:

mxfp_moe gemm1 read: prologue-hoisted crd2idx is op-identical yet ~5.5% (fp4) / ~11% (fp8) slower (instruction-scheduling sensitivity). Documented in-code.
fp8_4wave interleaved AGPR s2r (~1.3%); direct-DMA op stays inline-asm (no register fragment; matches preshuffle_gemm.py::dma_a_to_lds); StoreC (final store, needs full wave-decomp, ~0 payoff).

Caveats

conv BIG_IN (>2GB) is compile-validated only — no >2GB test exists; the rebase arithmetic is byte-identical to the prior manual path and the mechanism is runtime-proven by the non-BIG path (26/26).
Pre-existing (not this PR): test_conv3d_implicit_fp8.py [1-96-4-8-9-96-1-1] cold flakiness (C=96 partial-K split-K, rel_err ~0.3075) reproduces on main.

Test plan

FLYDSL_RUNTIME_ENABLE_CACHE=0 python3 -m pytest \
  tests/kernels/test_preshuffle_gemm.py tests/kernels/test_fp8_gemm_rowscale.py \
  tests/kernels/test_moe_gemm.py tests/kernels/test_conv3d_implicit.py \
  tests/kernels/test_conv3d_implicit_fp8.py -v

🤖 Generated with Claude Code

Remove buffer_ops; port the epilogue scale_a/scale_b (load_epi_operands) addressing to the layout API. Perf-parity confirmed (median-of-7, within gfx950 noise); cold-correct across fp8/int8/fp16/bf16/fp4/fp6/a8w8 + fused-epilogue modes + ragged-M. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

MMA phase via in-kernel make_tiled_mma; g2s + s2r LDS swizzle expressed as crd2idx over make_composed_layout(SwizzleType) (prologue-hoisted, ISA-identical hot loop). Direct-DMA BufferCopyLDS and split-16@64 pack kept (repo idiom). Scheduler-sensitive paths (interleaved AGPR s2r, StoreC) left manual by design, ISA-verified. fp8_gemm_utils.py untouched. Perf parity, cold both preshuffle modes. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

gemm2 A-LDS read uses make_composed_layout(SwizzleType)+crd2idx (bit-for-bit vs the manual XOR mask), covering all gemm2 variants incl a8w4 e2e stage-2. gemm1 read left manual with ISA evidence: prologue-hoisted crd2idx is op-count-identical yet ~5.5-11% slower from instruction-scheduling sensitivity. Cold 6/6 variants + 228/0 e2e; perf parity (fp4 1280.7, a8w4 828.4). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

im2col gather -> make_buffer_tensor + logical_divide + BufferCopyLDS128b with flat element index and OOB->sentinel padding; BIG_IN (>2GB) paths via make_buffer_ptr(rebased, num_records_bytes). fp8: Mfma16x16x128 -> in-kernel make_mma_atom(MFMA_Scale)+fx.gemm, split-16@64 pack kept. Cold: bf16 26/26; fp8 6/7 (1 case is a pre-existing main bug, not this port). BIG_IN is compile-validated (no >2GB test). Perf no-regression (median-of-7). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Cut verbose per-port rationale/ISA-evidence comments (~88 -> ~29 added lines) to terse one-liners; code byte-identical (comments-only). black+ruff clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

…face Same mechanical cleanup as the map kernels, applied to kernel_bf16_cmp: fx.Index(...) -> fx.Int32(...) (buffer_load wants i32; every fx.Index was forcing an index->i32 cast) arith.constant / cmpi -> fx.Int32(n) / typed compares arith.constant_vector -> fx.Vector.filled(4, 0.0, fx.Float32) const_expr(<python int>) -> the int (make_layout/vec_load take plain ints) dropped `lane = tid` alias and the duplicate mfma_row/mfma_col computation The EPILOGUE is left on fx.Index / vector.extract deliberately, with the reason in-code. Migrating it is op-count-identical but measured 185 us vs 170 us at C_OUT=512 (median-of-7, range 184.8-186.4 so well outside noise) -- an instruction-scheduling effect around the predicated atomics, the same class of exception PR ROCm#913 documents for mxfp_moe gemm1 and fp8_4wave. Bisected by reverting prologue/hot-loop/epilogue independently: the hot loop is free, the epilogue is not. Verified: LUT, in_rows and out_rows BIT-IDENTICAL to the pre-migration golden on four shapes incl. K=5; float output within the kernel's own atomic-order noise. 39/39 tests pass cold. Perf median-of-5 vs before: C=128 1.00x, C=256 1.02x, C=512 1.02x. Co-Authored-By: Claude <noreply@anthropic.com>

…face Same mechanical cleanup as the map kernels, applied to kernel_bf16_cmp: fx.Index(...) -> fx.Int32(...) (buffer_load wants i32; every fx.Index was forcing an index->i32 cast) arith.constant / cmpi -> fx.Int32(n) / typed compares arith.constant_vector -> fx.Vector.filled(4, 0.0, fx.Float32) const_expr(<python int>) -> the int (make_layout/vec_load take plain ints) dropped `lane = tid` alias and the duplicate mfma_row/mfma_col computation The EPILOGUE is left on fx.Index / vector.extract deliberately, with the reason in-code. Migrating it is op-count-identical but measured 185 us vs 170 us at C_OUT=512 (median-of-7, range 184.8-186.4 so well outside noise) -- an instruction-scheduling effect around the predicated atomics, the same class of exception PR ROCm#913 documents for mxfp_moe gemm1 and fp8_4wave. Bisected by reverting prologue/hot-loop/epilogue independently: the hot loop is free, the epilogue is not. Verified: LUT, in_rows and out_rows BIT-IDENTICAL to the pre-migration golden on four shapes incl. K=5; float output within the kernel's own atomic-order noise. 39/39 tests pass cold. Perf median-of-5 vs before: C=128 1.00x, C=256 1.02x, C=512 1.02x.

Picks up ROCm#913 ("Port GEMM/MoE/conv kernels to the layout API"), which landed after this branch was cut. conv3d_implicit.py merged automatically: ROCm#913's BufferCopyLDS128b copy atom + real-extent OOB sentinels (X_ELEMS / W_ELEMS) compose with this branch's channel padding, since the padded C is what reaches the kernel. conv3d_implicit_fp8.py conflicted as a whole-file rewrite (this branch replaced the kernel). Resolved by keeping the 8-wave kernel and porting ROCm#913's two layout-API changes onto it: - MMA: Mfma16x16x128 -> make_tiled_mma + fx.gemm via fp8_gemm_8wave's TiledMmaDriver, the same construction ROCm#913 gave the GEMM pipeline this kernel is a port of. - BIG_IN descriptor: drop the hand-rolled make_ptr/IntToPtrOp rebase (which reached into buffer_ops._get_buffer_flags) for the public make_buffer_ptr helper, keeping the explicit 2GB num_records. Verified on MI355X (gfx950), FLYDSL_RUNTIME_ENABLE_CACHE=0: - test_conv3d_implicit_fp8.py 35 passed - test_conv3d_implicit.py 57 passed - test_fp8_gemm_rowscale.py 8 passed - final ISA byte-identical before/after the fp8 port (perf-neutral) - check_python_style.sh clean

Rebasing onto main moved this kernel from flydsl 0.2.4 to 0.3.x, and it started returning numbers uncorrelated with the reference -- logits diff 0.98, i.e. not drift, wrong data -- while compiling and running clean. The cause is the inline-asm MMA's accumulator constraint. asm mode "2" emits v_mfma_f32_16x16x128_f8f6f4 with "=a,v,v,0": an AGPR-class output tied to srcC. On 0.2.4 that was correct; on 0.3.x it silently miscompiles, and the shape of the damage is distinctive -- with B = I, so the output must equal A, every fourth output row comes back entirely zero and the rest are ~25% populated. Four accumulator slots per lane map to four consecutive rows, so that is one whole accumulator slot lost. Isolated to the constraint: the operand types entering the asm are right (vector<8xi32>, vector<8xi32>, vector<4xf32> against a vector<4xf32> result), dropping the amdgpu-agpr-alloc="128,128" passthrough does not help, and the identical asm with "=v" (mode "3") is correct. So acc_mode defaults to "vgpr" -- still the in-place asm accumulate this kernel is built around, just in the VGPR file -- and "agpr" now asserts instead of returning wrong numbers, since nothing about the failure is visible at compile time. _raw -> as_ir_value on the asm operands is not the fix; it is the unwrap main moved to in #913 when it ported the parent fp8_gemm_4wave, and it is what the passing configuration runs. Measured on gfx950 with flydsl 0.3.2: 5 passed, and 2751 TFLOPS on 2x2048/n4096/k7168 against 2708 for the intrinsic MMA (agpr_inplace=False) -- the same within single-run noise, so the asm path keeps its reason to exist. Worth recording for whoever hits this next: the parent kernel does not use inline-asm MFMA at all, so this path had no coverage in main's 0.2.4 -> 0.3.x ports. Everything else about the kernel survived the jump.

Felix Li and others added 4 commits July 27, 2026 14:46

coderfeli marked this pull request as draft July 27, 2026 14:58

Trim porting comments to match repo density

2bc1b83

Cut verbose per-port rationale/ISA-evidence comments (~88 -> ~29 added lines) to terse one-liners; code byte-identical (comments-only). black+ruff clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

coderfeli marked this pull request as ready for review July 29, 2026 04:38

Merge branch 'main' into port/layout-api-ports

7c22411

coderfeli merged commit e07991f into main Jul 30, 2026
15 checks passed

coderfeli deleted the port/layout-api-ports branch July 30, 2026 07:58

coderfeli mentioned this pull request Jul 30, 2026

moe: mxfp_moe gemm1 fp4 A-LDS read -> crd2idx (fp8 kept manual) #929

Closed

jiacao-amd mentioned this pull request Jul 31, 2026

[Kernel] Add submanifold sparse 3D convolution (bf16 implicit GEMM, gfx950) #942

Closed

coderfeli mentioned this pull request Aug 6, 2026

[cleanup] Migrate kernels to fx.copy / .llvm_ptr; expand kernel-code-cleanup skill #974

Merged

Navigation Menu

Repository navigation

Port GEMM/MoE/conv kernels to the layout API - #913

Port GEMM/MoE/conv kernels to the layout API#913
coderfeli merged 6 commits into
mainfrom
port/layout-api-ports

coderfeli commented Jul 27, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

coderfeli commented Jul 27, 2026

Summary

Ports

Documented "manual by necessity" (kept, with ISA evidence)

Caveats

Test plan

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant