feat(gemm): optimize blockwise FP8 GEMM Triton kernels (gfx950 / MI355X)

[ChengYao-amd]

Optimize blockwise FP8 GEMM Triton kernels (gfx950 / MI355X)

PR branch: dev/yaoc/optimize_fp8_blockwise_gemm_triton
Base branch: main @ 06b8d3f (Add HYBRID FP8 format support for Triton backend in gemm and grouped_gemm)
Commit on this branch: opt: blockwise FP8 GEMM Triton kernels (gfx950 / MI355X) (4 files, +597 / −46)

1. TL;DR

This PR optimizes the blockwise FP8 GEMM Triton path on gfx950 (MI355X).
Source-level changes touch four files only:

• primus_turbo/triton/gemm/gemm_fp8_kernel.py
• primus_turbo/triton/quantization/quant_blockwise.py
• primus_turbo/pytorch/kernels/quantization/quantization_impl.py
• primus_turbo/pytorch/ops/gemm_fp8.py

Aggregate impact on the 84-shape bench_gemm_turbo.py suite
(PRIMUS_TURBO_GEMM_BACKEND=TRITON, --dtype fp8 --granularity blockwise,
geomean over the 83 shapes both branches pass correctness on):

Metric	main (gmean)	this PR (gmean)	delta
Forward TFLOPS	1470.6	1474.2	+0.25%
Backward TFLOPS	1018.7	1200.4	+17.83%
Combined-step TFLOPS¹	1134.3	1278.4	+12.71%

¹ Combined Step TFLOPS = 6·M·N·K / (fwd_ms + bwd_ms) / 1e9.

Additional wins:

• Correctness coverage: main FAILs TestID 47 (Llama-3.1-405B
  M=32768 N=106496 K=16384, dA=NaN). This PR PASSes all 84 shapes.
• Per-shape distribution is uniformly positive on combined-step:
  all 83 PASS-on-both shapes are ≥ +1% on combined-step
  (median +13.13%, p25 +10.95%, p75 +15.30%), 82 / 83 shapes are
  ≥ +5%. The backward delta is strictly positive on every shape
  (median +18.48%, min +3.22%, max +25.76%). Forward is centred near
  zero (median +0.63%, mean +0.28%) but 5 shapes carry a fwd
  regression ≤ −5% — see §4.1 for why these still come out net
  positive on combined-step.

CK / hipBLASLt / grouped-GEMM paths are not touched. gfx942 (MI300X /
MI325X) falls back to the unchanged shared kernel.

2. Scope

The PR contains exactly four files of source change:

File	LOC delta	Role
primus_turbo/triton/gemm/gemm_fp8_kernel.py	+450 / −27	Layout-specialised NT/NN/TN kernels, autotune extensions, EVEN_K fast path, wgrad transposed-store epilogue
primus_turbo/triton/quantization/quant_blockwise.py	+50 / −0	New single-launch row+col fused quant (quant_fp8_blockwise_dual)
primus_turbo/pytorch/kernels/quantization/quantization_impl.py	+53 / −2	Dual-quant impl + custom-op registration
primus_turbo/pytorch/ops/gemm_fp8.py	+44 / −17	Wires dual-quant fusion through Float8 forward, reuses cached col-major quant in backward

Out of scope (not in this PR):

• agent/, agent/historical_experience/ — agent-tooling assets are
  shipped in a separate PR.
• gemm_kernel.py — only the FP8 blockwise kernel file is touched; the
  shared bf16/fp16 GEMM kernel is unchanged.
• gemm_fp8_kernel.py shared kernel _blockwise_fp8_autotune_kernel is
  preserved (still used by grouped-GEMM); only NT/NN/TN dispatch is
  switched to the layout-specialised variants.

3. Optimization techniques (kept in this PR)

3.1 Layout-specialised kernels with EVEN_K fast path

Replace the single shared _blockwise_fp8_autotune_kernel (still kept
for grouped-GEMM) with three dispatch-specialised variants:

• _blockwise_fp8_nt_kernel — forward (A @ B.T)
• _blockwise_fp8_nn_kernel — dgrad (grad_out @ B)
• _blockwise_fp8_tn_kernel — wgrad (A.T @ grad_out), new in §3.4

Each variant hard-codes its layout constexpr (A_K_CONTIGUOUS,
B_K_CONTIGUOUS, SCALE_2D_B) at the kernel level, eliminating the
runtime per-call branch decisions and shrinking each binary.

EVEN_K fast path on NT/NN: when K % BLOCK_K == 0 (the common
training shape), the K-tail mask tl.load(..., mask=...) is replaced
with an unmasked path, removing two v_cmp_* + one v_cndmask_* per
K iteration.

3.2 Autotune space extension

The main autotune space is 32 configs
(BM ∈ {128,256}, BN=128, BK=128, kpack ∈ {1,2},
GM ∈ {4,8}, CHUNK ∈ {32,64}, num_stages ∈ {1,2},
num_warps = 4 if BM=128 else 8).

This PR extends it to 96 configs on gfx950 (forward / dgrad):
BLOCK_N ∈ {64,128} and num_stages ∈ {1,2,3} are added.

The wgrad path is restricted to num_stages ∈ {1,2} to work around a
Triton 3.7 AMD-backend assertion that fires only on the dual
strided-K load pattern combined with num_stages=3. This is the
single-line fix carried in the otherwise larger §3.1 change set;
main does not need it because main's autotune space does not
contain num_stages=3 in the first place.

3.3 NN-only matrix_instr_nonkdim=16 candidates on BLOCK_M=256

NN (dgrad) at BM=256 BN=128 nw=8 ns=3 overflows the 16 AGPR / warp
budget that gfx950's default v_mfma_f32_32x32x64_f8f6f4 requires for
its accumulator (405 v_accvgpr_* ops, ~24% of the NN kernel ASM in
the dev-branch baseline). Stacking matrix_instr_nonkdim=16 candidates
only on the BM=256 path lets Triton compile a
v_mfma_f32_16x16x128_f8f6f4 variant that uses 4 AGPR / warp; the
autotuner picks it on the shapes where the AGPR pressure dominates.

The BM=128 nw=4 path is left alone because it already only uses
16 AGPR / warp, and stacking a 16x16 MFMA candidate there would just
inflate cold-autotune time without any binary winning.

Empirical hit rate: nonkdim=16 is selected by 49 / 69 NN cache keys
(71%). Top wins land on Llama-2-70B M=16384 N=8192 K=28672 (bwd
+5.86%) and Qwen2.5-72B M=8192 N=8192 K=29568 (bwd +5.96%).

Cost: NN autotune candidate count 96 → 144 (+50% cold-bench time on a
fresh cache).

3.4 wgrad-specialised kernel with transposed-store epilogue

In the dev-branch baseline, wgrad (_blockwise_tn, trans_c=True)
shared the kernel with NT/NN and produced its (N, M) output via the
shared-kernel store path. With stride_cm=1, stride_cn=N, the BN
dimension is strided in physical memory while BM is contiguous, which
the compiler can only serialise into 64 × buffer_store_short (32
low-half + 32 high-half) per tile — about 6-16× lower effective store
throughput than the NT/NN path's 8 × buffer_store_dwordx4.

The new _blockwise_fp8_tn_kernel rewrites the epilogue:

c_ptrs_t = C_ptr + offs_n[:, None] * stride_cn + offs_m[None, :] * stride_cm
mask_t = (offs_n[:, None] < N) & (offs_m[None, :] < M)
acc_t = tl.trans(acc.to(C_ptr.type.element_ty))
tl.store(c_ptrs_t, acc_t, mask_t)

The tl.trans swaps BN to the outer dim, swapping pointer indexing
puts the contiguous-stride dim on the inner axis, and the store
naturally coalesces into dwordx4.

ASM-level verification on a fresh Triton cache:

$ find . -name '_blockwise_fp8_tn_kernel.amdgcn' \
    | xargs grep -l buffer_store_short | wc -l
0
$ find . -name '_blockwise_fp8_tn_kernel.amdgcn' \
    | xargs grep -l 'store_dwordx4' | wc -l
96

96 / 96 wgrad binaries now use (buffer|global)_store_dwordx4 × 4-8,
fully eliminating the buffer_store_short × 64 pattern.

The shared kernel is still used by grouped-GEMM (which uses
trans_c=False and already stores into a contiguous (M, N)
layout, so it does not need the transposed-store path).

3.5 Single-launch row + col fused quant (quant_fp8_blockwise_dual)

The forward path needs both row-major (for A @ B.T) and col-major
(saved for A.T @ grad_out in backward) FP8 blockwise quantisations
of the input activation. The dev-branch baseline launched two separate
quant kernels; this PR fuses them into a single quant_fp8_blockwise_dual
launch:

• One tl.load of the bf16 activation tile.
• Two scale computations (row + col) reusing the same loaded values.
• Two stores (row-major FP8 + col-major FP8 + their two scale tensors).

The col-major FP8 output is save_for_backward-ed and consumed
directly by wgrad without re-quantising.

This fusion is implemented at kernel level rather than as a
wrapper-level cache (no id()-keyed reuse of quant outputs across
calls), so the gain holds in real LLM training where activations
are fresh each iteration; it does not depend on the benchmark
harness's 100-iteration tensor-reuse pattern.

3.6 Llama-3.1-405B correctness fix

main FAILs TestID 47 (M=32768 N=106496 K=16384) with dA=NaN in
the dgrad NN path. This PR includes the targeted index-arithmetic fix
in _blockwise_fp8_nn_kernel (5 lines). The shape now passes with
dA=28.6 dB.

4. Performance results

4.1 Aggregate (PASS-on-both, 83 cases)

Metric	main	this PR	delta
Forward TFLOPS (gmean)	1470.6	1474.2	+0.25%
Backward TFLOPS (gmean)	1018.7	1200.4	+17.83%
Combined-step TFLOPS (gmean)	1134.3	1278.4	+12.71%

Per-shape delta distribution over the same 83 shapes:

Metric	median	mean	p25	p75	shapes ≥ +1%	shapes ≤ −1%	shapes ≥ +5%	shapes ≤ −5%
Fwd Δ%	+0.63	+0.28	−0.29	+1.80	34	14	0	5
Bwd Δ%	+18.48	+17.92	+15.30	+21.25	83	0	82	0
Comb Δ%	+13.13	+12.75	+10.95	+15.30	83	0	82	0

The combined-step delta is strictly positive on every shape that both
branches pass; the smallest combined gain is +3.29% (TestID 43,
Llama-3.1-405B M=16384 N=106496 K=16384) and the largest is +18.51%
(TestID 64, Qwen2.5-72B M=8192 N=8192 K=29568).

5 shapes carry a forward regression ≤ −5% — all of them have
(N, K) = (4096, 4096) or (N, K) = (4096, 11008):

TestID	Case	M	N	K	Δ fwd	Δ bwd	Δ comb
4	Llama-2-7B	4096	4096	11008	−8.79%	+14.40%	+7.14%
78	Mistral-7B	8192	4096	4096	−8.29%	+12.14%	+6.94%
26	Llama-3.1-8B	8192	4096	4096	−7.51%	+12.16%	+5.48%
74	Mistral-7B	4096	4096	4096	−7.48%	+25.76%	+13.51%
6	Llama-2-7B	8192	4096	4096	−6.85%	+12.63%	+6.94%

A further 5 shapes show a fwd regression in the −5% to −2.5% band
(TID 8, 14, 49, 2, 62). All 10 net-negative-fwd shapes remain net
positive on combined-step thanks to a +12-26% bwd uplift on the same
shapes.

4.2 Per-model (PASS-on-both)

