From 4a7a4f7260824e06caca18774e454780df9602b0 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 15 Oct 2026 00:22:52 +0000 Subject: [PATCH 01/10] add benchmark script for cutedsl Signed-off-by: Kaining Zhong --- .../benchmark_quantization_template.html | 385 ++++++++ .../benchmark/benchmark_quantize_mxfp8.py | 895 ++++++++++++++++++ 2 files changed, 1280 insertions(+) create mode 100644 transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html create mode 100644 transformer_engine/common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py diff --git a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html new file mode 100644 index 00000000000..43f3695bdda --- /dev/null +++ b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html @@ -0,0 +1,385 @@ + + + + + + +Quantization Benchmark + + + +
+

+

+
+
+
+

+
+
+
+
Data table
+
+
+ + + + diff --git a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py new file mode 100644 index 00000000000..156cef65cd6 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py @@ -0,0 +1,895 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Kernel-only benchmark of the MXFP8 quantize kernels (Nsight Compute, L2 flushed). + +CuTeDSL kernels come from this source tree (no rebuild needed); CUDA kernels come from the +installed TE (set PYTHONPATH= to use your own build). Writes a CSV. + +examples: + python benchmark_quantize_mxfp8.py + python benchmark_quantize_mxfp8.py --mode manual --shapes "4096x7168;16384x4096" --impl CUDA + python benchmark_quantize_mxfp8.py --compare # main vs HEAD + python benchmark_quantize_mxfp8.py --compare --old CUDA --mode manual --shapes "4096x7168" + python benchmark_quantize_mxfp8.py --fusion "none;dact;dactbias" --format html +""" + +import argparse +import contextlib +import ctypes +import datetime +import functools +import itertools +import json +import os +import re +import shlex +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path + +# Hidden sizes of current frontier models: gpt-oss (2880, also the non-128-aligned case), +# Qwen3-235B (4096), Llama 4 / Qwen3-32B / GLM-4.5 (5120), DeepSeek-V3 / Kimi K2 (7168), +# Llama 3 70B (8192). +DEFAULT_N = [2880, 4096, 5120, 7168, 8192] +DEFAULT_M = [1024, 2048, 4096, 8192, 16384, 32768] +DEFAULT_SHAPES = list(itertools.product(DEFAULT_M, DEFAULT_N)) +# Smallest set that still exercises every output: two M (a line per chart), two N (one of +# them not 128-aligned). +DEBUG_SHAPES = list(itertools.product([1024, 4096], [2880, 4096])) +MXFP8_BLOCK = 32 +DIRECTIONS = {"rowwise": (True, False), "colwise": (False, True), "both": (True, True)} +OUTPUT_FORMATS = ("csv", "png", "html", "raw") +IMPLS = {"cutedsl": "CuTeDSL", "cuda": "CUDA"} +FUSIONS = ("none", "act", "dact", "dbias", "dactbias") +ACTIVATIONS = ("gelu", "relu", "silu", "qgelu", "srelu") +COLUMNS = [ + "impl", + "M", + "N", + "direction", + "fusion", + "time_us", + "min_us", + "rel_std_pct", + "GB/s", + "kernel", +] +TIME_METRIC = "gpu__time_duration.sum" +CUDA_EXPLANATION = """\ +Comparing CUDA (old) vs CuTeDSL (new) MXFP8 quantize kernels: + old = the CUDA kernels compiled into the currently INSTALLED TE, i.e. whatever + `import transformer_engine` resolves to (set PYTHONPATH= to use a checkout): + package: {package} + version: {version} + library: {library} + source: {git} (the library may be built from an older commit, see + the + in the version) + new = the CuTeDSL kernels of {new} in this source tree ({repo}), JIT-compiled from + source. This is NOT the CuTeDSL code of the installed TE. +""" +# `-P` keeps the cwd off sys.path, so this resolves TE exactly like the CUDA worker does. +INSTALLED_TE_PROBE = """\ +import transformer_engine as te +from transformer_engine.common import _get_shared_object_file +print(te.__file__, te.__version__, _get_shared_object_file("core"), sep="\\n") +""" + + +def _installed_te(): + """Where the installed TE, and so the CUDA kernels, come from.""" + probe = subprocess.run( + [sys.executable, "-P", "-c", INSTALLED_TE_PROBE], check=True, capture_output=True, text=True + ) + init, version, library = probe.stdout.strip().splitlines()[-3:] + package = Path(init).parent + sha = _git(package, "rev-parse", "--short", "HEAD", check=False) + branch = _git(package, "branch", "--show-current", check=False).stdout.strip() + git = ( + f"{branch or 'detached'} @ {sha.stdout.strip()}" + if sha.returncode == 0 + else "not a checkout" + ) + return {"package": package, "version": version, "library": library, "git": git} + + +class _KeepOnly(list): + """nsight-python only supports an exact-name ignore list (checked via `name in list`), + so invert it: "ignore" every kernel whose name does not match the regex.""" + + def __init__(self, pattern): + super().__init__([pattern]) + self.regex = re.compile(pattern) + + def __contains__(self, name): + return self.regex.search(name) is None + + +def _round_up(x, y): + return (x + y - 1) // y * y + + +def gbps(t_ns, m, n, direction, fusion): + """Bytes moved (bf16 input(s), fp8 + e8m0 scales out per direction) per ns, i.e. GB/s.""" + num_inputs = 2 if "dact" in fusion else 1 + num_outputs = sum(DIRECTIONS[direction]) + nbytes = num_inputs * m * n * 2 + num_outputs * (m * n + m * n // 32) + return {"GB/s": (nbytes / t_ns, "GB/s")} + + +def _cutedsl_launcher(activation): + """Returns setup(x, act_input, direction, fusion) -> launch(), calling the compiled CuTeDSL + kernel directly.""" + # pylint: disable=import-outside-toplevel + import torch + + # Not needed for the direct call below, but keeps TE's own dispatch consistent with --impl. + os.environ["NVTE_ENABLE_CUTEDSL_BACKEND"] = "1" + from transformer_engine.common.CuTeDSL import get_mxfp8_quantization_function + + # Resolve the defining module instead of naming it: it was renamed between revisions. + qc = sys.modules[get_mxfp8_quantization_function.__module__] + + print(f"CuTeDSL kernels from {qc.__file__}") + + @functools.lru_cache(maxsize=None) + def compiled(direction, fusion): + rowwise, colwise = DIRECTIONS[direction] + with_dact = "dact" in fusion + cfg = qc.MXFP8QuantizeConfig( + "BFloat16", + "Float8E4M3", + rowwise, + colwise, + with_gemm_swizzled_scales=False, + with_amax=False, + with_dbias="bias" in fusion, + with_dact=with_dact, + with_act=fusion == "act", + activation={"act": activation}.get(fusion, f"d{activation}" if with_dact else "none"), + ) + fn = qc.compile_cutedsl_function_from_cfg(cfg) + return getattr(fn, "__tvm_ffi_object__", lambda: None)() or fn + + def setup(x, act_input, direction, fusion): + m, n = x.shape + rowwise, colwise = DIRECTIONS[direction] + empty = functools.partial(torch.empty, device="cuda") + fp8, e8m0 = torch.float8_e4m3fn, torch.float8_e8m0fnu + o_row = empty((m, n), dtype=fp8) if rowwise else None + s_row = ( + empty((_round_up(m, 128), _round_up(-(-n // 32), 4)), dtype=e8m0) if rowwise else None + ) + o_col = empty((m, n), dtype=fp8) if colwise else None + s_col = ( + empty((_round_up(-(-m // 32), 4), _round_up(n, 128)), dtype=e8m0) if colwise else None + ) + act_input = act_input if "dact" in fusion else None + # One partial-sum row per CTA; a CTA covers at least one 32-row MXFP8 block. + workspace = empty((-(-m // 32), n), dtype=torch.float32) if "bias" in fusion else None + fn = compiled(direction, fusion) + stream = torch.cuda.current_stream().cuda_stream + outs = (o_row, s_row, o_col, s_col) + return lambda: fn(x, *outs, None, None, act_input, workspace, stream) + + return setup + + +def _cuda_launcher(activation): + """Returns setup(x, act_input, direction, fusion) -> launch(), running the installed TE's + CUDA kernel.""" + # pylint: disable=import-outside-toplevel + os.environ["NVTE_ENABLE_CUTEDSL_BACKEND"] = "0" + import transformer_engine.pytorch # pylint: disable=unused-import + import transformer_engine_torch as tex + from transformer_engine.common import _get_shared_object_file + from transformer_engine.pytorch import MXFP8Quantizer + + core = _get_shared_object_file("core") + print(f"CUDA kernels from {core}") + lib = ctypes.CDLL(str(core)) + if hasattr(lib, "nvte_set_cutedsl_backend"): # in case TE was imported before the env var + lib.nvte_set_cutedsl_backend.argtypes = [ctypes.c_int] + lib.nvte_set_cutedsl_backend(0) + + def setup(x, act_input, direction, fusion): + rowwise, colwise = DIRECTIONS[direction] + quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=rowwise, columnwise=colwise + ) + if fusion == "none": + y = quantizer.make_empty(x.shape, dtype=x.dtype, device=x.device) + return lambda: tex.quantize(x, quantizer, y, None) + if fusion == "act": + return lambda: getattr(tex, activation)(x, quantizer) + if fusion == "dbias": + return lambda: tex.bgrad_quantize(x, quantizer) + op = getattr(tex, f"dbias_d{activation}" if fusion == "dactbias" else f"d{activation}") + return lambda: op(x, act_input, quantizer) + + return setup + + +def run_suite(shapes, directions, fusions, activation, runs, warmup, impl): + """Profile the `impl` quantize kernel for every shape, direction and fusion; returns a + DataFrame.""" + # pylint: disable=import-outside-toplevel + import nsight + import torch + + setup = (_cuda_launcher if impl == "CUDA" else _cutedsl_launcher)(activation) + + @functools.lru_cache(maxsize=1) + def make_inputs(m, n): + """The input (or incoming gradient) and, for dact, the forward activation input.""" + return [torch.randn((m, n), dtype=torch.bfloat16, device="cuda") for _ in range(2)] + + configs = [(m, n, d, f) for (m, n) in shapes for d in directions for f in fusions] + + @nsight.analyze.kernel( + configs=configs, + runs=runs, + metrics=[TIME_METRIC], + derive_metric=gbps, + # The TE path may launch helper kernels (e.g. zeroing scale padding); keep only the cast. + ignore_kernel_list=_KeepOnly("quantize") if impl == "CUDA" else None, + clock_control="none", # locking clocks is left to the user + cache_control="all", + ) + def quantize_mxfp8(m, n, direction, fusion): + launch = setup(*make_inputs(m, n), direction, fusion) + for _ in range(warmup): # module load / first-launch costs stay out of the profile + launch() + with nsight.annotate("quantize_mxfp8"): + launch() + + # nsight fills the arguments from `configs`. + df = quantize_mxfp8().to_dataframe() # pylint: disable=no-value-for-parameter + + keys = ["m", "n", "direction", "fusion"] + time = df[df["Metric"] == TIME_METRIC] + bw = df[df["Metric"] == "GB/s"][keys + ["AvgValue"]].rename(columns={"AvgValue": "GB/s"}) + out = time.merge(bw, on=keys) + out = out.rename(columns={"m": "M", "n": "N", "RelativeStdDevPct": "rel_std_pct"}) + out["time_us"] = out["AvgValue"] / 1e3 + out["min_us"] = out["MinValue"] / 1e3 + is_cutedsl = out["Kernel"].str.startswith("kernel_cutlass") + if (is_cutedsl != (impl == "CuTeDSL")).any(): + raise RuntimeError(f"Profiled kernels are not all {impl} kernels") + out["impl"] = impl + out["kernel"] = [ + m.group(1) if (m := re.search(r"(MXFP8Quantize\w+?Kernel)", k)) else k + for k in out["Kernel"] + ] + return out[COLUMNS] + + +def _git(repo, *args, check=True): + return subprocess.run( + ["git", "-C", str(repo), *args], check=check, capture_output=True, text=True + ) + + +@contextlib.contextmanager +def _workers(repo, worker_args): + """Yields run(label, impl) -> DataFrame, which runs the benchmark in a subprocess. + + The subprocesses run a copy of this script, so every run uses the same benchmark code, + and it survives checking out a revision that does not contain this file. + """ + # pylint: disable=import-outside-toplevel + import pandas as pd + + tmpdir = Path(tempfile.mkdtemp(prefix="bench_mxfp8_")) + script = tmpdir / Path(__file__).name + shutil.copy(__file__, script) + + def run(label, impl): + csv = tmpdir / f"{label}.csv" + cmd = [sys.executable, str(script), "--repo-root", str(repo)] + cmd += ["--impl", impl, *worker_args, "--worker-csv", str(csv)] + subprocess.run(cmd, check=True) + return pd.read_csv(csv).assign(revision=label) + + try: + yield run + finally: + shutil.rmtree(tmpdir, ignore_errors=True) + + +def _with_delta(frames, base, cand): + """Concatenates both runs and adds `delta`: % time reduction of `cand` vs `base` + (positive = `cand` is faster), the same value on both rows of a config.""" + # pylint: disable=import-outside-toplevel + import pandas as pd + + keys = ["M", "N", "direction", "fusion"] + df = pd.concat(frames, ignore_index=True) + table = df.pivot_table(index=keys, columns="revision", values="time_us")[[base, cand]] + table["delta"] = (table[base] - table[cand]) / table[base] * 100 + df = df.merge(table["delta"].reset_index(), on=keys) + return df[["revision", *COLUMNS, "delta"]] + + +@contextlib.contextmanager +def _in_place_checkouts(repo): + """Yields checkout(sha), which checks out a revision in the working tree, and checks out + the original branch again on exit. Refuses to run with uncommitted changes, which would + otherwise leak into every revision.""" + dirty = _git(repo, "status", "--porcelain", "--untracked-files=no").stdout.rstrip() + if dirty: + sys.exit(f"--compare checks out revisions in place; commit or stash these first:\n{dirty}") + head = _git(repo, "symbolic-ref", "--quiet", "--short", "HEAD", check=False).stdout.strip() + head = head or _git(repo, "rev-parse", "HEAD").stdout.strip() + + def checkout(sha): + result = _git(repo, "checkout", "--quiet", sha, check=False) + if result.returncode != 0: + sys.exit(f"git checkout {sha[:12]} failed:\n{result.stderr.strip()}") + + try: + yield checkout + finally: + _git(repo, "checkout", "--quiet", head) + print(f"Restored {head}") + + +def compare(repo, old, new, worker_args): + """Run the benchmark on `old` (a revision, or "CUDA") and `new` (a revision).""" + shas = { + rev: _git(repo, "rev-parse", "--verify", f"{rev}^{{commit}}").stdout.strip() + for rev in (old, new) + if rev != "CUDA" + } + explanation = "" + frames = [] + with _in_place_checkouts(repo) as checkout, _workers(repo, worker_args) as run: + if old == "CUDA": + explanation = CUDA_EXPLANATION.format(repo=repo, new=new, **_installed_te()) + print(explanation, flush=True) + for label, rev in (("old", old), ("new", new)): + if rev == "CUDA": + # Before any checkout: the installed TE may be this very source tree. + print(f"=== {label}: CUDA (installed TE) ===", flush=True) + frames.append(run(label, "CUDA")) + else: + print(f"=== {label}: CuTeDSL @ {rev} ({shas[rev][:12]}) ===", flush=True) + checkout(shas[rev]) + frames.append(run(label, "CuTeDSL")) + print(explanation) + return _with_delta(frames, "old", "new") + + +def build_charts(df, labels, activation="gelu"): + """Renderer-agnostic chart specs: one line chart of time vs M per (fusion, direction, N). + + `labels` maps each series key (the `revision` column when comparing, else the `impl` + column) to its legend name, in drawing order. Returns JSON-serializable dicts, so the + same specs can feed matplotlib (png) or a web chart library (html). A series' `role` + ("old" or "new"; a single run is "new") decides its look in every renderer. + """ + key_col = "revision" if "revision" in df.columns else "impl" + point_cols = ["M", "time_us", "min_us", "rel_std_pct", "GB/s", "kernel"] + charts = [] + for fusion, direction in itertools.product(FUSIONS, DIRECTIONS): + group = df[(df["fusion"] == fusion) & (df["direction"] == direction)] + for n in sorted(group["N"].unique()): + rows = group[group["N"] == n].sort_values("M") + series = [] + for key, name in labels.items(): + points = rows[rows[key_col] == key][point_cols] + points = points.astype(object).where(points.notna(), None) # NaN is not JSON + series.append( + { + "name": name, + "role": "old" if key == "old" else "new", + "x": points["M"].tolist(), + "y": points["time_us"].tolist(), + "points": points.to_dict("records"), + } + ) + subtitle = next(iter(labels.values())) if len(labels) == 1 else "" + if "delta" in rows.columns: + delta = rows.drop_duplicates("M")["delta"] + word = "faster" if delta.mean() >= 0 else "slower" + subtitle = ( + f"new is {abs(delta.mean()):.1f}% {word} than old on average" + f" (per M: {delta.min():+.1f}% to {delta.max():+.1f}%, + = faster)" + ) + charts.append( + { + "id": f"{fusion}_{direction}_N{n}", + "fusion": fusion, + "direction": direction, + "N": int(n), + "title": ( + f"{fusion}({activation}) · {direction} · N={n}" + if "act" in fusion + else f"{fusion} · {direction} · N={n}" + ), + "subtitle": subtitle, + "x_label": "M", + "y_label": "time (µs)", + "x_scale": "log2", + "y_scale": "log", + "series": series, + } + ) + return charts + + +# Default palette's first two categorical slots (validated light-mode pair) plus its chrome. +CHART_STYLE = { + "new": {"color": "#2a78d6", "linestyle": "-"}, + "old": {"color": "#eb6834", "linestyle": "--"}, + "surface": "#fcfcfb", + "ink": "#0b0b0b", + "ink_secondary": "#52514e", + "muted": "#898781", + "grid": "#e1e0d9", + "axis": "#c3c2b7", +} + + +def render_png(charts, out_dir): + """Writes one png per chart spec from build_charts().""" + # pylint: disable=import-outside-toplevel + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + from matplotlib.ticker import FixedLocator, LogLocator, NullLocator, ScalarFormatter + + st = CHART_STYLE + out_dir.mkdir(parents=True, exist_ok=True) + for chart in charts: + fig, ax = plt.subplots(figsize=(7, 4.2), dpi=150) + fig.patch.set_facecolor(st["surface"]) + ax.set_facecolor(st["surface"]) + xs = sorted({x for s in chart["series"] for x in s["x"]}) + # Spread the end-of-line labels apart when the lines end close together. + ends = sorted((s["y"][-1], i) for i, s in enumerate(chart["series"]) if s["y"]) + close = len(ends) > 1 and ends[-1][0] / ends[0][0] < 1.1 + nudge = { + i: (rank - (len(ends) - 1) / 2) * 9 if close else 0 for rank, (_, i) in enumerate(ends) + } + for i, series in enumerate(chart["series"]): + look = st[series["role"]] + ax.plot( + series["x"], + series["y"], + color=look["color"], + linestyle=look["linestyle"], + linewidth=2, + marker="o", + markersize=5, + label=series["name"], + ) + if series["x"] and len(chart["series"]) > 1: # direct label at the line end + ax.annotate( + series["role"], + (series["x"][-1], series["y"][-1]), + xytext=(6, nudge[i]), + textcoords="offset points", + va="center", + fontsize=8, + color=st["ink_secondary"], + ) + ax.set_xscale("log", base=2) + ax.set_yscale("log") + ax.xaxis.set_major_locator(FixedLocator(xs)) + ax.xaxis.set_minor_locator(NullLocator()) + ax.xaxis.set_major_formatter(ScalarFormatter()) + ax.yaxis.set_major_locator(LogLocator(base=10, subs=(1, 2, 5))) + ax.yaxis.set_minor_locator(NullLocator()) + ax.yaxis.set_major_formatter(ScalarFormatter()) + ax.set_title(chart["title"], loc="left", fontsize=11, color=st["ink"], pad=18) + ax.text( + 0, + 1.02, + chart["subtitle"], + transform=ax.transAxes, + fontsize=8, + color=st["ink_secondary"], + ) + ax.set_xlabel(chart["x_label"], color=st["muted"], fontsize=9) + ax.set_ylabel(chart["y_label"], color=st["muted"], fontsize=9) + ax.tick_params(colors=st["muted"], labelsize=8, which="both") + ax.grid(True, which="major", color=st["grid"], linewidth=0.6) + for side in ("top", "right"): + ax.spines[side].set_visible(False) + for side in ("left", "bottom"): + ax.spines[side].set_color(st["axis"]) + if len(chart["series"]) > 1: + ax.legend(frameon=False, fontsize=8, labelcolor=st["ink_secondary"]) + fig.tight_layout() + fig.savefig(out_dir / f"{chart['id']}.png", facecolor=st["surface"]) + plt.close(fig) + + +def render_html(charts, out_dir, meta): + """Writes data.json (the specs from build_charts() plus `meta`) and index.html, an + interactive page that also embeds the data so it opens straight from disk.""" + out_dir.mkdir(parents=True, exist_ok=True) + data = json.dumps({"meta": meta, "charts": charts}) + (out_dir / "data.json").write_text(data, encoding="utf-8") + template = Path(__file__).with_name("benchmark_quantization_template.html") + template = template.read_text(encoding="utf-8") + html = template.replace("/*__DATA__*/", data.replace("__N.png)") + else: + meta = { + "title": "MXFP8 quantize kernel benchmark", + "description": " vs ".join(labels.values()), + "compare": is_comparison, + "activation": activation, + "commit": sha, + "command": shlex.join(["python", *sys.argv]), + } + render_html(charts, out_dir, meta) + files += ["index.html (open in a browser)", "data.json"] + df.to_csv(out_dir / "results.csv", index=False) + if args.regenerate is not None: + _write_regenerated_config(out_dir, labels, activation, sha, sources, files + ["config.txt"]) + else: + _write_config(out_dir, args, fusions, shapes, repo, sha, labels, files + ["config.txt"]) + print(f"Results written to {out_dir.resolve()}/") + + +if __name__ == "__main__": + main() From 3b127f3601e4c02ff5a987d83d5e79f475ac7bdc Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 15 Oct 2026 00:39:53 +0000 Subject: [PATCH 02/10] support swizzled scale Signed-off-by: Kaining Zhong --- .../benchmark_quantization_template.html | 21 +++-- .../benchmark/benchmark_quantize_mxfp8.py | 88 ++++++++++++------- 2 files changed, 71 insertions(+), 38 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html index 43f3695bdda..132c971b049 100644 --- a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html +++ b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html @@ -8,11 +8,12 @@ benchmark_quantize_mxfp8.py --format html; reusable for e.g. NVFP4). A script fills it by replacing the DATA placeholder below with JSON of the form {"meta": {"title", "description", "compare", "activation", "commit", "command"}, - "charts": [{"fusion", "direction", "N", "subtitle", + "charts": [{"fusion", "direction", "swizzle", "N", "subtitle", "series": [{"name", "role": "old" | "new", "points": [{"M", "time_us", "min_us", "rel_std_pct", "GB/s", "kernel"}]}]}]} - One chart per (fusion, direction, N); the page offers Speedup Ratio when meta.compare. + One chart per (fusion, direction, swizzle, N); swizzle is "on"/"off" and optional. The + page offers Speedup Ratio when meta.compare. --> @@ -117,12 +118,14 @@

