feat: optimize w8a8 block fp8 matmul for Hopper

[zikong20]

PR Category

Operator (W8A8 Block-FP8 GEMM, Hopper)

Type of Change

Performance Optimization

Description

Optimize the Hopper w8a8_block_fp8_matmul operator (per-token/per-block FP8 E4M3 inputs
with [block_n, block_k] de-quant scales) for the decode/small-batch shapes that dominate
DeepSeek-V4-Flash inference traffic. The operator previously had only a single generic GEMM
kernel whose tile geometry is a poor fit for the tall-skinny (small M, large N/K) shapes
produced by MoE-style traffic, so it loses to DeepGEMM on exactly the shapes that are called
most often.

The core of this PR is a set of purpose-built GEMM kernels for those shapes, plus a dispatch in
general_w8a8_block_fp8_matmul that routes each shape by (M, N, K) to the best-fit kernel (the
generic kernel remains the fallback). All changes are confined to w8a8_block_fp8_matmul.py
(plus the two TLE helper modules it imports); the math and output dtype are unchanged.

Specialized kernels and the optimization idea

• short_k256 — for short-K projection shapes: fully unroll the K loop and pre-fold each
  N-group's scalar b_s into a_s, removing loop overhead and the per-tile scale multiply.
• swap_ab — for skinny GEMMs (small M, large N): swap A/B to compute C.T = B @ A.T,
  so the small M lands on the free dimension of the output tile and keeps the tensor core fed
  and the MMA issue efficient.
• swap_ab_splitk — for skinny shapes with long K: on top of swap-AB, split-K along K to
  parallelize the long reduction across more SMs (summed via atomic_add), relieving the SM
  under-occupancy at small M.
• general (+ TLE 2-consumer) — for the remaining large shapes: a TMA + warp-specialized
  2-consumer TLE kernel that uses async TMA loads and a producer/consumer pipeline to overlap
  memory and compute; otherwise it falls back to the generic tuned kernel.

Each specialized kernel's tile/pipeline meta-parameters (BLOCK_M/N/K, GROUP_M, num_warps,
num_stages, SPLIT_K) are tuned offline per shape by FlagTune expand-tune and stored in
w8a8_block_fp8_matmul_hopper_expand.yaml.

Numerical safety

• Output math is unchanged: every path accumulates in FP32 and casts to the requested output
  dtype (bf16 / fp16 / fp32) exactly as the generic kernel does.
• The split-K path uses atomic_add into a zero-initialized C, so the partials sum to the
  same value (bf16 atomic accumulation is not bit-reproducible across launches, but stays
  within FP8-GEMM tolerance).
• Every specialized path is entered only when its shape guard holds; any shape outside the
  guards falls through to the generic kernel, so there is no correctness regression for
  untested geometries.
• The TLE 2-consumer path is gated by _can_use_tle_2c (FP8 dtypes, contiguous strides,
  group_k == 128, N % TLE2C_BLOCK_N_HALF == 0, group_n == TLE2C_BLOCK_N_HALF) and only
  runs where TLE + device TMA are present; otherwise it transparently falls back.

Issue

N/A

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue.
• Change is fully covered by a UT.

Performance

Hardware: Hopper (H20-class), CUDA Graph pretune harness, warmup=5 / rep=5.
Reference (DeepGEMM): DeepGEMM latency. Gems: FlagGems Triton kernels (this PR's
specialized kernels). Speedup = T_DeepGEMM / T_Gems. In the tables, Default is
before-optimization (untuned default config) and Expand is after-optimization (this PR's
specialized kernels + per-shape tuning); Speedup Gain =
(Expand_speedup − Default_speedup) / Default_speedup, i.e. how much this optimization improves
the Gems-vs-DeepGEMM ratio. Shapes come from four real DeepSeek-V4-Flash traces (prefill = 1024 /
4096 / 32768 / 65536; N ∈ {512, 1536, 4096, 8192}, K ∈ {256, 1024, 4096}, M ∈ [1, 16384]).

Where the optimization helps (the gains). The largest gains land on the decode→mid-M band
(M ≈ 36–264) where the pre-optimization path sits below DeepGEMM and the specialized kernels
pull it back to a lead:

• K == 256 (short_k256): the biggest wins on the board — e.g. 1,184,4096,256 0.576 →
  1.107 (+92.19%), 1,214,4096,256 0.615 → 1.169 (+90.08%); the 218k-count
  1,100,4096,256 goes 0.696 → 1.079 (+55.03%).
• N == 512, K == 4096 (swap_ab_splitk): e.g. 1,158,512,4096 0.771 → 1.278 (+65.76%),
  the 216k-count 1,64,512,4096 1.120 → 1.338 (+19.46%).
• N == 8192, K == 1024 at small M (swap_ab/TLE): e.g. 1,2,8192,1024 1.375 → 1.806
  (+31.35%), 1,8,8192,1024 1.363 → 1.764 (+29.42%).

Where the geometry already suits the original kernel, the optimization correctly barely moves:
on M = 1 shapes (which make up almost all of the p65536 trace) the pre-optimization path
already runs at 1.5–2.1× vs DeepGEMM, and at M ≥ 2048 it already leads ~1.0–1.2×; in both
regions the after value matches it within noise, which is why the p32768/p65536 averages are
nearly flat.

What still needs optimizing (open gap). Even after this optimization, one region stays below
DeepGEMM (speedup < 1.0): the N ∈ {4096, 8192}, K == 1024 shapes in the mid-M band
(M ≈ 36–184) — 28/165 rows in p1024 and 27/185 in p4096, e.g. 1,100,8192,1024 0.799,
1,100,4096,1024 0.901, 1,64,8192,1024 0.857. These are not fixed by tuning
meta-parameters: the bottleneck is the per-CTA pipeline efficiency of the K = 1024 reduction
at this tile geometry, which is a kernel/compiler-layer problem (WGMMA issue + LSU overlap) left
for a follow-up. The p32768/p65536 traces have no below-parity shapes.

Per-trace summary

Trace	Shapes	Default avg speedup (before)	Expand avg speedup (after)	Avg Speedup Gain
p1024d1024	165	1.022	1.168	+18.01%
p4096d1024	185	1.181	1.286	+10.62%
p32768d1024	35	1.422	1.489	+5.07%
p65536d1024	20	1.384	1.383	−0.04%

Full per-shape data
DeepSeek-V4-Flash-p65536d1024_w8a8_block_fp8_matmul_deepgemm.xlsx
DeepSeek-V4-Flash-p1024d1024_w8a8_block_fp8_matmul_deepgemm.xlsx
DeepSeek-V4-Flash-p4096d1024_w8a8_block_fp8_matmul_deepgemm.xlsx
DeepSeek-V4-Flash-p32768d1024_w8a8_block_fp8_matmul_deepgemm.xlsx

[CLAassistant]

All committers have signed the CLA.