main numbers come from tmp/main/gemm_fp8_blockwise_triton_benchmark.csv.
This PR numbers come from tmp/new-branch/gemm_fp8_blockwise_triton_benchmark.csv.

Model	n	main cmb gmean	this PR cmb gmean	delta
Llama-2-7B	12	1105.3	1249.7	+13.07%
Llama-2-70B	12	1169.3	1317.4	+12.66%
Llama-3.1-8B	12	1112.8	1266.1	+13.78%
Llama-3.1-405B	11	1212.4	1323.4	+9.15%²
Qwen2.5-7B	12	1070.4	1211.7	+13.20%
Qwen2.5-72B	12	1183.2	1336.6	+12.97%
Mistral-7B	12	1100.1	1253.0	+13.90%

² Llama-3.1-405B excludes TestID 47 from the comparison because it
FAILs on main (PASS on this PR — see §3.6 and §4.4). The smaller
+9.15% on this model is dominated by the two large-N wgrad shapes
(TID 43 +3.29%, TID 39 +7.72% combined) where the BN tile count
explodes to N / 128 = 832; the other 9 Llama-3.1-405B shapes are at
+7.7% to +12.7% combined and match the other models.

4.3 Best combined movers (vs main)

TestID	Case	M	N	K	main comb	this PR comb	Δ comb	Δ fwd	Δ bwd
64	Qwen2.5-72B	8192	8192	29568	1148.1	1360.6	+18.51%	+3.18%	+25.00%
34	Llama-3.1-8B	32768	4096	4096	1024.4	1208.3	+17.95%	+1.63%	+25.04%
57	Qwen2.5-7B	32768	4608	3584	1054.2	1239.3	+17.56%	−0.88%	+25.58%
50	Qwen2.5-7B	8192	3584	3584	928.5	1088.6	+17.24%	+1.51%	+23.24%
82	Mistral-7B	16384	4096	4096	1043.8	1221.7	+17.04%	+0.27%	+23.18%
66	Qwen2.5-72B	16384	8192	8192	1184.4	1385.9	+17.02%	+0.80%	+23.66%
24	Llama-2-70B	16384	8192	28672	1182.3	1381.8	+16.88%	+1.49%	+23.56%
84	Mistral-7B	16384	4096	14336	1152.2	1345.6	+16.78%	+3.12%	+23.14%

The largest gains span both small-K + mid-M (K=4096, the
quant-fusion + wgrad transposed-store §3.5/§3.4 dominate) and
large-K + mid-M (K ∈ {28672, 29568}, the NN nonkdim=16 §3.3
plus the wgrad transposed-store §3.4 stack). The forward delta on
these top-bwd shapes is in the −1% to +3% range, confirming the gains
live primarily on the backward path.

4.4 Smallest combined gains (vs main)

The "worst" rows below are all still positive — no shape that PASSes
on both branches regresses on combined-step:

TestID	Case	M	N	K	main comb	this PR comb	Δ comb	Δ fwd	Δ bwd
43	Llama-3.1-405B	16384	106496	16384	1291.9	1334.4	+3.29%	+3.48%	+3.22%
26	Llama-3.1-8B	8192	4096	4096	1071.0	1129.6	+5.48%	−7.51%	+12.16%
71	Qwen2.5-72B	32768	59136	8192	1220.9	1299.2	+6.41%	+0.52%	+8.65%
6	Llama-2-7B	8192	4096	4096	1071.0	1145.3	+6.94%	−6.85%	+12.63%
78	Mistral-7B	8192	4096	4096	1071.0	1145.3	+6.94%	−8.29%	+12.14%

Two distinct populations show up in this Bottom-5:

• TID 43 and TID 71 are large-N (N ∈ {106496, 59136}) wgrad-bound
  shapes. The transposed-store epilogue (§3.4) brings less here
  because the BN tile count blows up (N / 128 = 462-832) and the
  per-tile tl.trans(acc) exhausts register budget; see §5.1.