const state = { fusion: CHARTS[0].fusion, direction: CHARTS[0].direction, + swizzle: CHARTS[0].swizzle ?? "off", N: CHARTS[0].N, metric: "time", }; const options = { fusion: uniq(CHARTS.map(c => c.fusion)), direction: uniq(CHARTS.map(c => c.direction)), + swizzle: uniq(CHARTS.map(c => c.swizzle ?? "off")), N: uniq(CHARTS.map(c => c.N)).sort((a, b) => a - b), metric: Object.keys(METRICS), }; @@ -142,8 +145,8 @@

// ---- controls ------------------------------------------------------------------------ function buildControls() { const box = document.getElementById("controls"); - const labels = { fusion: "Fusion", direction: "Direction", N: "N", metric: "Metric" }; - for (const key of ["fusion", "direction", "N", "metric"]) { + const labels = { fusion: "Fusion", direction: "Direction", swizzle: "Swizzle", N: "N", metric: "Metric" }; + for (const key of ["fusion", "direction", "swizzle", "N", "metric"]) { const ctl = document.createElement("div"); ctl.className = "control"; ctl.innerHTML = `${labels[key]}`; @@ -171,7 +174,8 @@

// ---- data for the current selection ------------------------------------------------ function currentChart() { - return CHARTS.find(c => c.fusion === state.fusion && c.direction === state.direction && c.N === state.N); + return CHARTS.find(c => c.fusion === state.fusion && c.direction === state.direction + && (c.swizzle ?? "off") === state.swizzle && c.N === state.N); } function seriesFor(chart) { if (state.metric !== "speedup") { @@ -216,9 +220,10 @@

