Skip to content

Commit 34322b5

Browse files
lizamdLiclaudeLi
authored
[moe training] Optimize FP8 MoE backward pass: fused colwise kernel + AMD tuning (#4069)
* Add fused FP8 rowwise scale+cast Triton kernel for MoE forward pass Adds a fused Triton kernel that replaces the 3-kernel sequence in the forward pass of _Float8GroupedMM: 1. tensor_to_scale(A, axiswise_dim=-1) 2. A_scaled = A.to(float32) * A_scales 3. A_fp8 = to_fp8_saturated(A_scaled, fp8_dtype) The fused kernel performs per-row absmax computation and FP8 cast in a single kernel launch with two passes, benefiting from L2 cache reuse on the second pass. Benchmarked on 8x MI300X with DeepSeek-MoE-16B (EP=8, batch=4, seq=4096): - Without fused kernel: 1,865 TPS - With fused kernel: 2,153 TPS (~15% improvement) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Add detailed comments explaining fused kernel design and usage Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Fix ruff formatting in float8_rowwise.py Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Add unit test and microbenchmark for triton_fp8_rowwise_2d_scale_and_cast * Improve microbenchmark: add compiled-graph-unfused column Add a third benchmark column that wraps tensor_to_scale with torch.compiler.disable, making it opaque inside a compiled graph. This simulates the actual MoE training context where torch.compile cannot fuse across the tensor_to_scale boundary, leaving 3 separate kernel launches — exactly what triton_fp8_rowwise_2d_scale_and_cast replaces. The 'speedup vs opaque' column shows the real-world benefit. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * benchmark: simplify to torch.compile vs triton comparison Remove the compile+opaque column per reviewer feedback. The correct baseline is torch.compile of the native 3-op sequence (tensor_to_scale + multiply + to_fp8_saturated) with fullgraph optimization, compared directly against the fused triton kernel. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix: use typing.List for Python 3.9 compatibility in benchmark list[...] as a generic type requires Python 3.10+. Use List from typing to support Python 3.9. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix: ruff format spacing in benchmark print statement Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * [moe training] fuse colwise FP8 quantization in backward pass Two complementary optimizations for the MoE backward pass: 1. Add triton_fp8_per_group_colwise_scales_dual kernel: quantizes two tensors (padded_grad_output and padded_A) in a single kernel launch. The row-iteration loops for both tensors are merged, so each row is visited once per pass instead of twice. When N1 == N2 (common for square hidden dims), all blocks process both tensors simultaneously. Reduces kernel launches by 2x and cuts per-row overhead. 2. Apply triton_fp8_rowwise_2d_scale_and_cast to grad_output rowwise quantization in backward (for grad_A = grad_output @ B), replacing the 3 unfused ops (tensor_to_scale + multiply + to_fp8_saturated) with a single fused Triton kernel launch. Together these reduce backward kernel launches from 5 per step to 2: Before: tensor_to_scale + multiply + to_fp8_saturated (grad_A path) + colwise_scales(grad_output) + colwise_scales(A) After: rowwise_2d_scale_and_cast(grad_output) + colwise_scales_dual(grad_output, A) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * [moe training][AMD] fix runtime issues and eliminate Triton autotune sync overhead Three fixes for running FP8 MoE training on AMD MI300X: 1. Add FP8GroupedMMConfig alias in config.py torchtitan uses the old name from the Feb 2026 refactor (commit 4a42d32). Add a backward-compat alias so training doesn't crash on import. 2. Disable token group padding for FP8 on AMD (utils.py) fused_pad_token_groups_cuda is not available on ROCm, so pad_token_groups() falls back to torch_pad_token_groups which does a D2H sync (group_sizes.tolist()) that breaks torch.compile. FP8 grouped GEMM doesn't require 16-alignment padding the way MXFP8 requires 32-alignment, so pass pad_token_groups_for_grouped_mm=False. 3. Use single fixed Triton config on AMD (jagged_float8_scales.py) Multiple autotune configs trigger hipDeviceSynchronize for every unique (K, N_GROUPS) shape seen during training. With ~33 unique shapes per step, 3 configs = 100 D2H syncs/step that dominate the entire backward pass. Fix: use one fixed config (BLOCK_SIZE=128, BLOCK_SIZE_ITER=128, num_warps=8) for all three AMD kernels (rowwise, colwise, dual colwise). This eliminates all autotuning overhead at the cost of not finding potentially better configs for each shape — an acceptable tradeoff on MI300X where the overhead cost far exceeds any per-shape tuning benefit. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * [moe training] add pad_token_groups_for_grouped_mm to Float8TrainingOpConfig Add the pad_token_groups_for_grouped_mm field to Float8TrainingOpConfig (defaulting to False) instead of hardcoding False in utils.py. This matches the pattern already used by MXFP8TrainingOpConfig and lets callers opt-in to padding when running on hardware where the CUDA padding kernel is available. The default of False is safe on AMD/ROCm where the torch fallback triggers a D2H sync (group_sizes.tolist()) that breaks torch.compile. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * [moe training][AMD] tune colwise kernel block config on MI300X Sweep BLOCK_SIZE x BLOCK_SIZE_ITER x num_warps on representative DeepSeek-MoE-16B backward shapes (M=16640, K=2048/5120, E=64/128). Previous fixed config (BS=128, BSI=128, warps=8) was chosen by reasoning, not measurement. Benchmark shows it's 2-3x slower than the optimum. Results (MI300X, float8_e4m3fnuz): K=2048 E=64: 128/32/8 best (200 us vs 573 us old = 2.9x) K=5120 E=64: 32/128/4 best (492 us vs 1339 us old = 2.7x) K=2048 E=128: 32/32/8 best (217 us vs 417 us old = 1.9x) K=5120 E=128: 64/32/8 best (486 us vs 1005 us old = 2.1x) Best single compromise across all shapes: BS=32, BSI=128, num_warps=4 (geomean closest to per-shape optima). Also add bench_colwise_block_configs.py sweep benchmark. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * [moe training] add benchmark for dual colwise FP8 scales kernel Benchmarks triton_fp8_per_group_colwise_scales_dual vs two sequential triton_fp8_per_group_colwise_scales calls, across representative MoE backward shapes. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Remove unrelated files accidentally included in branch Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Fix ruff formatting in bench_colwise_block_configs.py Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Remove FP8GroupedMMConfig BC alias, will update torchtitan separately Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * [moe training] replace subprocess benchmarking with standard do_bench approach Removes multi-process BENCH_CFG mechanism in bench_colwise_block_configs.py and replaces with standard in-process do_bench pattern consistent with other benchmark files in this directory. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * [moe training] Assert pad_token_groups_for_grouped_mm is False (not yet supported) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * [moe training] Skip distributed tests on ROCm CI (requires world_size=4) ROCm CI has only 1 device, but test_distributed.py requires world_size=4. Add a module-level pytest.skip for ROCm to prevent CI failures. The ep/ and mxfp8/ distributed tests are already implicitly skipped on ROCm via their is_sm_at_least_100() / CUDA SM100 guards. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: Li <lizli102@ctr2-alola-login-01.amd.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Li <lizli102@ctr2-alola-ctrl-01.amd.com>
1 parent d5894a7 commit 34322b5

11 files changed

Lines changed: 1305 additions & 50 deletions

File tree

‎benchmarks/prototype/moe_training/fp8_rowwise/bench_colwise_block_configs.py‎

Lines changed: 268 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,268 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD 3-Clause license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
#
7+
# Sweeps (BLOCK_SIZE, BLOCK_SIZE_ITER, num_warps) for the colwise FP8 kernel
8+
# on representative DeepSeek-MoE-16B training shapes (MI300X).
9+
#
10+
# Usage:
11+
# cd ~/ao
12+
# python benchmarks/prototype/moe_training/fp8_rowwise/bench_colwise_block_configs.py
13+
14+
import itertools
15+
from dataclasses import dataclass
16+
from typing import List
17+
18+
import torch
19+
import triton
20+
import triton.language as tl
21+
from tabulate import tabulate
22+
from triton.testing import do_bench
23+
24+
from torchao.prototype.moe_training.utils import generate_jagged_offs
25+
26+
EPS = 1e-12
27+
FP8_DTYPE_MAP = {
28+
torch.float8_e4m3fn: tl.float8e4nv,
29+
torch.float8_e4m3fnuz: tl.float8e4b8,
30+
}
31+
32+
device = torch.device("cuda")
33+
34+
35+
# ---------------------------------------------------------------------------
36+
# Standalone colwise kernel (no autotune wrapper) so we can inject any config.
37+
# Mirrors the production _triton_fp8_per_group_colwise_scales_kernel exactly.
38+
# Input shape: (K, N) where K=token rows (jagged dim), N=hidden cols.
39+
# BLOCK_SIZE tiles N; BLOCK_SIZE_ITER tiles K in the inner loop.
40+
# ---------------------------------------------------------------------------
41+
@triton.jit
42+
def _colwise_kernel(
43+
input_ptr,
44+
offsets_ptr,
45+
out_ptr,
46+
scales_ptr,
47+
K: tl.int64,
48+
N: tl.int64,
49+
N_GROUPS: tl.int64,
50+
str_ir: tl.int64,
51+
str_ic: tl.int64,
52+
str_or: tl.int64,
53+
str_oc: tl.int64,
54+
fp8_min: tl.constexpr,
55+
fp8_max: tl.constexpr,
56+
in_dtype: tl.constexpr,
57+
out_dtype: tl.constexpr,
58+
BLOCK_SIZE: tl.constexpr,
59+
BLOCK_SIZE_ITER: tl.constexpr,
60+
EPS: tl.constexpr,
61+
):
62+
bcol = tl.program_id(0)
63+
gidx = tl.program_id(1)
64+
rs = tl.load(offsets_ptr + gidx - 1, mask=gidx > 0, other=0)
65+
re = tl.load(offsets_ptr + gidx)
66+
col_offs = (bcol * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)).to(tl.int64)
67+
amax = tl.zeros((BLOCK_SIZE,), dtype=in_dtype)
68+
for r0 in range(rs, re, BLOCK_SIZE_ITER):
69+
row_offs = (r0 + tl.arange(0, BLOCK_SIZE_ITER)).to(tl.int64)
70+
offs = row_offs[:, None] * str_ir + col_offs[None, :] * str_ic
71+
mask = (row_offs[:, None] < re) & (col_offs[None, :] < N)
72+
d = tl.load(input_ptr + offs, mask=mask, other=0.0).to(in_dtype)
73+
amax = tl.maximum(amax, tl.max(tl.abs(d), axis=0)).to(in_dtype)
74+
amax = amax.to(tl.float64)
75+
s = (fp8_max / tl.clamp(amax, min=EPS, max=float("inf"))).to(tl.float32)
76+
s = tl.exp2(tl.floor(tl.log2(s)))
77+
sc_offs = col_offs + N * gidx
78+
sc_mask = tl.arange(0, BLOCK_SIZE) < N
79+
tl.store(scales_ptr + sc_offs, s, mask=sc_mask)
80+
for r0 in range(rs, re, BLOCK_SIZE_ITER):
81+
row_offs = (r0 + tl.arange(0, BLOCK_SIZE_ITER)).to(tl.int64)
82+
offs = row_offs[:, None] * str_ir + col_offs[None, :] * str_ic
83+
mask = (row_offs[:, None] < re) & (col_offs[None, :] < N)
84+
d = tl.load(input_ptr + offs, mask=mask, other=0.0).to(in_dtype)
85+
sd = d * s[None, :]
86+
fp8d = tl.clamp(sd, min=fp8_min, max=fp8_max).to(out_dtype)
87+
o_offs = row_offs[:, None] * str_or + col_offs[None, :] * str_oc
88+
tl.store(out_ptr + o_offs, fp8d, mask=mask)
89+
90+
91+
@dataclass(frozen=True)
92+
class ExperimentConfig:
93+
M: int # total token rows (jagged dim, K_triton)
94+
K: int # hidden cols (N_triton)
95+
n_groups: int
96+
block_size: int
97+
block_size_iter: int
98+
num_warps: int
99+
fp8_dtype: torch.dtype
100+
101+
102+
@dataclass(frozen=True)
103+
class ExperimentResult:
104+
time_us: float
105+
106+
107+
@dataclass(frozen=True)
108+
class Experiment:
109+
config: ExperimentConfig
110+
result: ExperimentResult
111+
112+
113+
SHAPES = [
114+
(16640, 2048, 64),
115+
(16640, 5120, 64),
116+
(16640, 2048, 128),
117+
(16640, 5120, 128),
118+
]
119+
120+
BLOCK_SIZES = [32, 64, 128, 256]
121+
BLOCK_SIZE_ITERS = [32, 64, 128, 256]
122+
NUM_WARPS_LIST = [4, 8]
123+
124+
125+
def get_configs(fp8_dtype: torch.dtype) -> List[ExperimentConfig]:
126+
configs = []
127+
for (M, K, n_groups), bs, bsi, nw in itertools.product(
128+
SHAPES, BLOCK_SIZES, BLOCK_SIZE_ITERS, NUM_WARPS_LIST
129+
):
130+
configs.append(
131+
ExperimentConfig(
132+
M=M,
133+
K=K,
134+
n_groups=n_groups,
135+
block_size=bs,
136+
block_size_iter=bsi,
137+
num_warps=nw,
138+
fp8_dtype=fp8_dtype,
139+
)
140+
)
141+
return configs
142+
143+
144+
def run_experiment(config: ExperimentConfig) -> ExperimentResult:
145+
M, K, n_groups = config.M, config.K, config.n_groups
146+
bs, bsi, nw = config.block_size, config.block_size_iter, config.num_warps
147+
fp8_dtype = config.fp8_dtype
148+
tl_fp8_dtype = FP8_DTYPE_MAP[fp8_dtype]
149+
150+
inp = torch.randn(M, K, dtype=torch.bfloat16, device=device)
151+
offs = generate_jagged_offs(n_groups, M, multiple_of=16)
152+
fp8_min = torch.finfo(fp8_dtype).min
153+
fp8_max = torch.finfo(fp8_dtype).max
154+
155+
out = torch.empty_like(inp, dtype=fp8_dtype).as_strided(inp.size(), (1, M))
156+
sc = torch.empty(K * n_groups, dtype=torch.float32, device=device)
157+
grid = (triton.cdiv(K, bs), n_groups)
158+
159+
def run():
160+
_colwise_kernel[grid](
161+
inp,
162+
offs,
163+
out,
164+
sc,
165+
M,
166+
K,
167+
n_groups,
168+
inp.stride(0),
169+
inp.stride(1),
170+
out.stride(0),
171+
out.stride(1),
172+
fp8_min,
173+
fp8_max,
174+
tl.bfloat16,
175+
tl_fp8_dtype,
176+
BLOCK_SIZE=bs,
177+
BLOCK_SIZE_ITER=bsi,
178+
EPS=EPS,
179+
num_warps=nw,
180+
num_stages=2,
181+
)
182+
183+
t_us = do_bench(run, return_mode="median") * 1e3
184+
return ExperimentResult(time_us=t_us)
185+
186+
187+
def print_results(experiments: List[Experiment]):
188+
# Group by (M, K, n_groups) and print best configs
189+
from collections import defaultdict
190+
191+
groups = defaultdict(list)
192+
for e in experiments:
193+
key = (e.config.M, e.config.K, e.config.n_groups)
194+
groups[key].append(e)
195+
196+
summary_rows = []
197+
for key, exps in groups.items():
198+
M, K, n_groups = key
199+
best = min(exps, key=lambda e: e.result.time_us)
200+
summary_rows.append(
201+
[
202+
M,
203+
K,
204+
n_groups,
205+
best.config.block_size,
206+
best.config.block_size_iter,
207+
best.config.num_warps,
208+
f"{best.result.time_us:.1f}",
209+
]
210+
)
211+
212+
print()
213+
print("=" * 70)
214+
print("Colwise FP8 Block Config Sweep — Best per Shape")
215+
print("=" * 70)
216+
headers = [
217+
"M",
218+
"K",
219+
"n_groups",
220+
"BLOCK_SIZE",
221+
"BLOCK_SIZE_ITER",
222+
"num_warps",
223+
"time (us)",
224+
]
225+
print(tabulate(summary_rows, headers=headers))
226+
227+
# Full results table
228+
print()
229+
print("=" * 70)
230+
print("Full Results")
231+
print("=" * 70)
232+
all_rows = []
233+
for e in experiments:
234+
c, r = e.config, e.result
235+
all_rows.append(
236+
[
237+
c.M,
238+
c.K,
239+
c.n_groups,
240+
c.block_size,
241+
c.block_size_iter,
242+
c.num_warps,
243+
f"{r.time_us:.1f}",
244+
]
245+
)
246+
print(tabulate(all_rows, headers=headers))
247+
248+
249+
def main():
250+
fp8_dtype = (
251+
torch.float8_e4m3fnuz if torch.version.hip is not None else torch.float8_e4m3fn
252+
)
253+
print(f"GPU : {torch.cuda.get_device_name()}")
254+
print(f"ROCm : {torch.version.hip}")
255+
print(f"FP8 : {fp8_dtype}")
256+
print()
257+
258+
configs = get_configs(fp8_dtype)
259+
experiments = []
260+
for config in configs:
261+
result = run_experiment(config)
262+
experiments.append(Experiment(config=config, result=result))
263+
264+
print_results(experiments)
265+
266+
267+
if __name__ == "__main__":
268+
main()

0 commit comments

Comments
 (0)