amd-agi/geak/flydsl-fp8-blockscale-gemm

When to use

Trigger on the problem signature, not a specific shape: any CK a8w8 fp8 block-scale dense GEMM (weight_block_size [128,128], arbitrary fp32 per-[128,128] block scales, bf16 out; the profiled hotspot name matches kernel_gemm_xdl*blockscale / gemm_a8w8_blockscale) being rewritten CK → FlyDSL on MI355X / gfx950 (CDNA4, OCP fp8_e4m3), where the FlyDSL native fp8 scaled-MFMA path fails parity against the CK oracle (the symptom that blocks the port) and CK config-JSON / tuned-CSV tiling is already near its ceiling (~1.0–1.06×). It is one kernel over arbitrary M/N/K — whatever shape hits this signature triggers the recipe; nothing is special-cased or enumerated.

Selection is by shape, never by layer/function name. The kernel only sees an (M, N, K) + quant layout; the model-side labels below (q_up_proj, down_proj, …) are just where each shape was measured and carry zero weight in matching or recipe choice — the same M/N/K in any other model gets the same recipe. The rows are illustrative instances seen so far, not a closed set; they span the regimes, which is why the tile re-sweep (not the mechanism) matters per shape:

(…and any future / decode -m/-nk variant of the same operator). Two stages, and only the first is universal: (1) the parity fix — do the block-scale in software (never E8M0) — is the same for every shape and is what unblocks the port; (2) the perf lever is shape-dependent — locality-limited very-wide-N shapes need a 4-wave core + XCD swizzle, grid-saturated K-light shapes need the 4-wave tile_n=256/tile_k=128 geometry but not XCD, and narrow-N/deep-K shapes win with an 8-wave "cluster" core. Three shapes are already tuned end-to-end; use their configs as the per-regime starting point (see Per-shape recipes below) rather than sweeping blind.

Mechanism

A. Universal — the parity fix (same for every shape; this is what unblocks the port)

The descriptive *why* now lives as a default-available fact card (read by author / tech-lead without use_expert_skills): operators/dense_gemm/numerics.md → "fp8 a8w8 block-scale: arbitrary fp32 ≠ E8M0", plus quantization/block_scaling_mxfp.md. In one line: CK a8w8 carries arbitrary fp32 per-[128,128] block scales, but the HW scaled-MFMA scale is E8M0 (power-of-two only) and cannot represent them — folding CK's scale into mfma_scale_*_f8f6f4 silently rounds it → parity fails. It is representational, not a knob, which is why CK→FlyDSL auto-ports on this family fail unless they pick the software path below.

The parity key (the one universal action): keep the MFMA unscaled low-precision fp8 and apply the fp32 block scales in software after the MFMA — promote the fp8 partials to fp32, ×per-block fp32 scale, accumulate. Bit-comparable to CK (err=0, cos=1.0 across seeds), and irreducible (the scale changes every 128 in K → one promote per K-block). FlyDSL already ships two cores that do exactly this software promote — pick one by shape (§ Per-shape recipes):

• a 4-wave blockscale-preshuffle fp8 GEMM core;
• an 8-wave ping-pong fp8 blockscale GEMM core.

Both pin the E8M0 HW scale to 1.0 and do the fp32 scale in software, so no target-backend change is needed — the fix is choosing a software-scale core over the native scaled-MFMA. Locate the equivalents by behaviour in whatever FlyDSL build is present; do not hardcode their paths or symbols.