host.querySelector(".empty")?.remove(); tooltip.style.display = "none"; const metric = METRICS[state.metric]; - const title = state.fusion.includes("act") && META.activation - ? `${state.fusion}(${META.activation}) · ${state.direction} · N=${state.N}` - : `${state.fusion} · ${state.direction} · N=${state.N}`; + const fusion = state.fusion.includes("act") && META.activation + ? `${state.fusion}(${META.activation})` : state.fusion; + const title = [fusion, state.direction, ...(state.swizzle === "on" ? ["swizzled"] : []), `N=${state.N}`] + .join(" · "); document.getElementById("chart-title").textContent = `${metric.label} — ${title}`; document.getElementById("chart-sub").textContent = chart ? chart.subtitle : ""; const legend = document.getElementById("legend"); diff --git a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py index 156cef65cd6..31b347bf453 100644 --- a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py @@ -45,6 +45,7 @@ OUTPUT_FORMATS = ("csv", "png", "html", "raw") IMPLS = {"cutedsl": "CuTeDSL", "cuda": "CUDA"} FUSIONS = ("none", "act", "dact", "dbias", "dactbias") +SWIZZLES = ("off", "on") # scales in the GEMM-ready swizzled layout ACTIVATIONS = ("gelu", "relu", "silu", "qgelu", "srelu") COLUMNS = [ "impl", @@ -52,6 +53,7 @@ "N", "direction", "fusion", + "swizzle", "time_us", "min_us", "rel_std_pct", @@ -112,7 +114,7 @@ def _round_up(x, y): return (x + y - 1) // y * y -def gbps(t_ns, m, n, direction, fusion): +def gbps(t_ns, m, n, direction, fusion, _swizzle): """Bytes moved (bf16 input(s), fp8 + e8m0 scales out per direction) per ns, i.e. GB/s.""" num_inputs = 2 if "dact" in fusion else 1 num_outputs = sum(DIRECTIONS[direction]) @@ -121,7 +123,7 @@ def gbps(t_ns, m, n, direction, fusion): def _cutedsl_launcher(activation): - """Returns setup(x, act_input, direction, fusion) -> launch(), calling the compiled CuTeDSL + """Returns setup(x, act_input, direction, fusion, swizzle) -> launch(), calling the compiled CuTeDSL kernel directly.""" # pylint: disable=import-outside-toplevel import torch @@ -136,7 +138,7 @@ def _cutedsl_launcher(activation): print(f"CuTeDSL kernels from {qc.__file__}") @functools.lru_cache(maxsize=None) - def compiled(direction, fusion): + def compiled(direction, fusion, swizzle): rowwise, colwise = DIRECTIONS[direction] with_dact = "dact" in fusion cfg = qc.MXFP8QuantizeConfig( @@ -144,7 +146,7 @@ def compiled(direction, fusion): "Float8E4M3", rowwise, colwise, - with_gemm_swizzled_scales=False, + with_gemm_swizzled_scales=swizzle == "on", with_amax=False, with_dbias="bias" in fusion, with_dact=with_dact, @@ -154,7 +156,7 @@ def compiled(direction, fusion): fn = qc.compile_cutedsl_function_from_cfg(cfg) return getattr(fn, "__tvm_ffi_object__", lambda: None)() or fn - def setup(x, act_input, direction, fusion): + def setup(x, act_input, direction, fusion, swizzle): m, n = x.shape rowwise, colwise = DIRECTIONS[direction] empty = functools.partial(torch.empty, device="cuda") @@ -170,7 +172,7 @@ def setup(x, act_input, direction, fusion): act_input = act_input if "dact" in fusion else None # One partial-sum row per CTA; a CTA covers at least one 32-row MXFP8 block. workspace = empty((-(-m // 32), n), dtype=torch.float32) if "bias" in fusion else None - fn = compiled(direction, fusion) + fn = compiled(direction, fusion, swizzle) stream = torch.cuda.current_stream().cuda_stream outs = (o_row, s_row, o_col, s_col) return lambda: fn(x, *outs, None, None, act_input, workspace, stream) @@ -179,7 +181,7 @@ def setup(x, act_input, direction, fusion): def _cuda_launcher(activation): - """Returns setup(x, act_input, direction, fusion) -> launch(), running the installed TE's + """Returns setup(x, act_input, direction, fusion, swizzle) -> launch(), running the installed TE's CUDA kernel.""" # pylint: disable=import-outside-toplevel os.environ["NVTE_ENABLE_CUTEDSL_BACKEND"] = "0" @@ -195,11 +197,12 @@ def _cuda_launcher(activation): lib.nvte_set_cutedsl_backend.argtypes = [ctypes.c_int] lib.nvte_set_cutedsl_backend(0) - def setup(x, act_input, direction, fusion): + def setup(x, act_input, direction, fusion, swizzle): rowwise, colwise = DIRECTIONS[direction] quantizer = MXFP8Quantizer( fp8_dtype=tex.DType.kFloat8E4M3, rowwise=rowwise, columnwise=colwise ) + quantizer.optimize_for_gemm = swizzle == "on" if fusion == "none": y = quantizer.make_empty(x.shape, dtype=x.dtype, device=x.device) return lambda: tex.quantize(x, quantizer, y, None) @@ -213,9 +216,9 @@ def setup(x, act_input, direction, fusion): return setup -def run_suite(shapes, directions, fusions, activation, runs, warmup, impl): - """Profile the `impl` quantize kernel for every shape, direction and fusion; returns a - DataFrame.""" +def run_suite(shapes, directions, fusions, swizzles, activation, runs, warmup, impl): + """Profile the `impl` quantize kernel for every shape, direction, fusion and swizzle; + returns a DataFrame.""" # pylint: disable=import-outside-toplevel import nsight import torch @@ -227,7 +230,13 @@ def make_inputs(m, n): """The input (or incoming gradient) and, for dact, the forward activation input.""" return [torch.randn((m, n), dtype=torch.bfloat16, device="cuda") for _ in range(2)] - configs = [(m, n, d, f) for (m, n) in shapes for d in directions for f in fusions] + configs = [ + (m, n, d, f, sw) + for (m, n) in shapes + for d in directions + for f in fusions + for sw in swizzles + ] @nsight.analyze.kernel( configs=configs, @@ -239,8 +248,8 @@ def make_inputs(m, n): clock_control="none", # locking clocks is left to the user cache_control="all", ) - def quantize_mxfp8(m, n, direction, fusion): - launch = setup(*make_inputs(m, n), direction, fusion) + def quantize_mxfp8(m, n, direction, fusion, swizzle): + launch = setup(*make_inputs(m, n), direction, fusion, swizzle) for _ in range(warmup): # module load / first-launch costs stay out of the profile launch() with nsight.annotate("quantize_mxfp8"): @@ -249,7 +258,7 @@ def quantize_mxfp8(m, n, direction, fusion): # nsight fills the arguments from `configs`. df = quantize_mxfp8().to_dataframe() # pylint: disable=no-value-for-parameter - keys = ["m", "n", "direction", "fusion"] + keys = ["m", "n", "direction", "fusion", "swizzle"] time = df[df["Metric"] == TIME_METRIC] bw = df[df["Metric"] == "GB/s"][keys + ["AvgValue"]].rename(columns={"AvgValue": "GB/s"}) out = time.merge(bw, on=keys) @@ -261,7 +270,7 @@ def quantize_mxfp8(m, n, direction, fusion): raise RuntimeError(f"Profiled kernels are not all {impl} kernels") out["impl"] = impl out["kernel"] = [ - m.group(1) if (m := re.search(r"(MXFP8Quantize\w+?Kernel)", k)) else k + m.group(1) if (m := re.search(r"(MXFP8Quantize\w*?Kernel)", k)) else k for k in out["Kernel"] ] return out[COLUMNS] @@ -306,7 +315,7 @@ def _with_delta(frames, base, cand): # pylint: disable=import-outside-toplevel import pandas as pd - keys = ["M", "N", "direction", "fusion"] + keys = ["M", "N", "direction", "fusion", "swizzle"] df = pd.concat(frames, ignore_index=True) table = df.pivot_table(index=keys, columns="revision", values="time_us")[[base, cand]] table["delta"] = (table[base] - table[cand]) / table[base] * 100 @@ -364,7 +373,8 @@ def compare(repo, old, new, worker_args): def build_charts(df, labels, activation="gelu"): - """Renderer-agnostic chart specs: one line chart of time vs M per (fusion, direction, N). + """Renderer-agnostic chart specs: one line chart of time vs M per (fusion, direction, + swizzle, N). `labels` maps each series key (the `revision` column when comparing, else the `impl` column) to its legend name, in drawing order. Returns JSON-serializable dicts, so the @@ -374,8 +384,10 @@ def build_charts(df, labels, activation="gelu"): key_col = "revision" if "revision" in df.columns else "impl" point_cols = ["M", "time_us", "min_us", "rel_std_pct", "GB/s", "kernel"] charts = [] - for fusion, direction in itertools.product(FUSIONS, DIRECTIONS): - group = df[(df["fusion"] == fusion) & (df["direction"] == direction)] + for fusion, direction, swizzle in itertools.product(FUSIONS, DIRECTIONS, SWIZZLES): + group = df[ + (df["fusion"] == fusion) & (df["direction"] == direction) & (df["swizzle"] == swizzle) + ] for n in sorted(group["N"].unique()): rows = group[group["N"] == n].sort_values("M") series = [] @@ -401,14 +413,15 @@ def build_charts(df, labels, activation="gelu"): ) charts.append( { - "id": f"{fusion}_{direction}_N{n}", + "id": f"{fusion}_{direction}{'_swizzled' if swizzle == 'on' else ''}_N{n}", "fusion": fusion, "direction": direction, + "swizzle": swizzle, "N": int(n), - "title": ( - f"{fusion}({activation}) · {direction} · N={n}" - if "act" in fusion - else f"{fusion} · {direction} · N={n}" + "title": " · ".join( + [f"{fusion}({activation})" if "act" in fusion else fusion, direction] + + (["swizzled"] if swizzle == "on" else []) + + [f"N={n}"] ), "subtitle": subtitle, "x_label": "M", @@ -539,7 +552,7 @@ def _gpu(): return f"{name} (driver {driver})" -def _write_config(out_dir, args, fusions, shapes, repo, commit, labels, files): +def _write_config(out_dir, args, fusions, swizzles, shapes, repo, commit, labels, files): """Writes config.txt: what was benchmarked, on which commit and GPU.""" def rev(name): @@ -579,6 +592,7 @@ def installed(): f"shapes: {shape_text}", f"directions: {', '.join(args.direction)}", f"fusions: {', '.join(fusions)}", + f"swizzle: {', '.join(swizzles)}", ] if {"act", "dact", "dactbias"} & set(fusions): lines.append(f"activation: {args.activation or 'gelu'}") @@ -606,7 +620,10 @@ def _load_results(folder): fields = {} for key, value in re.findall(r"^(\w+):\s+(.*)$", text, re.M): fields.setdefault(key, value) # the top-level entry, not an embedded source config - return pd.read_csv(folder / "results.csv"), fields, text + df = pd.read_csv(folder / "results.csv") + if "swizzle" not in df.columns: # written before --swizzle existed + df.insert(df.columns.get_loc("fusion") + 1, "swizzle", "off") + return df, fields, text def _regenerate_inputs(args): @@ -729,6 +746,13 @@ def main(): metavar="all|F;F;...", help=f"all, or a ;-separated subset of {';'.join(FUSIONS)} (default: none)", ) + parser.add_argument( + "--swizzle", + type=str.lower, + choices=["off", "on", "both"], + default="off", + help="scales in the GEMM-ready swizzled layout (default: off)", + ) parser.add_argument( "--activation", choices=ACTIVATIONS, help="only with act/dact/dactbias (default: gelu)" ) @@ -788,6 +812,7 @@ def main(): f"--fusion takes all or a ;-separated list of {', '.join(FUSIONS)}, got" f" {args.fusion!r}" ) + swizzles = list(SWIZZLES) if args.swizzle == "both" else [args.swizzle] if args.activation and not {"act", "dact", "dactbias"} & set(fusions): parser.error("--activation only applies to the act, dact and dactbias fusions") if not args.compare and (args.old or args.new): @@ -827,7 +852,7 @@ def main(): worker_args = ["--mode", args.mode, *(["--shapes", args.shapes] if args.shapes else [])] worker_args += ["--runs", str(args.runs), "--warmup", str(args.warmup)] worker_args += ["--direction", *args.direction] - worker_args += ["--fusion", ";".join(fusions)] + worker_args += ["--fusion", ";".join(fusions), "--swizzle", args.swizzle] worker_args += ["--activation", args.activation] if args.activation else [] df = compare(repo, args.old, args.new, worker_args) else: @@ -839,6 +864,7 @@ def main(): shapes, args.direction, fusions, + swizzles, args.activation or "gelu", args.runs, args.warmup, @@ -851,7 +877,7 @@ def main(): is_comparison = "revision" in df.columns if args.format == "raw": if is_comparison: - keys = ["M", "N", "direction", "fusion"] + keys = ["M", "N", "direction", "fusion", "swizzle"] table = df.pivot_table(index=keys, columns="revision", values="time_us")[["old", "new"]] table["delta"] = df.groupby(keys)["delta"].first() print( @@ -887,7 +913,9 @@ def main(): if args.regenerate is not None: _write_regenerated_config(out_dir, labels, activation, sha, sources, files + ["config.txt"]) else: - _write_config(out_dir, args, fusions, shapes, repo, sha, labels, files + ["config.txt"]) + _write_config( + out_dir, args, fusions, swizzles, shapes, repo, sha, labels, files + ["config.txt"] + ) print(f"Results written to {out_dir.resolve()}/") From ea9a4360dded434dc4aeee1ae55112290bc3742f Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 15 Oct 2026 00:42:04 +0000 Subject: [PATCH 03/10] fix Signed-off-by: Kaining Zhong --- .../CuTeDSL/benchmark/benchmark_quantization_template.html | 4 ++-- .../common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html index 132c971b049..5ce88343ecc 100644 --- a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html +++ b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html @@ -145,7 +145,7 @@

