Port GEMM/MoE/conv kernels to the layout API

[coderfeli]

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