• TID 26, 6, 78 are (8192, 4096, 4096) square-ish shapes whose
  combined gain is held back by a fwd regression of −6.85% to −8.29%
  (cold-autotune drift, see §5.3). Their bwd is still +12-13%, so
  combined remains net positive.

TestID 47 (Llama-3.1-405B M=32768 N=106496 K=16384) FAILs on main
and does not appear in the apples-to-apples table. On this PR it
PASSes with combined = 1347.6 TFLOPS (fwd 1626.5 TFLOPS / 70.30 ms,
bwd 1241.1 TFLOPS / 184.27 ms).

5. Known limitations

5.1 Large-N wgrad shapes get smaller combined gain

Three shapes with N ≥ 57344 produce noticeably smaller combined
gain than the +12.71% gmean:

TestID	Case	M	N	K	Δ comb	Δ fwd	Δ bwd
43	Llama-3.1-405B	16384	106496	16384	+3.29%	+3.48%	+3.22%
71	Qwen2.5-72B	32768	59136	8192	+6.41%	+0.52%	+8.65%
39	Llama-3.1-405B	8192	106496	16384	+7.72%	+4.21%	+9.11%

Root cause: on these shapes the BN tile count explodes to
N / 128 = 462-832 and the per-tile tl.trans(acc) in the wgrad
transposed-store epilogue (§3.4) exhausts register budget, spilling
the transpose to LDS-exchange. The shapes are still strictly net
positive vs main because main's wgrad starts from a much lower
baseline (Bwd TFLOPS in the 1100-1190 range, vs the post-PR 1200+).

Mitigation, deferred to a follow-up PR: add a
WGRAD_TRANS_STORE: tl.constexpr switch and let the autotuner pick
between the transposed-store path and the legacy short-store path
per shape; expected to recover the regressed subspace without giving
up the +17.83% bwd gmean on the other 80 shapes.

5.2 Cold-autotune wall-time

This PR's autotune candidate count is larger than main's:

• NT (forward): 32 → 96 (3× from BLOCK_N=64 and num_stages=3)
• NN (dgrad): 32 → 144 (3× + nonkdim=16 stack on BM=256)
• TN (wgrad): 32 → 96 (3× same as NT, new specialised kernel)

End-to-end cold benchmark wall-time on the 84-shape suite goes from
~25 min (main) to ~120 min (this PR). For production deployment we
recommend caching the per-shape best config into a lookup table after
a one-time autotune sweep; this is not in this PR but is straight
to add as a follow-up.

5.3 Cold-autotune drift on near-tie configs

The 5 shapes with Δ fwd ≤ −5% listed in §4.1 are all in the
(*, 4096, 4096) / (4096, 4096, *) square-ish family. On these
shapes several autotune candidates are within ~1% of each other on
fwd time, and the cold-cache pick can flip between sweeps; this
shows up as a per-shape ±5-10% fwd drift between independent runs
even when the underlying kernel is unchanged. The drift does not
affect the PR-vs-main aggregate (every drifting shape is still net
positive on combined-step thanks to the +12-26% bwd uplift), but
pinning the best config per shape (§5.2) is the recommended fix in
production.

6. Reproducing the numbers

# Build (editable install, primus_turbo)
pip install -e .

# Run the 84-shape benchmark
CUDA_VISIBLE_DEVICES=0 PRIMUS_TURBO_GEMM_BACKEND=TRITON \
  python benchmark/ops/bench_gemm_turbo.py \
  --dtype fp8 --granularity blockwise \
  -o pr_fp8_blockwise_triton.csv

# Correctness on the FP8 blockwise subset
pytest tests/pytorch/ops/test_gemm_fp8.py -v -k 'blockwise and TRITON'

Environment used for the numbers in this report:

Item	Value
GPU	AMD Instinct MI355X (gfx950, single card)
PyTorch	2.10.0a0+git449b176
Triton	3.7.0+gitf4a3db9e
primus_turbo	editable install at dev/yaoc/optimize_fp8_blockwise_gemm_triton
Backend selection	PRIMUS_TURBO_GEMM_BACKEND=TRITON
Quantisation	blockwise, FP8 E4M3, block_size=128
Iterations	20 warmup + 100 timed (fwd / bwd separately)

[ChengYao-amd]

(gfx942 / MI300X) Performance Report

Comparison: tmp/main (baseline) vs tmp/opt (optimized)
Workload: gemm_fp8_blockwise_triton_benchmark
Platform: ROCm / AMD MI300X
Dtype: FP8, Blockwise granularity
Total test cases: 84 (7 model families × 12 shapes each)

---

1. Executive Summary

Metric	Mean	Median	Min	Max
Forward time reduction	+1.80%	+1.49%	-2.83%	+7.14%
Backward time reduction	+9.53%	+9.72%	+1.64%	+15.11%
Forward TFLOPS gain	+1.91%	+1.57%	-2.75%	+9.05%
Backward TFLOPS gain	+10.64%	+10.83%	+1.66%	+18.14%

Aggregate average	main	opt	Δ
Forward TFLOPS	627.96	639.32	+11.37
Backward TFLOPS	449.38	496.34	+46.96

Highlights

• Backward kernel achieves a substantial ~10% mean speedup with the maximum reaching +15.11%.
• Forward kernel shows a modest ~2% mean improvement.
• All 84 cases produce correct results; one previously failing case (TestID 47, Llama-3.1-405B) now passes.

---

2. Per-Model Summary (Mean Improvement)

Model	#Cases	Forward Speedup	Backward Speedup	Forward TFLOPS	Backward TFLOPS
Mistral-7B	12	+2.96%	+11.60%	+3.16%	+13.14%
Llama-3.1-8B	12	+1.28%	+11.06%	+1.16%	+12.43%
Llama-2-7B	12	+2.26%	+10.52%	+2.52%	+11.92%
Qwen2.5-7B	12	+1.39%	+10.49%	+1.57%	+11.81%
Llama-2-70B	12	+1.50%	+8.82%	+1.53%	+9.72%
Qwen2.5-72B	12	+0.76%	+7.79%	+0.82%	+8.49%
Llama-3.1-405B	12	+2.45%	+6.44%	+2.59%	+6.97%

The optimization benefits small- and mid-sized models more uniformly. Large-model shapes (405B / 72B) show a smaller but still consistently positive backward speedup.

---

3. Top 10 Backward Improvements

TestID	Model	Shape (M, N, K)	Bwd Main (ms)	Bwd Opt (ms)	Speedup	TFLOPS Gain
49	Qwen2.5-7B	(8192, 4608, 3584)	1.39	1.18	+15.11%	+18.14%
77	Mistral-7B	(8192, 6144, 4096)	2.04	1.75	+14.22%	+16.29%
25	Llama-3.1-8B	(8192, 6144, 4096)	2.05	1.76	+14.15%	+16.12%
81	Mistral-7B	(16384, 6144, 4096)	4.00	3.44	+14.00%	+16.50%
29	Llama-3.1-8B	(16384, 6144, 4096)	3.98	3.43	+13.82%	+15.83%
53	Qwen2.5-7B	(16384, 4608, 3584)	2.69	2.33	+13.38%	+15.52%
10	Llama-2-7B	(16384, 4096, 4096)	2.70	2.34	+13.33%	+15.32%
78	Mistral-7B	(8192, 4096, 4096)	1.37	1.19	+13.14%	+15.25%
26	Llama-3.1-8B	(8192, 4096, 4096)	1.40	1.22	+12.86%	+14.89%
54	Qwen2.5-7B	(16384, 3584, 3584)	2.11	1.85	+12.32%	+13.98%

These top backward gains share the pattern of moderate N/K dimensions (N ≤ 6144, K ≤ 4096) combined with relatively large M. The backward pass computes dB = X^T · dY where the contraction dimension equals M, so kernels in this regime benefit most from the new tiling/scheduling.