// ---- controls ------------------------------------------------------------------------ function buildControls() { const box = document.getElementById("controls"); - const labels = { fusion: "Fusion", direction: "Direction", swizzle: "Swizzle", N: "N", metric: "Metric" }; + const labels = { fusion: "Fusion", direction: "Direction", swizzle: "Swizzled Scale", N: "N", metric: "Metric" }; for (const key of ["fusion", "direction", "swizzle", "N", "metric"]) { const ctl = document.createElement("div"); ctl.className = "control"; @@ -222,7 +222,7 @@

const metric = METRICS[state.metric]; const fusion = state.fusion.includes("act") && META.activation ? `${state.fusion}(${META.activation})` : state.fusion; - const title = [fusion, state.direction, ...(state.swizzle === "on" ? ["swizzled"] : []), `N=${state.N}`] + const title = [fusion, state.direction, ...(state.swizzle === "on" ? ["swizzled scale"] : []), `N=${state.N}`] .join(" · "); document.getElementById("chart-title").textContent = `${metric.label} — ${title}`; document.getElementById("chart-sub").textContent = chart ? chart.subtitle : ""; diff --git a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py index 31b347bf453..3147f3d6c4f 100644 --- a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py @@ -245,7 +245,7 @@ def make_inputs(m, n): derive_metric=gbps, # The TE path may launch helper kernels (e.g. zeroing scale padding); keep only the cast. ignore_kernel_list=_KeepOnly("quantize") if impl == "CUDA" else None, - clock_control="none", # locking clocks is left to the user + clock_control="base", # repeatable, but base clocks are below the usual boost clock cache_control="all", ) def quantize_mxfp8(m, n, direction, fusion, swizzle): @@ -420,7 +420,7 @@ def build_charts(df, labels, activation="gelu"): "N": int(n), "title": " · ".join( [f"{fusion}({activation})" if "act" in fusion else fusion, direction] - + (["swizzled"] if swizzle == "on" else []) + + (["swizzled scale"] if swizzle == "on" else []) + [f"N={n}"] ), "subtitle": subtitle, @@ -600,7 +600,7 @@ def installed(): f"runs: {args.runs} profiled + {args.warmup} warm-up launches per config", ( "timing: Nsight Compute gpu__time_duration.sum, L2 flushed before each run," - " clocks as set by the user" + " GPU clocks locked to base by Nsight Compute" ), f"files: {', '.join(files)}", ] From 1a507fdafa95ed953f868a860129467c4117eccc Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 15 Oct 2026 00:49:22 +0000 Subject: [PATCH 04/10] fix Signed-off-by: Kaining Zhong --- .../common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py index 3147f3d6c4f..0c8a8ffc382 100644 --- a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py @@ -245,7 +245,7 @@ def make_inputs(m, n): derive_metric=gbps, # The TE path may launch helper kernels (e.g. zeroing scale padding); keep only the cast. ignore_kernel_list=_KeepOnly("quantize") if impl == "CUDA" else None, - clock_control="base", # repeatable, but base clocks are below the usual boost clock + clock_control="none", # locking clocks is left to the user cache_control="all", ) def quantize_mxfp8(m, n, direction, fusion, swizzle): @@ -600,7 +600,7 @@ def installed(): f"runs: {args.runs} profiled + {args.warmup} warm-up launches per config", ( "timing: Nsight Compute gpu__time_duration.sum, L2 flushed before each run," - " GPU clocks locked to base by Nsight Compute" + " clocks as set by the user" ), f"files: {', '.join(files)}", ] From 4827d0d83b6cda7b55649544a5831d4cfc92c407 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 15 Oct 2026 18:48:21 +0000 Subject: [PATCH 05/10] switch to torch profiler which is faster Signed-off-by: Kaining Zhong --- .../benchmark_quantization_template.html | 15 +- .../benchmark/benchmark_quantize_mxfp8.py | 225 ++++++++++++------ 2 files changed, 161 insertions(+), 79 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html index 5ce88343ecc..b3d4cc4f930 100644 --- a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html +++ b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html @@ -8,11 +8,12 @@ benchmark_quantize_mxfp8.py --format html; reusable for e.g. NVFP4). A script fills it by replacing the DATA placeholder below with JSON of the form {"meta": {"title", "description", "compare", "activation", "commit", "command"}, - "charts": [{"fusion", "direction", "swizzle", "N", "subtitle", + "charts": [{"fusion", "direction", "swizzle", "input", "N", "subtitle", "series": [{"name", "role": "old" | "new", "points": [{"M", "time_us", "min_us", "rel_std_pct", "GB/s", "kernel"}]}]}]} - One chart per (fusion, direction, swizzle, N); swizzle is "on"/"off" and optional. The + One chart per (fusion, direction, swizzle, input, N); swizzle ("on"/"off") and input (dtype + name) are optional. The page offers Speedup Ratio when meta.compare. --> @@ -119,6 +120,7 @@

fusion: CHARTS[0].fusion, direction: CHARTS[0].direction, swizzle: CHARTS[0].swizzle ?? "off", + input: CHARTS[0].input ?? "bf16", N: CHARTS[0].N, metric: "time", }; @@ -126,6 +128,7 @@

fusion: uniq(CHARTS.map(c => c.fusion)), direction: uniq(CHARTS.map(c => c.direction)), swizzle: uniq(CHARTS.map(c => c.swizzle ?? "off")), + input: uniq(CHARTS.map(c => c.input ?? "bf16")), N: uniq(CHARTS.map(c => c.N)).sort((a, b) => a - b), metric: Object.keys(METRICS), }; @@ -145,8 +148,8 @@

// ---- controls ------------------------------------------------------------------------ function buildControls() { const box = document.getElementById("controls"); - const labels = { fusion: "Fusion", direction: "Direction", swizzle: "Swizzled Scale", N: "N", metric: "Metric" }; - for (const key of ["fusion", "direction", "swizzle", "N", "metric"]) { + const labels = { fusion: "Fusion", direction: "Direction", swizzle: "Swizzled Scale", input: "Input", N: "N", metric: "Metric" }; + for (const key of ["input", "fusion", "direction", "swizzle", "N", "metric"]) { const ctl = document.createElement("div"); ctl.className = "control"; ctl.innerHTML = `${labels[key]}`; @@ -175,7 +178,7 @@

// ---- data for the current selection ------------------------------------------------ function currentChart() { return CHARTS.find(c => c.fusion === state.fusion && c.direction === state.direction - && (c.swizzle ?? "off") === state.swizzle && c.N === state.N); + && (c.swizzle ?? "off") === state.swizzle && (c.input ?? "bf16") === state.input && c.N === state.N); } function seriesFor(chart) { if (state.metric !== "speedup") { @@ -222,7 +225,7 @@