B. Shape-dependent — the perf lever (do NOT assume q_up's answer transfers)

The correctness fix is universal; which perf lever wins inverts with the shape, because the two shape classes are bound by different things. Measured on the three tuned shapes:

• Very-wide-N / locality-limited (q_up N=65536/512 MiB C): **XCD-aware grid rasterization is the

lever. MI355X has 8 XCDs, each with its own L2 slice; remapping CTA `(bx,by)` so CTAs sharing a B/N-panel land on one XCD lifts L2 hit-rate `74.8%→83.0%`. The `xcd0→xcd8` ablation is worth ~1.23×**; xcd0 loses to CK while xcd8 wins. The 8-wave family spills on this accumulator footprint.

• K-light / grid-saturated (down_proj M=16384,N=7168,K=768): compute-bound with 7168 output tiles

over the device, so the decisive lever is the 4-wave `tile_n=256/tile_k=128` geometry that keeps the short six-K-tile loop fed. XCD remapping and fused_promote are measured non-load-bearing on this shape, and the 8-wave port is a measured dead-end (147 us ceiling vs the 4-wave winner near 129 us).

• Narrow-N / deep-K / tiny-C (qkv N=2048/16 MiB C/K=7168): grid+LDS-bound at **1 block/CU, 2

waves/SIMD, no cold tail → xcd_swizzle is a measured no-op and the 4-wave core is ~19% slower. The win comes from the 8-wave ping-pong kernel with the cluster schedule** — move all four fp32 promotes to after the last MFMA of the K-iter so none sits on a barrier critical path (MFMA util 45.8%→47.9%). 16-wave (higher occupancy) loses (register wall → spill), and BLOCK_M=256 is 14× slower (only 128 CTAs for 256 CUs).

Per-shape recipes (validated winners)

Three shapes are tuned end-to-end (err=0 / cos=1.0 all). Selection is by shape class, not by the layer label — match your (M, N, K) to the nearest row via the decision guide, apply its config, then re-sweep only the shape-dependent knobs. The seen as column is provenance only and never a selector.

Decision guide (pick the kernel family purely from the shape):

• Very-wide-N with measured locality pressure → 4-wave blockscale-preshuffle core, tile

t64×256×128 wpe2, then sweep XCD; xcd8 is the validated q_up winner.

• K-light with enough tiles to saturate the grid → the same 4-wave core at t64×256×128 wpe2;

lock the tile geometry first. On validated down_proj, XCD/fused-promote are non-load-bearing and the 8-wave family cannot beat the 4-wave winner.

• Narrow-N (N ≲ 2048) + deep-K + tiny C output → 8-wave ping-pong core with the `cluster` schedule,

BLOCK_M=128, BLOCK_N=256. (8-wave hides the promote; xcd_swizzle is a no-op; don't bother with it.)

Notes: the narrow-N/deep-K shape (4096×2048×7168) is the marginal-speedup case — it sits ~1.05–1.10× because it is already near its promote-free rowscale ceiling (~1.37×, a different numeric scheme, not a drop-in); treat ~1.05× as the family floor, not a regression. An untuned shape such as 16384×1536×7168 (narrow-N + deep-K, happens to be gate_up_proj) → start from the 8-wave `cluster` recipe by its shape class, not from the wide-N one — again a shape decision, not a name decision.

Bottleneck diagnosis (validated method — think before you sweep)

For a new/untuned shape, work out what limits it before touching knobs. These moves were validated across the three shapes and repeatedly corrected wrong reads — they are the transferable part, more than any single config:

1. Ablation is ground truth; util% is not. To test "is X the bottleneck?", remove X and re-measure —

never trust a utilization counter alone. On the wide-N shape rocprof showed VALU≈MFMA (~30%) and looked promote/VALU-bound, but deleting the entire per-block promote bought only ~1.7% → it was hidden under stalls all along. util% ≠ boundedness. Confirm every suspected bottleneck by ablation before optimizing it.

1. Build the promote-free "ceiling twin". Compare against an identical-MFMA-work kernel that scales

cheaply (a rowscale epilogue, or a plain fp8 GEMM). Same v_mfma count / same MFMA-busy-cycles but no per-block promote isolates the software-scale tax: on narrow-N the blockscale core ran 45.8% MFMA-busy vs the twin's 66.9% at bit-identical MFMA work → the gap is the promote-on-critical-path tax (not compute, not memory), and the twin's ~1.37× is the true ceiling (different numerics, not a drop-in). The plain-GEMM floor gives the shape's intrinsic limit.

1. Latency- vs throughput-bound: read both roofs at once. If compute-util and HBM-BW-util are both

low simultaneously (e.g. ~25% MFMA + ~22% HBM on the K-light shape), the kernel is latency-bound — idling, not out of compute or bandwidth (a throughput-bound kernel must saturate one roof). The FLOP-roofline "headroom" is then fictional; the real ceiling is the occupancy-latency one.

1. Quantify the reachable ceiling: `MFMA-duty ≈ occupancy × per-wave-MFMA-fraction`. Split wave lifetime

with the exact identity wave_cycles = active + wait_any + wait_inst (holds to ~0.01%). If occupancy is register-capped at 2 waves/SIMD and per-wave MFMA-exec ≈16%, the register-limited MFU ceiling ≈ ~31.6% — and you may already be at ~87% of it, so no in-kernel lever can help much. This tells you whether a lever can even exist before you spend time on it.

1. Localize with identical-work counters. Hold MFMA work fixed and vary one thing. MemUnitStalled=0

kills the memory hypothesis; AccVGPR=0 / no spill / LDS-bound occupancy kills the register hypothesis; equal MFMA-busy-cycles + higher total cycles pins the tax to scheduling/stalls, not work.

1. Phase-decompose via a K-sweep. Fit time ≈ fixed + slope·nKtiles. A dominant mainloop → target it

(8-wave / scheduling); a fixed cost already near its C-write BW floor → epilog-overlap is low-ROI.

1. Only levers on the *measured* critical path help. Narrow-N was promote-on-barrier-path → the win was

scheduling (cluster the promotes off the barrier); memory levers (scale-prefetch) were measured no-ops because MemUnitStalled=0. Never apply a memory lever to a latency/scheduling problem, or occupancy to a register-capped one.

1. Know when to stop. Once the winner is within a couple % of the ablation / ceiling-twin number (all

three shapes are), stop kernel-level tuning — the only lever left is E2E epilog fusion (cut the C round-trip), which is invisible in a standalone kernel and only bankable in the model graph.

Honesty caveats. Some byte counters (e.g. TCC/L2) can deadlock the queue on some ROCm builds → model the traffic instead and check robustness to ~2× error. GPU drifts ~5% run-to-run → quote same-session interleaved A/B only, and correctness gates every candidate (a fast-but-wrong kernel scores 0).