---

4. Top 10 Forward Improvements

TestID	Model	Shape (M, N, K)	Fwd Main (ms)	Fwd Opt (ms)	Speedup	TFLOPS Gain
2	Llama-2-7B	(4096, 4096, 4096)	0.28	0.26	+7.14%	+7.55%
74	Mistral-7B	(4096, 4096, 4096)	0.28	0.26	+7.14%	+6.64%
73	Mistral-7B	(4096, 6144, 4096)	0.42	0.39	+7.14%	+9.05%
49	Qwen2.5-7B	(8192, 4608, 3584)	0.50	0.47	+6.00%	+7.26%
84	Mistral-7B	(16384, 4096, 14336)	3.08	2.90	+5.84%	+6.04%
1	Llama-2-7B	(4096, 12288, 4096)	0.77	0.73	+5.19%	+5.63%
47	Llama-3.1-405B	(32768, 106496, 16384)	180.93	171.81	+5.04%	+5.31%
45	Llama-3.1-405B	(32768, 18432, 16384)	31.56	30.07	+4.72%	+4.98%
39	Llama-3.1-405B	(8192, 106496, 16384)	45.19	43.21	+4.38%	+4.59%
6	Llama-2-7B	(8192, 4096, 4096)	0.49	0.47	+4.08%	+5.17%

The largest forward gains land on (a) small square shapes (≤4096) and (b) very wide N shapes (N ≥ 100k) of the 405B model, suggesting the optimized kernel handles both ends of the shape spectrum better than the baseline.

---

5. Forward Regressions

Ten shapes show a small forward slowdown. All ten still report a healthy positive backward speedup.

TestID	Model	Shape (M, N, K)	Fwd Speedup	Bwd Speedup
48	Llama-3.1-405B	(32768, 16384, 53248)	-2.83%	+1.64%
71	Qwen2.5-72B	(32768, 59136, 8192)	-2.45%	+5.95%
44	Llama-3.1-405B	(16384, 16384, 53248)	-1.38%	+8.81%
59	Qwen2.5-7B	(32768, 37888, 3584)	-0.88%	+7.39%
58	Qwen2.5-7B	(32768, 3584, 3584)	-0.71%	+10.90%
80	Mistral-7B	(8192, 4096, 14336)	-0.65%	+10.26%
70	Qwen2.5-72B	(32768, 8192, 8192)	-0.63%	+9.86%
52	Qwen2.5-7B	(8192, 3584, 18944)	-0.57%	+11.57%
34	Llama-3.1-8B	(32768, 4096, 4096)	-0.55%	+11.36%
24	Llama-2-70B	(16384, 8192, 28672)	-0.35%	+8.63%

Observations on the regression pattern:

• Most regressions occur on shapes with very large K (≥ 14336) or very large M (≥ 32768).
• Baseline forward throughput on these shapes is already close to peak (>650 TFLOPS), leaving limited headroom.
• The shape distribution suggests autotuning configurations are sub-optimal at the extremes; revisiting the tuning grid for large-K / large-M shapes is a likely follow-up.

---

6. Correctness Check

One test case has a different correctness status between the two builds.

TestID	Model	Shape (M, N, K)	main	opt
47	Llama-3.1-405B	(32768, 106496, 16384)	FAIL	PASS

The optimized kernel resolves the previously failing accuracy check on the largest 405B shape. All remaining 83 cases keep PASS status on both sides.

---

7. Conclusion

• The optimized FP8 blockwise GEMM Triton kernel delivers a mean backward speedup of 9.53% and a mean forward speedup of 1.80% across the 84-shape sweep.
• Backward TFLOPS rise from 449.38 to 496.34 (+46.96), reflecting a measurable improvement on the dominant cost component of training.
• Small- and mid-sized models gain the most; very large M / K shapes see minor forward regressions worth re-tuning.
• The optimization also fixes a numerical-correctness failure on the largest 405B shape.

---