const metric = METRICS[state.metric]; const fusion = state.fusion.includes("act") && META.activation ? `${state.fusion}(${META.activation})` : state.fusion; - const title = [fusion, state.direction, ...(state.swizzle === "on" ? ["swizzled scale"] : []), `N=${state.N}`] + const title = [state.input, fusion, state.direction, ...(state.swizzle === "on" ? ["swizzled scale"] : []), `N=${state.N}`] .join(" · "); document.getElementById("chart-title").textContent = `${metric.label} — ${title}`; document.getElementById("chart-sub").textContent = chart ? chart.subtitle : ""; diff --git a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py index 0c8a8ffc382..90022432ae3 100644 --- a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py +++ b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantize_mxfp8.py @@ -2,7 +2,7 @@ # # See LICENSE for license information. -"""Kernel-only benchmark of the MXFP8 quantize kernels (Nsight Compute, L2 flushed). +"""Kernel-only benchmark of the MXFP8 quantize kernels (CUPTI kernel timing, L2 flushed). CuTeDSL kernels come from this source tree (no rebuild needed); CUDA kernels come from the installed TE (set PYTHONPATH= to use your own build). Writes a CSV. @@ -26,6 +26,7 @@ import re import shlex import shutil +import statistics import subprocess import sys import tempfile @@ -46,6 +47,12 @@ IMPLS = {"cutedsl": "CuTeDSL", "cuda": "CUDA"} FUSIONS = ("none", "act", "dact", "dbias", "dactbias") SWIZZLES = ("off", "on") # scales in the GEMM-ready swizzled layout +# --input name -> (torch dtype name, CuTeDSL dtype name, bytes per element) +INPUTS = { + "bf16": ("bfloat16", "BFloat16", 2), + "fp16": ("float16", "Float16", 2), + "fp32": ("float32", "Float32", 4), +} ACTIVATIONS = ("gelu", "relu", "silu", "qgelu", "srelu") COLUMNS = [ "impl", @@ -54,13 +61,13 @@ "direction", "fusion", "swizzle", + "input", "time_us", "min_us", "rel_std_pct", "GB/s", "kernel", ] -TIME_METRIC = "gpu__time_duration.sum" CUDA_EXPLANATION = """\ Comparing CUDA (old) vs CuTeDSL (new) MXFP8 quantize kernels: old = the CUDA kernels compiled into the currently INSTALLED TE, i.e. whatever @@ -98,28 +105,38 @@ def _installed_te(): return {"package": package, "version": version, "library": library, "git": git} -class _KeepOnly(list): - """nsight-python only supports an exact-name ignore list (checked via `name in list`), - so invert it: "ignore" every kernel whose name does not match the regex.""" +def l2_cache_bytes(device=None): + """L2 cache size in bytes of `device` (the current CUDA device by default).""" + import torch # pylint: disable=import-outside-toplevel + + device = torch.cuda.current_device() if device is None else device + return torch.cuda.get_device_properties(device).L2_cache_size + - def __init__(self, pattern): - super().__init__([pattern]) - self.regex = re.compile(pattern) +def make_l2_flush(device=None, multiple=2): + """Returns flush(), which evicts L2 by reading a buffer `multiple` times its size; call it + before each timed launch so the kernel reads its inputs from DRAM. + + It must be a read: a write (e.g. zero_) leaves L2 full of dirty lines, and the next kernel + pays for writing them back (measured +10% on some MXFP8 quantize kernels). + """ + import torch # pylint: disable=import-outside-toplevel - def __contains__(self, name): - return self.regex.search(name) is None + device = torch.cuda.current_device() if device is None else device + buffer = torch.ones(multiple * l2_cache_bytes(device) // 4, dtype=torch.float32, device=device) + return lambda: buffer.sum() def _round_up(x, y): return (x + y - 1) // y * y -def gbps(t_ns, m, n, direction, fusion, _swizzle): - """Bytes moved (bf16 input(s), fp8 + e8m0 scales out per direction) per ns, i.e. GB/s.""" +def gbps(t_us, m, n, direction, fusion, input_dtype): + """Bytes moved (input(s), fp8 + e8m0 scales out per direction) per ns, i.e. GB/s.""" num_inputs = 2 if "dact" in fusion else 1 num_outputs = sum(DIRECTIONS[direction]) - nbytes = num_inputs * m * n * 2 + num_outputs * (m * n + m * n // 32) - return {"GB/s": (nbytes / t_ns, "GB/s")} + nbytes = num_inputs * m * n * INPUTS[input_dtype][2] + num_outputs * (m * n + m * n // 32) + return nbytes / (t_us * 1e3) def _cutedsl_launcher(activation): @@ -138,11 +155,11 @@ def _cutedsl_launcher(activation): print(f"CuTeDSL kernels from {qc.__file__}") @functools.lru_cache(maxsize=None) - def compiled(direction, fusion, swizzle): + def compiled(dtype, direction, fusion, swizzle): rowwise, colwise = DIRECTIONS[direction] with_dact = "dact" in fusion cfg = qc.MXFP8QuantizeConfig( - "BFloat16", + dtype, "Float8E4M3", rowwise, colwise, @@ -172,7 +189,8 @@ def setup(x, act_input, direction, fusion, swizzle): act_input = act_input if "dact" in fusion else None # One partial-sum row per CTA; a CTA covers at least one 32-row MXFP8 block. workspace = empty((-(-m // 32), n), dtype=torch.float32) if "bias" in fusion else None - fn = compiled(direction, fusion, swizzle) + dtype = next(v[1] for v in INPUTS.values() if getattr(torch, v[0]) == x.dtype) + fn = compiled(dtype, direction, fusion, swizzle) stream = torch.cuda.current_stream().cuda_stream outs = (o_row, s_row, o_col, s_col) return lambda: fn(x, *outs, None, None, act_input, workspace, stream) @@ -216,64 +234,90 @@ def setup(x, act_input, direction, fusion, swizzle): return setup -def run_suite(shapes, directions, fusions, swizzles, activation, runs, warmup, impl): - """Profile the `impl` quantize kernel for every shape, direction, fusion and swizzle; - returns a DataFrame.""" +def _base_name(kernel): + """`void ns::quantize_mxfp8_kernel<...>(...)` -> `quantize_mxfp8_kernel`.""" + return re.split(r"[<(]", kernel.removeprefix("void "), maxsplit=1)[0].split("::")[-1] + + +def _short_name(base): + """CuTeDSL kernel symbols encode the module path and arguments; keep only the class.""" + match = re.search(r"(MXFP8Quantize\w*?Kernel)", base) + return match.group(1) if match else base + + +def run_suite(shapes, inputs, directions, fusions, swizzles, activation, runs, warmup, impl): + """Times the `impl` quantize kernel for every shape, direction, fusion and swizzle with + CUPTI kernel timestamps (via torch.profiler); returns a DataFrame.""" # pylint: disable=import-outside-toplevel - import nsight + import pandas as pd import torch + from torch.profiler import ProfilerActivity, profile setup = (_cuda_launcher if impl == "CUDA" else _cutedsl_launcher)(activation) + # Helper kernels such as scale-padding zeroing or the dbias reduction are not timed. + prefix = "quantize" if impl == "CUDA" else "kernel_cutlass" @functools.lru_cache(maxsize=1) - def make_inputs(m, n): + def make_inputs(m, n, input_dtype): """The input (or incoming gradient) and, for dact, the forward activation input.""" - return [torch.randn((m, n), dtype=torch.bfloat16, device="cuda") for _ in range(2)] + dtype = getattr(torch, INPUTS[input_dtype][0]) + return [torch.randn((m, n), dtype=dtype, device="cuda") for _ in range(2)] + flush_l2 = make_l2_flush() configs = [ - (m, n, d, f, sw) + (m, n, inp, d, f, sw) for (m, n) in shapes + for inp in inputs for d in directions for f in fusions for sw in swizzles ] - - @nsight.analyze.kernel( - configs=configs, - runs=runs, - metrics=[TIME_METRIC], - derive_metric=gbps, - # The TE path may launch helper kernels (e.g. zeroing scale padding); keep only the cast. - ignore_kernel_list=_KeepOnly("quantize") if impl == "CUDA" else None, - clock_control="none", # locking clocks is left to the user - cache_control="all", + with profile(activities=[ProfilerActivity.CUDA]) as prof: + for m, n, input_dtype, direction, fusion, swizzle in configs: + launch = setup(*make_inputs(m, n, input_dtype), direction, fusion, swizzle) + for _ in range(warmup): # JIT compile and first-launch costs stay out of the timing + launch() + for _ in range(runs): + flush_l2() + launch() + torch.cuda.synchronize() + + kernels = sorted( + (e for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA), + key=lambda e: e.time_range.start, ) - def quantize_mxfp8(m, n, direction, fusion, swizzle): - launch = setup(*make_inputs(m, n), direction, fusion, swizzle) - for _ in range(warmup): # module load / first-launch costs stay out of the profile - launch() - with nsight.annotate("quantize_mxfp8"): - launch() - - # nsight fills the arguments from `configs`. - df = quantize_mxfp8().to_dataframe() # pylint: disable=no-value-for-parameter - - keys = ["m", "n", "direction", "fusion", "swizzle"] - time = df[df["Metric"] == TIME_METRIC] - bw = df[df["Metric"] == "GB/s"][keys + ["AvgValue"]].rename(columns={"AvgValue": "GB/s"}) - out = time.merge(bw, on=keys) - out = out.rename(columns={"m": "M", "n": "N", "RelativeStdDevPct": "rel_std_pct"}) - out["time_us"] = out["AvgValue"] / 1e3 - out["min_us"] = out["MinValue"] / 1e3 - is_cutedsl = out["Kernel"].str.startswith("kernel_cutlass") - if (is_cutedsl != (impl == "CuTeDSL")).any(): - raise RuntimeError(f"Profiled kernels are not all {impl} kernels") - out["impl"] = impl - out["kernel"] = [ - m.group(1) if (m := re.search(r"(MXFP8Quantize\w*?Kernel)", k)) else k - for k in out["Kernel"] - ] - return out[COLUMNS] + kernels = [(base, e.device_time_total) for e in kernels if (base := _base_name(e.name))] + kernels = [(base, t) for base, t in kernels if base.startswith(prefix)] + per_config = warmup + runs + if len(kernels) != per_config * len(configs): + names = sorted({_short_name(base) for base, _ in kernels}) + raise RuntimeError( + f"Expected {per_config} {impl} quantize kernels per config, got" + f" {len(kernels) / len(configs):.2f} (kernels: {names}); did it fall back to" + " another implementation or launch extra kernels?" + ) + rows = [] + for i, (m, n, input_dtype, direction, fusion, swizzle) in enumerate(configs): + timed = kernels[i * per_config + warmup : (i + 1) * per_config] + times = [t for _, t in timed] + mean = statistics.fmean(times) + rows.append( + { + "impl": impl, + "M": m, + "N": n, + "direction": direction, + "fusion": fusion, + "swizzle": swizzle, + "input": input_dtype, + "time_us": mean, + "min_us": min(times), + "rel_std_pct": statistics.stdev(times) / mean * 100 if runs > 1 else None, + "GB/s": gbps(mean, m, n, direction, fusion, input_dtype), + "kernel": _short_name(timed[-1][0]), + } + ) + return pd.DataFrame(rows, columns=COLUMNS) def _git(repo, *args, check=True): @@ -315,7 +359,7 @@ def _with_delta(frames, base, cand): # pylint: disable=import-outside-toplevel import pandas as pd - keys = ["M", "N", "direction", "fusion", "swizzle"] + keys = ["M", "N", "direction", "fusion", "swizzle", "input"] df = pd.concat(frames, ignore_index=True) table = df.pivot_table(index=keys, columns="revision", values="time_us")[[base, cand]] table["delta"] = (table[base] - table[cand]) / table[base] * 100 @@ -374,7 +418,7 @@ def compare(repo, old, new, worker_args): def build_charts(df, labels, activation="gelu"): """Renderer-agnostic chart specs: one line chart of time vs M per (fusion, direction, - swizzle, N). + swizzle, input, N). `labels` maps each series key (the `revision` column when comparing, else the `impl` column) to its legend name, in drawing order. Returns JSON-serializable dicts, so the @@ -384,9 +428,14 @@ def build_charts(df, labels, activation="gelu"): key_col = "revision" if "revision" in df.columns else "impl" point_cols = ["M", "time_us", "min_us", "rel_std_pct", "GB/s", "kernel"] charts = [] - for fusion, direction, swizzle in itertools.product(FUSIONS, DIRECTIONS, SWIZZLES): + for fusion, direction, swizzle, input_dtype in itertools.product( + FUSIONS, DIRECTIONS, SWIZZLES, INPUTS + ): group = df[ - (df["fusion"] == fusion) & (df["direction"] == direction) & (df["swizzle"] == swizzle) + (df["fusion"] == fusion) + & (df["direction"] == direction) + & (df["swizzle"] == swizzle) + & (df["input"] == input_dtype) ] for n in sorted(group["N"].unique()): rows = group[group["N"] == n].sort_values("M") @@ -413,13 +462,18 @@ def build_charts(df, labels, activation="gelu"): ) charts.append( { - "id": f"{fusion}_{direction}{'_swizzled' if swizzle == 'on' else ''}_N{n}", + "id": ( + f"{input_dtype}_{fusion}_{direction}" + f"{'_swizzled' if swizzle == 'on' else ''}_N{n}" + ), "fusion": fusion, "direction": direction, "swizzle": swizzle, + "input": input_dtype, "N": int(n), "title": " · ".join( - [f"{fusion}({activation})" if "act" in fusion else fusion, direction] + [input_dtype, f"{fusion}({activation})" if "act" in fusion else fusion] + + [direction] + (["swizzled scale"] if swizzle == "on" else []) + [f"N={n}"] ), @@ -552,7 +606,7 @@ def _gpu(): return f"{name} (driver {driver})" -def _write_config(out_dir, args, fusions, swizzles, shapes, repo, commit, labels, files): +def _write_config(out_dir, args, fusions, swizzles, inputs, shapes, repo, commit, labels, files): """Writes config.txt: what was benchmarked, on which commit and GPU.""" def rev(name): @@ -593,14 +647,15 @@ def installed(): f"directions: {', '.join(args.direction)}", f"fusions: {', '.join(fusions)}", f"swizzle: {', '.join(swizzles)}", + f"input: {', '.join(inputs)}", ] if {"act", "dact", "dactbias"} & set(fusions): lines.append(f"activation: {args.activation or 'gelu'}") lines += [ f"runs: {args.runs} profiled + {args.warmup} warm-up launches per config", ( - "timing: Nsight Compute gpu__time_duration.sum, L2 flushed before each run," - " clocks as set by the user" + "timing: CUPTI kernel timestamps (torch.profiler), L2 flushed (read) before each" + " run, clocks as set by the user" ), f"files: {', '.join(files)}", ] @@ -623,6 +678,8 @@ def _load_results(folder): df = pd.read_csv(folder / "results.csv") if "swizzle" not in df.columns: # written before --swizzle existed df.insert(df.columns.get_loc("fusion") + 1, "swizzle", "off") + if "input" not in df.columns: # written before --input existed + df.insert(df.columns.get_loc("swizzle") + 1, "input", "bf16") return df, fields, text @@ -753,10 +810,16 @@ def main(): default="off", help="scales in the GEMM-ready swizzled layout (default: off)", ) + parser.add_argument( + "--input", + default="bf16", + metavar="DTYPE;...", + help=f"input dtypes, a ;-separated subset of {';'.join(INPUTS)} (default: bf16)", + ) parser.add_argument( "--activation", choices=ACTIVATIONS, help="only with act/dact/dactbias (default: gelu)" ) - parser.add_argument("--runs", type=int, help="profiled runs per config (default: 5; debug: 1)") + parser.add_argument("--runs", type=int, help="timed runs per config (default: 20; debug: 3)") parser.add_argument( "--warmup", type=int, help="untimed launches per config (default: 5; debug: 1)" ) @@ -801,7 +864,7 @@ def main(): shapes = DEBUG_SHAPES if args.mode == "debug" else DEFAULT_SHAPES except ValueError as e: parser.error(str(e)) - args.runs = args.runs or (1 if args.mode == "debug" else 5) + args.runs = args.runs or (3 if args.mode == "debug" else 20) args.warmup = args.warmup or (1 if args.mode == "debug" else 5) if args.fusion.strip().lower() == "all": fusions = list(FUSIONS) @@ -813,6 +876,11 @@ def main(): f" {args.fusion!r}" ) swizzles = list(SWIZZLES) if args.swizzle == "both" else [args.swizzle] + inputs = [i.strip().lower() for i in args.input.split(";") if i.strip()] + if not inputs or any(i not in INPUTS for i in inputs): + parser.error( + f"--input takes a ;-separated list of {', '.join(INPUTS)}, got {args.input!r}" + ) if args.activation and not {"act", "dact", "dactbias"} & set(fusions): parser.error("--activation only applies to the act, dact and dactbias fusions") if not args.compare and (args.old or args.new): @@ -853,6 +921,7 @@ def main(): worker_args += ["--runs", str(args.runs), "--warmup", str(args.warmup)] worker_args += ["--direction", *args.direction] worker_args += ["--fusion", ";".join(fusions), "--swizzle", args.swizzle] + worker_args += ["--input", ";".join(inputs)] worker_args += ["--activation", args.activation] if args.activation else [] df = compare(repo, args.old, args.new, worker_args) else: @@ -862,6 +931,7 @@ def main(): sys.path.insert(0, str(repo)) df = run_suite( shapes, + inputs, args.direction, fusions, swizzles, @@ -877,7 +947,7 @@ def main(): is_comparison = "revision" in df.columns if args.format == "raw": if is_comparison: - keys = ["M", "N", "direction", "fusion", "swizzle"] + keys = ["M", "N", "direction", "fusion", "swizzle", "input"] table = df.pivot_table(index=keys, columns="revision", values="time_us")[["old", "new"]] table["delta"] = df.groupby(keys)["delta"].first() print( @@ -914,7 +984,16 @@ def main(): _write_regenerated_config(out_dir, labels, activation, sha, sources, files + ["config.txt"]) else: _write_config( - out_dir, args, fusions, swizzles, shapes, repo, sha, labels, files + ["config.txt"] + out_dir, + args, + fusions, + swizzles, + inputs, + shapes, + repo, + sha, + labels, + files + ["config.txt"], ) print(f"Results written to {out_dir.resolve()}/") From 94ca567ccb2eebcc65a183ab2226f731f44f9bbe Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 15 Oct 2026 22:33:48 +0000 Subject: [PATCH 06/10] update Signed-off-by: Kaining Zhong --- .../benchmark_quantization_template.html | 232 +++++++++++------- 1 file changed, 149 insertions(+), 83 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html index b3d4cc4f930..1821aa6ae04 100644 --- a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html +++ b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html @@ -25,13 +25,17 @@ :root { --page: #f9f9f7; --surface: #fcfcfb; --ink: #0b0b0b; --ink-2: #52514e; --muted: #898781; --grid: #e1e0d9; --axis: #c3c2b7; --border: rgba(11,11,11,0.10); - --new: #2a78d6; --old: #eb6834; --ref: #898781; + --ref: #898781; + /* categorical slots for one line per N, in fixed order */ + --c1: #2a78d6; --c2: #eb6834; --c3: #1baf7a; --c4: #eda100; + --c5: #e87ba4; --c6: #008300; --c7: #6250d6; --c8: #e34948; } @media (prefers-color-scheme: dark) { :root { --page: #0d0d0d; --surface: #1a1a19; --ink: #f0efec; --ink-2: #c3c2b7; --muted: #898781; --grid: #2c2c2a; --axis: #383835; --border: rgba(255,255,255,0.10); - --new: #3987e5; --old: #d95926; + --c1: #3987e5; --c2: #d95926; --c3: #199e70; --c4: #c98500; + --c5: #d55181; --c6: #008300; --c7: #9085e9; --c8: #e66767; } } * { box-sizing: border-box; } @@ -42,6 +46,7 @@ .meta { color: var(--ink-2); font-size: 12px; margin: 0; } .meta code { font-size: 11px; word-break: break-all; } .controls { display: flex; flex-wrap: wrap; gap: 12px 20px; align-items: center; margin: 20px 0 12px; } + #controls { margin-bottom: 0; } .control { display: flex; align-items: center; gap: 8px; } .control > span { color: var(--muted); font-size: 12px; text-transform: uppercase; letter-spacing: .04em; } .control { flex-wrap: wrap; max-width: 100%; } @@ -51,6 +56,8 @@ font-size: 13px; cursor: pointer; } .seg button + button { border-left: 1px solid var(--border); } .seg button[aria-pressed="true"] { background: var(--ink); color: var(--surface); } + .seg button:disabled { opacity: .35; cursor: not-allowed; } + .control > span.hint { color: var(--muted); font-size: 12px; text-transform: none; letter-spacing: 0; } .card { background: var(--surface); border: 1px solid var(--border); border-radius: 12px; padding: 16px; } .chart-head { display: flex; flex-wrap: wrap; justify-content: space-between; gap: 8px 16px; } .chart-title { font-size: 15px; font-weight: 600; margin: 0; } @@ -86,6 +93,7 @@

+

@@ -114,6 +122,8 @@

METRICS.speedup = { label: "Speedup Ratio", axis: "speedup (old time / new time)", scale: "linear", fmt: v => `${v.toFixed(3)}×` }; } +const MULTI_N_METRICS = ["throughput", "speedup"]; // comparable across N; time is not +const MAX_NS = 8; // one categorical color per N, never reused const uniq = xs => [...new Set(xs)]; const state = { @@ -121,7 +131,7 @@

direction: CHARTS[0].direction, swizzle: CHARTS[0].swizzle ?? "off", input: CHARTS[0].input ?? "bf16", - N: CHARTS[0].N, + Ns: [CHARTS[0].N], metric: "time", }; const options = { @@ -132,6 +142,8 @@

N: uniq(CHARTS.map(c => c.N)).sort((a, b) => a - b), metric: Object.keys(METRICS), }; +// Each N keeps its color whatever else is selected. +const N_COLOR = Object.fromEntries(options.N.map((n, i) => [n, `var(--c${(i % MAX_NS) + 1})`])); function fmtNum(v) { if (v == null || Number.isNaN(v)) return "–"; @@ -146,51 +158,99 @@

} // ---- controls ------------------------------------------------------------------------ +function addControl(box, key, label) { + const ctl = document.createElement("div"); + ctl.className = "control"; + ctl.innerHTML = `${label}`; + const seg = document.createElement("div"); + seg.className = "seg"; + seg.dataset.key = key; + for (const opt of options[key]) { + const b = document.createElement("button"); + b.type = "button"; + b.textContent = key === "metric" ? METRICS[opt].label : opt; + b.dataset.value = opt; + b.addEventListener("click", () => select(key, opt)); + seg.appendChild(b); + } + ctl.appendChild(seg); + box.appendChild(ctl); + return ctl; +} function buildControls() { const box = document.getElementById("controls"); - const labels = { fusion: "Fusion", direction: "Direction", swizzle: "Swizzled Scale", input: "Input", N: "N", metric: "Metric" }; - for (const key of ["input", "fusion", "direction", "swizzle", "N", "metric"]) { - const ctl = document.createElement("div"); - ctl.className = "control"; - ctl.innerHTML = `${labels[key]}`; - const seg = document.createElement("div"); - seg.className = "seg"; - seg.dataset.key = key; - for (const opt of options[key]) { - const b = document.createElement("button"); - b.type = "button"; - b.textContent = key === "metric" ? METRICS[opt].label : opt; - b.dataset.value = opt; - b.addEventListener("click", () => { state[key] = key === "N" ? Number(opt) : opt; update(); }); - seg.appendChild(b); - } - ctl.appendChild(seg); - box.appendChild(ctl); + const labels = { input: "Input", fusion: "Fusion", direction: "Direction", swizzle: "Swizzled Scale", N: "N" }; + for (const [key, label] of Object.entries(labels)) addControl(box, key, label); + const metricCtl = addControl(document.getElementById("metric-controls"), "metric", "Metric"); + metricCtl.insertAdjacentHTML("beforeend", ``); +} +function select(key, opt) { + if (key === "N") { + const n = Number(opt), Ns = state.Ns; + if (Ns.includes(n)) { if (Ns.length > 1) state.Ns = Ns.filter(v => v !== n); } + else if (Ns.length < MAX_NS) state.Ns = [...Ns, n].sort((a, b) => a - b); + if (state.Ns.length > 1 && !MULTI_N_METRICS.includes(state.metric)) state.metric = "throughput"; + } else { + state[key] = opt; } + update(); } function syncControls() { + const multi = state.Ns.length > 1; for (const seg of document.querySelectorAll(".seg")) { const key = seg.dataset.key; - for (const b of seg.children) b.setAttribute("aria-pressed", String(b.dataset.value === String(state[key]))); + for (const b of seg.children) { + const on = key === "N" ? state.Ns.includes(Number(b.dataset.value)) : b.dataset.value === String(state[key]); + b.setAttribute("aria-pressed", String(on)); + if (key === "metric") { + b.disabled = multi && !MULTI_N_METRICS.includes(b.dataset.value); + b.title = b.disabled ? "Time is only shown for a single N" : ""; + } + if (key === "N") b.disabled = !on && state.Ns.length >= MAX_NS; + } } + document.getElementById("metric-hint").textContent = + multi ? "Several N selected: Time is unavailable, select a single N to see it." : ""; } // ---- data for the current selection ------------------------------------------------ -function currentChart() { +function chartFor(n) { return CHARTS.find(c => c.fusion === state.fusion && c.direction === state.direction - && (c.swizzle ?? "off") === state.swizzle && (c.input ?? "bf16") === state.input && c.N === state.N); + && (c.swizzle ?? "off") === state.swizzle && (c.input ?? "bf16") === state.input && c.N === n); } -function seriesFor(chart) { - if (state.metric !== "speedup") { - const m = METRICS[state.metric]; - return chart.series.map(s => ({ name: s.name, role: s.role, - pts: s.points.map(p => ({ x: p.M, y: m.value(p), p })) })); - } +function speedupPoints(chart) { const old = chart.series.find(s => s.role === "old"), nw = chart.series.find(s => s.role === "new"); const byM = new Map(old.points.map(p => [p.M, p])); - const pts = nw.points.filter(p => byM.has(p.M)) + return nw.points.filter(p => byM.has(p.M)) .map(p => ({ x: p.M, y: byM.get(p.M).time_us / p.time_us, p, old: byM.get(p.M) })); - return [{ name: "new vs old", role: "new", pts }]; +} +// Each series: {name, color, dash, N, role, pts}. Color identifies the N; in a comparison the +// old line is dashed and the new one solid, in the same color. +function seriesFor(charts) { + const multi = charts.length > 1, out = []; + for (const chart of charts) { + const tag = multi ? `N=${chart.N}` : ""; + if (state.metric === "speedup") { + out.push({ name: tag || "new vs old", color: N_COLOR[chart.N], dash: false, + N: chart.N, role: "new", pts: speedupPoints(chart) }); + continue; + } + const m = METRICS[state.metric]; + for (const s of chart.series) { + const name = multi ? (COMPARE ? `${tag} ${s.role}` : tag) : s.name; + out.push({ name, color: N_COLOR[chart.N], dash: s.role === "old", + N: chart.N, role: s.role, pts: s.points.map(p => ({ x: p.M, y: m.value(p), p })) }); + } + } + return out; +} +function subtitleFor(charts) { + if (charts.length === 1) return charts[0].subtitle; + if (!COMPARE) return charts[0].series[0].name; + return "new vs old, mean time saved: " + charts.map(c => { + const pts = speedupPoints(c), d = pts.reduce((a, p) => a + (1 - 1 / p.y), 0) / pts.length * 100; + return `N=${c.N} ${d >= 0 ? "+" : ""}${d.toFixed(1)}%`; + }).join(" · "); } // ---- scales ---------------------------------------------------------------------------- @@ -215,8 +275,28 @@

} // ---- render ---------------------------------------------------------------------------- +function swatch(color, dash) { + return ``; +} +function renderLegend(charts, series) { + const legend = document.getElementById("legend"); + legend.innerHTML = ""; + const item = (color, dash, text) => legend.insertAdjacentHTML("beforeend", `${swatch(color, dash)}${text}`); + if (charts.length === 1) { + if (series.length > 1) for (const s of series) item(s.color, s.dash, s.name); + return; + } + for (const c of charts) item(N_COLOR[c.N], false, `N=${c.N}`); + if (COMPARE && state.metric !== "speedup") { // the dash, not the color, tells old from new + const names = Object.fromEntries(charts[0].series.map(s => [s.role, s.name])); + item("var(--muted)", true, names.old); + item("var(--muted)", false, names.new); + } +} + function render() { - const chart = currentChart(); + const charts = state.Ns.map(chartFor).filter(Boolean); const host = document.getElementById("chart"); const tooltip = document.getElementById("tooltip"); host.querySelector("svg")?.remove(); @@ -225,28 +305,21 @@

const metric = METRICS[state.metric]; const fusion = state.fusion.includes("act") && META.activation ? `${state.fusion}(${META.activation})` : state.fusion; - const title = [state.input, fusion, state.direction, ...(state.swizzle === "on" ? ["swizzled scale"] : []), `N=${state.N}`] - .join(" · "); + const title = [state.input, fusion, state.direction, ...(state.swizzle === "on" ? ["swizzled scale"] : []), + `N=${state.Ns.join(", ")}`].join(" · "); document.getElementById("chart-title").textContent = `${metric.label} — ${title}`; - document.getElementById("chart-sub").textContent = chart ? chart.subtitle : ""; - const legend = document.getElementById("legend"); - legend.innerHTML = ""; - if (!chart) { + document.getElementById("chart-sub").textContent = charts.length ? subtitleFor(charts) : ""; + if (!charts.length) { + document.getElementById("legend").innerHTML = ""; host.insertAdjacentHTML("beforeend", `
This combination was not benchmarked.
`); document.getElementById("table").innerHTML = ""; return; } - const series = seriesFor(chart); - if (series.length > 1) { - for (const s of series) { - legend.insertAdjacentHTML("beforeend", - `${s.name}`); - } - } + const series = seriesFor(charts); + renderLegend(charts, series); const W = Math.max(320, host.clientWidth), H = Math.round(Math.min(440, Math.max(260, W * 0.5))); - const m = { top: 12, right: 56, bottom: 40, left: 64 }; + const m = { top: 12, right: 64, bottom: 40, left: 64 }; const iw = W - m.left - m.right, ih = H - m.top - m.bottom; const xs = uniq(series.flatMap(s => s.pts.map(p => p.x))).sort((a, b) => a - b); const ys = series.flatMap(s => s.pts.map(p => p.y)).filter(v => v != null && v > 0); @@ -285,18 +358,20 @@

for (const s of series) { const pts = s.pts.filter(p => p.y != null && p.y > 0); el("path", { d: pts.map((p, i) => `${i ? "L" : "M"}${X(p.x)},${Y(p.y)}`).join(""), fill: "none", - stroke: `var(--${s.role})`, "stroke-width": 2, "stroke-linejoin": "round", - ...(s.role === "old" ? { "stroke-dasharray": "6 4" } : {}) }, svg); + stroke: s.color, "stroke-width": 2, "stroke-linejoin": "round", + ...(s.dash ? { "stroke-dasharray": "6 4" } : {}) }, svg); for (const p of pts) { - markers.push({ x: p.x, c: el("circle", { cx: X(p.x), cy: Y(p.y), r: 4, fill: `var(--${s.role})`, + markers.push({ x: p.x, c: el("circle", { cx: X(p.x), cy: Y(p.y), r: 4, fill: s.color, stroke: "var(--surface)", "stroke-width": 2 }, svg) }); } - if (series.length > 1 && pts.length) ends.push({ role: s.role, y: Y(pts[pts.length - 1].y) }); + if (series.length > 1 && pts.length) { + ends.push({ text: charts.length > 1 ? s.name : s.role, y: Y(pts[pts.length - 1].y) }); + } } ends.sort((a, b) => a.y - b.y); for (let i = 1; i < ends.length; i++) ends[i].y = Math.max(ends[i].y, ends[i - 1].y + 13); for (const e of ends) { - el("text", { x: m.left + iw + 8, y: e.y, class: "end-label", "dominant-baseline": "middle" }, svg).textContent = e.role; + el("text", { x: m.left + iw + 8, y: e.y, class: "end-label", "dominant-baseline": "middle" }, svg).textContent = e.text; } // Hover: crosshair at the nearest M, tooltip with every series' numbers there. @@ -308,7 +383,7 @@

const x = xs.reduce((a, b) => (Math.abs(X(b) - px) < Math.abs(X(a) - px) ? b : a)); cross.setAttribute("x1", X(x)); cross.setAttribute("x2", X(x)); cross.setAttribute("visibility", "visible"); for (const mk of markers) mk.c.setAttribute("r", mk.x === x ? 6 : 4); - tooltip.innerHTML = tooltipHtml(series, x); + tooltip.innerHTML = tooltipHtml(charts, series, x); tooltip.style.display = "block"; const tx = X(x) * (r.width / W), tw = tooltip.offsetWidth; tooltip.style.left = `${tx + 16 + tw > r.width ? tx - 16 - tw : tx + 16}px`; @@ -322,39 +397,26 @@

tooltip.style.display = "none"; }); - renderTable(chart); + renderTable(charts); } -function swatch(role) { - return ``; -} function detail(p) { const std = p.rel_std_pct == null ? "" : ` ± ${p.rel_std_pct.toFixed(1)}%`; return `${fmtNum(p.time_us)} µs${std} (min ${fmtNum(p.min_us)}) · ${fmtNum(p["GB/s"])} GB/s` + `
${p.kernel}`; } -function tooltipHtml(series, x) { +function tooltipHtml(charts, series, x) { let html = `
M = ${x}
`; - if (state.metric === "speedup") { - const pt = series[0].pts.find(p => p.x === x); - if (!pt) return html; - const byRole = Object.fromEntries(currentChart().series.map(s => [s.role, s.name])); - html += `
speedup${pt.y.toFixed(3)}×
`; - for (const [role, p] of [["old", pt.old], ["new", pt.p]]) { - html += `
${swatch(role)}${byRole[role]}` - + `${detail(p)}
`; - } - return html; - } const metric = METRICS[state.metric]; for (const s of series) { const pt = s.pts.find(p => p.x === x); if (!pt) continue; - html += `
${swatch(s.role)}${s.name}${metric.fmt(pt.y)}` - + `${detail(pt.p)}
`; + html += `
${swatch(s.color, s.dash)}${s.name}${metric.fmt(pt.y)}`; + html += state.metric === "speedup" + ? `old ${detail(pt.old)}
new ${detail(pt.p)}
` + : `${detail(pt.p)}
`; } - if (COMPARE) { + if (COMPARE && charts.length === 1 && state.metric !== "speedup") { const t = Object.fromEntries(series.map(s => [s.role, s.pts.find(p => p.x === x)?.p.time_us])); if (t.old && t.new) { const d = (t.old - t.new) / t.old * 100; @@ -365,16 +427,20 @@

return html; } -function renderTable(chart) { +function renderTable(charts) { const cols = ["time_us", "min_us", "rel_std_pct", "GB/s", "kernel"]; - const head = chart.series.length > 1 - ? chart.series.map(s => cols.map(c => `${s.role} ${c}`).join("")).join("") : cols.map(c => `${c}`).join(""); - const Ms = uniq(chart.series.flatMap(s => s.points.map(p => p.M))).sort((a, b) => a - b); - const rows = Ms.map(M => `${M}` + chart.series.map(s => { - const p = s.points.find(q => q.M === M) || {}; - return cols.map(c => `${typeof p[c] === "number" ? fmtNum(p[c]) : (p[c] ?? "–")}`).join(""); - }).join("") + "").join(""); - document.getElementById("table").innerHTML = `${head}${rows}
M
`; + const roles = charts[0].series.map(s => s.role); + const head = roles.length > 1 + ? roles.map(r => cols.map(c => `${r} ${c}`).join("")).join("") : cols.map(c => `${c}`).join(""); + const rows = charts.flatMap(chart => { + const Ms = uniq(chart.series.flatMap(s => s.points.map(p => p.M))).sort((a, b) => a - b); + return Ms.map(M => `${chart.N}${M}` + chart.series.map(s => { + const p = s.points.find(q => q.M === M) || {}; + return cols.map(c => `${typeof p[c] === "number" ? fmtNum(p[c]) : (p[c] ?? "–")}`).join(""); + }).join("") + ""); + }).join(""); + document.getElementById("table").innerHTML = + `${head}${rows}
NM
`; } function update() { syncControls(); render(); } From bc1924f484e428acafd1239eb28ec5fa4eda0400 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 15 Oct 2026 22:39:54 +0000 Subject: [PATCH 07/10] Revert "update" This reverts commit 94ca567ccb2eebcc65a183ab2226f731f44f9bbe. --- .../benchmark_quantization_template.html | 232 +++++++----------- 1 file changed, 83 insertions(+), 149 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html index 1821aa6ae04..b3d4cc4f930 100644 --- a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html +++ b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html @@ -25,17 +25,13 @@ :root { --page: #f9f9f7; --surface: #fcfcfb; --ink: #0b0b0b; --ink-2: #52514e; --muted: #898781; --grid: #e1e0d9; --axis: #c3c2b7; --border: rgba(11,11,11,0.10); - --ref: #898781; - /* categorical slots for one line per N, in fixed order */ - --c1: #2a78d6; --c2: #eb6834; --c3: #1baf7a; --c4: #eda100; - --c5: #e87ba4; --c6: #008300; --c7: #6250d6; --c8: #e34948; + --new: #2a78d6; --old: #eb6834; --ref: #898781; } @media (prefers-color-scheme: dark) { :root { --page: #0d0d0d; --surface: #1a1a19; --ink: #f0efec; --ink-2: #c3c2b7; --muted: #898781; --grid: #2c2c2a; --axis: #383835; --border: rgba(255,255,255,0.10); - --c1: #3987e5; --c2: #d95926; --c3: #199e70; --c4: #c98500; - --c5: #d55181; --c6: #008300; --c7: #9085e9; --c8: #e66767; + --new: #3987e5; --old: #d95926; } } * { box-sizing: border-box; } @@ -46,7 +42,6 @@ .meta { color: var(--ink-2); font-size: 12px; margin: 0; } .meta code { font-size: 11px; word-break: break-all; } .controls { display: flex; flex-wrap: wrap; gap: 12px 20px; align-items: center; margin: 20px 0 12px; } - #controls { margin-bottom: 0; } .control { display: flex; align-items: center; gap: 8px; } .control > span { color: var(--muted); font-size: 12px; text-transform: uppercase; letter-spacing: .04em; } .control { flex-wrap: wrap; max-width: 100%; } @@ -56,8 +51,6 @@ font-size: 13px; cursor: pointer; } .seg button + button { border-left: 1px solid var(--border); } .seg button[aria-pressed="true"] { background: var(--ink); color: var(--surface); } - .seg button:disabled { opacity: .35; cursor: not-allowed; } - .control > span.hint { color: var(--muted); font-size: 12px; text-transform: none; letter-spacing: 0; } .card { background: var(--surface); border: 1px solid var(--border); border-radius: 12px; padding: 16px; } .chart-head { display: flex; flex-wrap: wrap; justify-content: space-between; gap: 8px 16px; } .chart-title { font-size: 15px; font-weight: 600; margin: 0; } @@ -93,7 +86,6 @@

-

@@ -122,8 +114,6 @@

METRICS.speedup = { label: "Speedup Ratio", axis: "speedup (old time / new time)", scale: "linear", fmt: v => `${v.toFixed(3)}×` }; } -const MULTI_N_METRICS = ["throughput", "speedup"]; // comparable across N; time is not -const MAX_NS = 8; // one categorical color per N, never reused const uniq = xs => [...new Set(xs)]; const state = { @@ -131,7 +121,7 @@

direction: CHARTS[0].direction, swizzle: CHARTS[0].swizzle ?? "off", input: CHARTS[0].input ?? "bf16", - Ns: [CHARTS[0].N], + N: CHARTS[0].N, metric: "time", }; const options = { @@ -142,8 +132,6 @@

N: uniq(CHARTS.map(c => c.N)).sort((a, b) => a - b), metric: Object.keys(METRICS), }; -// Each N keeps its color whatever else is selected. -const N_COLOR = Object.fromEntries(options.N.map((n, i) => [n, `var(--c${(i % MAX_NS) + 1})`])); function fmtNum(v) { if (v == null || Number.isNaN(v)) return "–"; @@ -158,99 +146,51 @@

} // ---- controls ------------------------------------------------------------------------ -function addControl(box, key, label) { - const ctl = document.createElement("div"); - ctl.className = "control"; - ctl.innerHTML = `${label}`; - const seg = document.createElement("div"); - seg.className = "seg"; - seg.dataset.key = key; - for (const opt of options[key]) { - const b = document.createElement("button"); - b.type = "button"; - b.textContent = key === "metric" ? METRICS[opt].label : opt; - b.dataset.value = opt; - b.addEventListener("click", () => select(key, opt)); - seg.appendChild(b); - } - ctl.appendChild(seg); - box.appendChild(ctl); - return ctl; -} function buildControls() { const box = document.getElementById("controls"); - const labels = { input: "Input", fusion: "Fusion", direction: "Direction", swizzle: "Swizzled Scale", N: "N" }; - for (const [key, label] of Object.entries(labels)) addControl(box, key, label); - const metricCtl = addControl(document.getElementById("metric-controls"), "metric", "Metric"); - metricCtl.insertAdjacentHTML("beforeend", ``); -} -function select(key, opt) { - if (key === "N") { - const n = Number(opt), Ns = state.Ns; - if (Ns.includes(n)) { if (Ns.length > 1) state.Ns = Ns.filter(v => v !== n); } - else if (Ns.length < MAX_NS) state.Ns = [...Ns, n].sort((a, b) => a - b); - if (state.Ns.length > 1 && !MULTI_N_METRICS.includes(state.metric)) state.metric = "throughput"; - } else { - state[key] = opt; + const labels = { fusion: "Fusion", direction: "Direction", swizzle: "Swizzled Scale", input: "Input", N: "N", metric: "Metric" }; + for (const key of ["input", "fusion", "direction", "swizzle", "N", "metric"]) { + const ctl = document.createElement("div"); + ctl.className = "control"; + ctl.innerHTML = `${labels[key]}`; + const seg = document.createElement("div"); + seg.className = "seg"; + seg.dataset.key = key; + for (const opt of options[key]) { + const b = document.createElement("button"); + b.type = "button"; + b.textContent = key === "metric" ? METRICS[opt].label : opt; + b.dataset.value = opt; + b.addEventListener("click", () => { state[key] = key === "N" ? Number(opt) : opt; update(); }); + seg.appendChild(b); + } + ctl.appendChild(seg); + box.appendChild(ctl); } - update(); } function syncControls() { - const multi = state.Ns.length > 1; for (const seg of document.querySelectorAll(".seg")) { const key = seg.dataset.key; - for (const b of seg.children) { - const on = key === "N" ? state.Ns.includes(Number(b.dataset.value)) : b.dataset.value === String(state[key]); - b.setAttribute("aria-pressed", String(on)); - if (key === "metric") { - b.disabled = multi && !MULTI_N_METRICS.includes(b.dataset.value); - b.title = b.disabled ? "Time is only shown for a single N" : ""; - } - if (key === "N") b.disabled = !on && state.Ns.length >= MAX_NS; - } + for (const b of seg.children) b.setAttribute("aria-pressed", String(b.dataset.value === String(state[key]))); } - document.getElementById("metric-hint").textContent = - multi ? "Several N selected: Time is unavailable, select a single N to see it." : ""; } // ---- data for the current selection ------------------------------------------------ -function chartFor(n) { +function currentChart() { return CHARTS.find(c => c.fusion === state.fusion && c.direction === state.direction - && (c.swizzle ?? "off") === state.swizzle && (c.input ?? "bf16") === state.input && c.N === n); + && (c.swizzle ?? "off") === state.swizzle && (c.input ?? "bf16") === state.input && c.N === state.N); } -function speedupPoints(chart) { +function seriesFor(chart) { + if (state.metric !== "speedup") { + const m = METRICS[state.metric]; + return chart.series.map(s => ({ name: s.name, role: s.role, + pts: s.points.map(p => ({ x: p.M, y: m.value(p), p })) })); + } const old = chart.series.find(s => s.role === "old"), nw = chart.series.find(s => s.role === "new"); const byM = new Map(old.points.map(p => [p.M, p])); - return nw.points.filter(p => byM.has(p.M)) + const pts = nw.points.filter(p => byM.has(p.M)) .map(p => ({ x: p.M, y: byM.get(p.M).time_us / p.time_us, p, old: byM.get(p.M) })); -} -// Each series: {name, color, dash, N, role, pts}. Color identifies the N; in a comparison the -// old line is dashed and the new one solid, in the same color. -function seriesFor(charts) { - const multi = charts.length > 1, out = []; - for (const chart of charts) { - const tag = multi ? `N=${chart.N}` : ""; - if (state.metric === "speedup") { - out.push({ name: tag || "new vs old", color: N_COLOR[chart.N], dash: false, - N: chart.N, role: "new", pts: speedupPoints(chart) }); - continue; - } - const m = METRICS[state.metric]; - for (const s of chart.series) { - const name = multi ? (COMPARE ? `${tag} ${s.role}` : tag) : s.name; - out.push({ name, color: N_COLOR[chart.N], dash: s.role === "old", - N: chart.N, role: s.role, pts: s.points.map(p => ({ x: p.M, y: m.value(p), p })) }); - } - } - return out; -} -function subtitleFor(charts) { - if (charts.length === 1) return charts[0].subtitle; - if (!COMPARE) return charts[0].series[0].name; - return "new vs old, mean time saved: " + charts.map(c => { - const pts = speedupPoints(c), d = pts.reduce((a, p) => a + (1 - 1 / p.y), 0) / pts.length * 100; - return `N=${c.N} ${d >= 0 ? "+" : ""}${d.toFixed(1)}%`; - }).join(" · "); + return [{ name: "new vs old", role: "new", pts }]; } // ---- scales ---------------------------------------------------------------------------- @@ -275,28 +215,8 @@

} // ---- render ---------------------------------------------------------------------------- -function swatch(color, dash) { - return ``; -} -function renderLegend(charts, series) { - const legend = document.getElementById("legend"); - legend.innerHTML = ""; - const item = (color, dash, text) => legend.insertAdjacentHTML("beforeend", `${swatch(color, dash)}${text}`); - if (charts.length === 1) { - if (series.length > 1) for (const s of series) item(s.color, s.dash, s.name); - return; - } - for (const c of charts) item(N_COLOR[c.N], false, `N=${c.N}`); - if (COMPARE && state.metric !== "speedup") { // the dash, not the color, tells old from new - const names = Object.fromEntries(charts[0].series.map(s => [s.role, s.name])); - item("var(--muted)", true, names.old); - item("var(--muted)", false, names.new); - } -} - function render() { - const charts = state.Ns.map(chartFor).filter(Boolean); + const chart = currentChart(); const host = document.getElementById("chart"); const tooltip = document.getElementById("tooltip"); host.querySelector("svg")?.remove(); @@ -305,21 +225,28 @@

const metric = METRICS[state.metric]; const fusion = state.fusion.includes("act") && META.activation ? `${state.fusion}(${META.activation})` : state.fusion; - const title = [state.input, fusion, state.direction, ...(state.swizzle === "on" ? ["swizzled scale"] : []), - `N=${state.Ns.join(", ")}`].join(" · "); + const title = [state.input, fusion, state.direction, ...(state.swizzle === "on" ? ["swizzled scale"] : []), `N=${state.N}`] + .join(" · "); document.getElementById("chart-title").textContent = `${metric.label} — ${title}`; - document.getElementById("chart-sub").textContent = charts.length ? subtitleFor(charts) : ""; - if (!charts.length) { - document.getElementById("legend").innerHTML = ""; + document.getElementById("chart-sub").textContent = chart ? chart.subtitle : ""; + const legend = document.getElementById("legend"); + legend.innerHTML = ""; + if (!chart) { host.insertAdjacentHTML("beforeend", `
This combination was not benchmarked.
`); document.getElementById("table").innerHTML = ""; return; } - const series = seriesFor(charts); - renderLegend(charts, series); + const series = seriesFor(chart); + if (series.length > 1) { + for (const s of series) { + legend.insertAdjacentHTML("beforeend", + `${s.name}`); + } + } const W = Math.max(320, host.clientWidth), H = Math.round(Math.min(440, Math.max(260, W * 0.5))); - const m = { top: 12, right: 64, bottom: 40, left: 64 }; + const m = { top: 12, right: 56, bottom: 40, left: 64 }; const iw = W - m.left - m.right, ih = H - m.top - m.bottom; const xs = uniq(series.flatMap(s => s.pts.map(p => p.x))).sort((a, b) => a - b); const ys = series.flatMap(s => s.pts.map(p => p.y)).filter(v => v != null && v > 0); @@ -358,20 +285,18 @@

for (const s of series) { const pts = s.pts.filter(p => p.y != null && p.y > 0); el("path", { d: pts.map((p, i) => `${i ? "L" : "M"}${X(p.x)},${Y(p.y)}`).join(""), fill: "none", - stroke: s.color, "stroke-width": 2, "stroke-linejoin": "round", - ...(s.dash ? { "stroke-dasharray": "6 4" } : {}) }, svg); + stroke: `var(--${s.role})`, "stroke-width": 2, "stroke-linejoin": "round", + ...(s.role === "old" ? { "stroke-dasharray": "6 4" } : {}) }, svg); for (const p of pts) { - markers.push({ x: p.x, c: el("circle", { cx: X(p.x), cy: Y(p.y), r: 4, fill: s.color, + markers.push({ x: p.x, c: el("circle", { cx: X(p.x), cy: Y(p.y), r: 4, fill: `var(--${s.role})`, stroke: "var(--surface)", "stroke-width": 2 }, svg) }); } - if (series.length > 1 && pts.length) { - ends.push({ text: charts.length > 1 ? s.name : s.role, y: Y(pts[pts.length - 1].y) }); - } + if (series.length > 1 && pts.length) ends.push({ role: s.role, y: Y(pts[pts.length - 1].y) }); } ends.sort((a, b) => a.y - b.y); for (let i = 1; i < ends.length; i++) ends[i].y = Math.max(ends[i].y, ends[i - 1].y + 13); for (const e of ends) { - el("text", { x: m.left + iw + 8, y: e.y, class: "end-label", "dominant-baseline": "middle" }, svg).textContent = e.text; + el("text", { x: m.left + iw + 8, y: e.y, class: "end-label", "dominant-baseline": "middle" }, svg).textContent = e.role; } // Hover: crosshair at the nearest M, tooltip with every series' numbers there. @@ -383,7 +308,7 @@

const x = xs.reduce((a, b) => (Math.abs(X(b) - px) < Math.abs(X(a) - px) ? b : a)); cross.setAttribute("x1", X(x)); cross.setAttribute("x2", X(x)); cross.setAttribute("visibility", "visible"); for (const mk of markers) mk.c.setAttribute("r", mk.x === x ? 6 : 4); - tooltip.innerHTML = tooltipHtml(charts, series, x); + tooltip.innerHTML = tooltipHtml(series, x); tooltip.style.display = "block"; const tx = X(x) * (r.width / W), tw = tooltip.offsetWidth; tooltip.style.left = `${tx + 16 + tw > r.width ? tx - 16 - tw : tx + 16}px`; @@ -397,26 +322,39 @@

tooltip.style.display = "none"; }); - renderTable(charts); + renderTable(chart); } +function swatch(role) { + return ``; +} function detail(p) { const std = p.rel_std_pct == null ? "" : ` ± ${p.rel_std_pct.toFixed(1)}%`; return `${fmtNum(p.time_us)} µs${std} (min ${fmtNum(p.min_us)}) · ${fmtNum(p["GB/s"])} GB/s` + `
${p.kernel}`; } -function tooltipHtml(charts, series, x) { +function tooltipHtml(series, x) { let html = `
M = ${x}
`; + if (state.metric === "speedup") { + const pt = series[0].pts.find(p => p.x === x); + if (!pt) return html; + const byRole = Object.fromEntries(currentChart().series.map(s => [s.role, s.name])); + html += `
speedup${pt.y.toFixed(3)}×
`; + for (const [role, p] of [["old", pt.old], ["new", pt.p]]) { + html += `
${swatch(role)}${byRole[role]}` + + `${detail(p)}
`; + } + return html; + } const metric = METRICS[state.metric]; for (const s of series) { const pt = s.pts.find(p => p.x === x); if (!pt) continue; - html += `
${swatch(s.color, s.dash)}${s.name}${metric.fmt(pt.y)}`; - html += state.metric === "speedup" - ? `old ${detail(pt.old)}
new ${detail(pt.p)}
` - : `${detail(pt.p)}
`; + html += `
${swatch(s.role)}${s.name}${metric.fmt(pt.y)}` + + `${detail(pt.p)}
`; } - if (COMPARE && charts.length === 1 && state.metric !== "speedup") { + if (COMPARE) { const t = Object.fromEntries(series.map(s => [s.role, s.pts.find(p => p.x === x)?.p.time_us])); if (t.old && t.new) { const d = (t.old - t.new) / t.old * 100; @@ -427,20 +365,16 @@

return html; } -function renderTable(charts) { +function renderTable(chart) { const cols = ["time_us", "min_us", "rel_std_pct", "GB/s", "kernel"]; - const roles = charts[0].series.map(s => s.role); - const head = roles.length > 1 - ? roles.map(r => cols.map(c => `${r} ${c}`).join("")).join("") : cols.map(c => `${c}`).join(""); - const rows = charts.flatMap(chart => { - const Ms = uniq(chart.series.flatMap(s => s.points.map(p => p.M))).sort((a, b) => a - b); - return Ms.map(M => `${chart.N}${M}` + chart.series.map(s => { - const p = s.points.find(q => q.M === M) || {}; - return cols.map(c => `${typeof p[c] === "number" ? fmtNum(p[c]) : (p[c] ?? "–")}`).join(""); - }).join("") + ""); - }).join(""); - document.getElementById("table").innerHTML = - `${head}${rows}
NM
`; + const head = chart.series.length > 1 + ? chart.series.map(s => cols.map(c => `${s.role} ${c}`).join("")).join("") : cols.map(c => `${c}`).join(""); + const Ms = uniq(chart.series.flatMap(s => s.points.map(p => p.M))).sort((a, b) => a - b); + const rows = Ms.map(M => `${M}` + chart.series.map(s => { + const p = s.points.find(q => q.M === M) || {}; + return cols.map(c => `${typeof p[c] === "number" ? fmtNum(p[c]) : (p[c] ?? "–")}`).join(""); + }).join("") + "").join(""); + document.getElementById("table").innerHTML = `${head}${rows}
M
`; } function update() { syncControls(); render(); } From f80a628980519d9e751a8b90b8aeeab1e6dc6b9a Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Thu, 15 Oct 2026 23:09:37 +0000 Subject: [PATCH 08/10] update Signed-off-by: Kaining Zhong --- .../benchmark_quantization_template.html | 198 +++++++++++------- .../benchmark/benchmark_quantize_mxfp8.py | 4 +- 2 files changed, 124 insertions(+), 78 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html index b3d4cc4f930..c444415e9d3 100644 --- a/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html +++ b/transformer_engine/common/CuTeDSL/benchmark/benchmark_quantization_template.html @@ -13,8 +13,7 @@ "points": [{"M", "time_us", "min_us", "rel_std_pct", "GB/s", "kernel"}]}]}]} One chart per (fusion, direction, swizzle, input, N); swizzle ("on"/"off") and input (dtype - name) are optional. The - page offers Speedup Ratio when meta.compare. + name) are optional. When meta.compare, each chart shows old / new time above every M. --> @@ -25,7 +24,7 @@ :root { --page: #f9f9f7; --surface: #fcfcfb; --ink: #0b0b0b; --ink-2: #52514e; --muted: #898781; --grid: #e1e0d9; --axis: #c3c2b7; --border: rgba(11,11,11,0.10); - --new: #2a78d6; --old: #eb6834; --ref: #898781; + --new: #2a78d6; --old: #eb6834; } @media (prefers-color-scheme: dark) { :root { @@ -38,6 +37,10 @@ body { margin: 0; background: var(--page); color: var(--ink); font: 14px/1.45 system-ui, -apple-system, "Segoe UI", sans-serif; } main { max-width: 1040px; margin: 0 auto; padding: 24px 16px 48px; } + .selection { font-size: 16px; font-weight: 600; margin: 4px 0 12px; } + /* All selected charts share one row (columns set from the selection); stacked on phones. */ + .grid { display: grid; gap: 12px; grid-template-columns: repeat(var(--cols, 1), minmax(0, 1fr)); } + @media (max-width: 640px) { .grid { grid-template-columns: minmax(0, 1fr); } } h1 { font-size: 20px; margin: 0 0 4px; font-weight: 600; } .meta { color: var(--ink-2); font-size: 12px; margin: 0; } .meta code { font-size: 11px; word-break: break-all; } @@ -51,7 +54,14 @@ font-size: 13px; cursor: pointer; } .seg button + button { border-left: 1px solid var(--border); } .seg button[aria-pressed="true"] { background: var(--ink); color: var(--surface); } - .card { background: var(--surface); border: 1px solid var(--border); border-radius: 12px; padding: 16px; } + .seg button[aria-pressed="true"] + button[aria-pressed="true"] { + border-left-color: color-mix(in srgb, var(--surface) 35%, transparent); } + .card { background: var(--surface); border: 1px solid var(--border); border-radius: 12px; padding: 16px; + min-width: 0; } + .grid.dense .card { padding: 10px; } + .grid.dense .legend span { white-space: normal; } + .shared-legend { margin: -4px 0 12px; } + .shared-legend:empty { display: none; } .chart-head { display: flex; flex-wrap: wrap; justify-content: space-between; gap: 8px 16px; } .chart-title { font-size: 15px; font-weight: 600; margin: 0; } .chart-sub { color: var(--ink-2); font-size: 12px; margin: 2px 0 0; } @@ -63,6 +73,7 @@ .tick { fill: var(--muted); font-size: 11px; } .axis-label { fill: var(--muted); font-size: 12px; } .end-label { fill: var(--ink-2); font-size: 11px; } + .speedup-label { fill: var(--ink-2); font-size: 11px; font-weight: 600; font-variant-numeric: tabular-nums; } .tooltip { position: absolute; pointer-events: none; background: var(--surface); color: var(--ink); border: 1px solid var(--border); border-radius: 8px; padding: 8px 10px; font-size: 12px; box-shadow: 0 4px 16px rgba(0,0,0,.12); min-width: 220px; display: none; z-index: 2; } @@ -86,14 +97,19 @@

-
-
-

-
-
-
-
Data table
-
+

+
+
+