From c3818128ff4c4a5196c3e2e706c9e85455cf5201 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 12 Jun 2026 22:13:08 +0000 Subject: [PATCH 01/32] start draft Signed-off-by: Kaining Zhong --- build_tools/pytorch.py | 28 +- pyproject.toml | 2 +- setup.py | 7 +- tests/pytorch/mxfp8/bench_mxfp8_cutedsl.py | 517 +++++ tests/pytorch/mxfp8/run_mxfp8_benchmark.py | 477 ++++ tests/pytorch/mxfp8/run_nsys_profile.sh | 220 ++ .../mxfp8/test_mxfp8_cutedsl_backend.py | 246 +++ transformer_engine/common/CMakeLists.txt | 34 + transformer_engine/common/CuTeDSL/__init__.py | 29 + .../common/CuTeDSL/activations.py | 97 + .../common/CuTeDSL/cast/__init__.py | 5 + .../common/CuTeDSL/cast/mxfp8/__init__.py | 5 + .../CuTeDSL/cast/mxfp8/quantize_mxfp8.py | 1916 +++++++++++++++++ transformer_engine/common/CuTeDSL/utils.py | 238 ++ .../common/CuTeDSL/utils_fp8.py | 320 +++ .../common/cast/dispatch/quantize.cuh | 27 +- .../cast/mxfp8/quantize_mxfp8_cutedsl.cuh | 257 +++ .../common/transformer_engine.cpp | 12 + transformer_engine/common/tvm_ffi_bridge.h | 286 +++ transformer_engine/pytorch/__init__.py | 1 + .../pytorch/csrc/extensions/cast.cpp | 2 + .../pytorch/csrc/extensions/pybind.cpp | 14 + 22 files changed, 4732 insertions(+), 8 deletions(-) create mode 100644 tests/pytorch/mxfp8/bench_mxfp8_cutedsl.py create mode 100644 tests/pytorch/mxfp8/run_mxfp8_benchmark.py create mode 100755 tests/pytorch/mxfp8/run_nsys_profile.sh create mode 100644 tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py create mode 100644 transformer_engine/common/CuTeDSL/__init__.py create mode 100644 transformer_engine/common/CuTeDSL/activations.py create mode 100644 transformer_engine/common/CuTeDSL/cast/__init__.py create mode 100644 transformer_engine/common/CuTeDSL/cast/mxfp8/__init__.py create mode 100644 transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py create mode 100644 transformer_engine/common/CuTeDSL/utils.py create mode 100644 transformer_engine/common/CuTeDSL/utils_fp8.py create mode 100644 transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh create mode 100644 transformer_engine/common/tvm_ffi_bridge.h diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index 2bb238c522..35968f70aa 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -4,6 +4,7 @@ """PyTorch related extensions.""" +import importlib.util import os from pathlib import Path from importlib import metadata @@ -67,6 +68,25 @@ def setup_pytorch_extension( ] ) + # apache-tvm-ffi: headers for the C++ API (Module / Function / TensorView) + # and libtvm_ffi.so for symbol resolution. Used by tvm_ffi_bridge.h / + # applyTVMFunction. Python registers AOT-compiled CuTeDSL kernels into + # the global registry; TE C++ looks them up via Function::GetGlobalRequired. + tvm_ffi_spec = importlib.util.find_spec("tvm_ffi") + if tvm_ffi_spec is None or not tvm_ffi_spec.submodule_search_locations: + raise RuntimeError( + "apache-tvm-ffi package not found; install it (e.g. " + "`pip install apache-tvm-ffi`) — required for the TVM FFI bridge." + ) + tvm_ffi_root = Path(tvm_ffi_spec.submodule_search_locations[0]) + tvm_ffi_include = tvm_ffi_root / "include" + tvm_ffi_lib_dir = tvm_ffi_root / "lib" + if not tvm_ffi_include.is_dir() or not (tvm_ffi_lib_dir / "libtvm_ffi.so").exists(): + raise RuntimeError( + f"apache-tvm-ffi assets missing at {tvm_ffi_root} (need include/ and lib/libtvm_ffi.so)" + ) + include_dirs.append(tvm_ffi_include) + # Compiler flags cxx_flags = ["-O3", "-fvisibility=hidden"] if debug_build_enabled(): @@ -96,8 +116,11 @@ def setup_pytorch_extension( # it needs those macros visible. cxx_flags.append("-DUSE_NCCL") - library_dirs = [] - libraries = [] + library_dirs = [tvm_ffi_lib_dir] + libraries = ["tvm_ffi"] + # rpath pinned to the pip install dir so the loader finds libtvm_ffi.so + # without LD_LIBRARY_PATH at runtime. + extra_link_args = [f"-Wl,-rpath,{tvm_ffi_lib_dir}"] if bool(int(os.getenv("NVTE_ENABLE_NVSHMEM", 0))): assert ( os.getenv("NVSHMEM_HOME") is not None @@ -121,6 +144,7 @@ def setup_pytorch_extension( sources=[str(src) for src in sources], include_dirs=[str(inc) for inc in include_dirs], extra_compile_args={"cxx": cxx_flags}, + extra_link_args=extra_link_args, libraries=[str(lib) for lib in libraries], library_dirs=[str(lib_dir) for lib_dir in library_dirs], ) diff --git a/pyproject.toml b/pyproject.toml index 2c9f224c14..924ca49a28 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ # See LICENSE for license information. [build-system] -requires = ["setuptools>=61.0", "cmake>=3.21", "wheel", "pybind11[global]", "ninja", "pip", "torch>=2.1", "jax>=0.5.0", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.25.0"] +requires = ["setuptools>=61.0", "cmake>=3.21", "wheel", "pybind11[global]", "ninja", "pip", "torch>=2.1", "jax>=0.5.0", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.25.0", "apache-tvm-ffi>=0.1.12"] # Use legacy backend to import local packages in setup.py build-backend = "setuptools.build_meta:__legacy__" diff --git a/setup.py b/setup.py index 150d92969c..a252e18888 100644 --- a/setup.py +++ b/setup.py @@ -122,6 +122,7 @@ def setup_requirements() -> Tuple[List[str], List[str]]: "pydantic", "importlib-metadata>=1.0", "packaging", + "apache-tvm-ffi>=0.1.12", ] test_reqs: List[str] = ["pytest>=8.2.1"] @@ -359,13 +360,17 @@ def git_check_submodules() -> None: "core_cu13": [f"transformer_engine_cu13=={__version__}"], "pytorch": [f"transformer_engine_torch=={__version__}"], "jax": [f"transformer_engine_jax=={__version__}"], + "cutedsl": ["nvidia-cutlass-dsl>=4.2.0"], } else: install_requires, test_requires = setup_requirements() ext_modules = [setup_common_extension()] package_data = {"": ["VERSION.txt"]} include_package_data = True - extras_require = {"test": test_requires} + extras_require = { + "test": test_requires, + "cutedsl": ["nvidia-cutlass-dsl>=4.2.0"], + } if not bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))): if "pytorch" in frameworks: diff --git a/tests/pytorch/mxfp8/bench_mxfp8_cutedsl.py b/tests/pytorch/mxfp8/bench_mxfp8_cutedsl.py new file mode 100644 index 0000000000..2524b4a60f --- /dev/null +++ b/tests/pytorch/mxfp8/bench_mxfp8_cutedsl.py @@ -0,0 +1,517 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Benchmark MXFP8 quantization through the standard TE dispatch path. + +This bench measures ONE backend per process. The kernel that actually runs is +decided entirely by the env var the process is launched with: + + NVTE_ENABLE_CUTEDSL_QUANT_BACKEND=0 -> C++ CUDA kernel (label "cpp") + NVTE_ENABLE_CUTEDSL_QUANT_BACKEND=1 -> CuTeDSL kernel (label "dsl") + +The Python code path is identical either way (`tex.quantize` / `tex.`), so +the two runs are directly comparable — only the backend under the hood differs. +Use run_nsys_profile.sh to run this script twice (disabled vs enabled) and align +the results. + +Timing paths: + --gpu-nsys GPU time from nsys. Warms up (incl. the CuTeDSL JIT + compile), then runs `iters` cold-L2 iterations with each + fn() call wrapped in a same-named QBENCH|tag|MxN|dir NVTX + range, and prints the per-workload byte count. The kernel + time is read from nsys's NVTX Range Kernel Summary + (nvtx_kern_sum) by run_mxfp8_benchmark.py — NOT here. The + whole matrix can run in ONE nsys process (ranges, not + kernel names, separate the workloads). + --no-evict-l2 CPU time, warm cache. Tight launch loop, NO sync and NO + L2 flush — host dispatch cost / warm-cache throughput. + --evict-l2 (default) GPU time, cold cache, via CUDA events (sync per iter). + Kept for standalone use; the driver uses --gpu-nsys for + GPU timing since CUDA events add a small measurement skew. + --single One-shot cold-cache wall-clock. Overrides --iters. + +Produces NVTX-tagged iterations for Nsight Systems timeline profiling. +""" + +import argparse +import csv +import os +import sys +import time + +import torch +import torch.cuda.nvtx as nvtx + +import transformer_engine.pytorch as te # must precede transformer_engine_torch +import transformer_engine_torch as tex +from transformer_engine.pytorch import MXFP8Quantizer + +# Shape presets — names map to lists of (M, N) pairs. +# All shapes are multiples of 64 (the CuTeDSL kernel's CHUNK_DIM). +SHAPE_PRESETS = { + "tiny": [(128, 128), (256, 256), (512, 512)], + "small": [(1024, 1024), (2048, 2048), (4096, 4096)], + "medium": [(8192, 8192), (8192, 4096), (4096, 8192)], + "large": [(16384, 8192), (16384, 16384), (32768, 8192)], + "square": [(1024, 1024), (2048, 2048), (4096, 4096), + (8192, 8192), (16384, 16384)], + # LLM-typical shapes: (batch*seq, hidden) for common hidden sizes + "llm": [(2048, 5120), (2048, 8192), (4096, 12288), + (8192, 14336), (16384, 16384)], + # Aspect-ratio sweep: tall-narrow and short-wide + "aspect": [(1024, 16384), (4096, 4096), (16384, 1024), + (512, 32768), (32768, 512)], + "default": [(1024, 1024), (4096, 4096), (8192, 8192), + (16384, 8192), (16384, 16384)], + "test": [(2048, 5120)], +} + + +def backend_label(): + """Mirror the C++ gate in tvm_ffi_bridge.h: enabled iff the env var is set + and its first char is not '0'.""" + v = os.environ.get("NVTE_ENABLE_CUTEDSL_QUANT_BACKEND", "") + return "dsl" if (len(v) > 0 and v[0] != "0") else "cpp" + + +def parse_shapes(shapes_str: str): + """Parse a ';'-separated list of 'M,N' pairs.""" + shapes = [] + for pair in shapes_str.split(";"): + m, n = pair.strip().split(",") + shapes.append((int(m), int(n))) + return shapes + + +# All five elementwise activations CuTeDSL supports (gated geglu/qgeglu excluded +# — they don't route through this MXFP8 quantize path). +_ACTS = ["gelu", "silu", "relu", "qgelu", "srelu"] +_ACT_FNS = {a: getattr(tex, a) for a in _ACTS} +_DACT_FNS = {"d" + a: getattr(tex, "d" + a) for a in _ACTS} +_DBIAS_DACT_FNS = {"dbias_d" + a: getattr(tex, "dbias_d" + a) for a in _ACTS} + +# combo -> kind, which selects the tex entry point and how many inputs are read: +# plain tex.quantize(x) 1 input -> C++ SPECIALIZED cast-only kernel +# act tex.(x) 1 input -> C++ standard kernel (IS_ACT) +# dact tex.(grad, act_in) 2 inputs -> C++ standard kernel (IS_DACT) +# dbias tex.bgrad_quantize(grad) 1 input -> C++ standard kernel (IS_DBIAS) +# dbias_dact tex.dbias_(grad, act_in) 2 inputs -> standard (IS_DBIAS|IS_DACT) +# Only "plain" hits the specialized kernel (specialized::hasSpec is true solely for +# IS_DBIAS=IS_DACT=IS_ACT=false). Every other combo forces the standard MXFP8 kernel +# that the CuTeDSL backend mirrors — the apples-to-apples comparison. +COMBO_KIND = {"plain": "plain", "dbias": "dbias"} +for _a in _ACTS: + COMBO_KIND[_a] = "act" # fwd activation, no dbias + COMBO_KIND["d" + _a] = "dact" # dactivation, no dbias + COMBO_KIND["dbias_d" + _a] = "dbias_dact" # dactivation + dbias + +_FP8_DTYPES = { + "e4m3": tex.DType.kFloat8E4M3, + "e5m2": tex.DType.kFloat8E5M2, +} +_TORCH_IN_DTYPES = { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, +} + + +def _make_quantizer(rowwise, colwise, fp8_dtype, swizzle): + q = MXFP8Quantizer( + fp8_dtype=_FP8_DTYPES[fp8_dtype], + rowwise=rowwise, + columnwise=colwise, + ) + q.internal = True + if swizzle: + q.optimize_for_gemm = True + return q + + +def make_fn(combo, x, act_input, rowwise, colwise, fp8_dtype="e4m3", swizzle=False): + """Return a 0-arg callable that quantizes `x` via the standard TE dispatch. + + Whether this lands on the C++ or CuTeDSL kernel is governed by + NVTE_ENABLE_CUTEDSL_QUANT_BACKEND in the environment — not by this code. + `act_input` is the second input (fwd activation input) for dact combos. + + Uses the direct `tex.*` pybinds (not `MXFP8Quantizer.__call__`, which wraps + the result in a Float8Tensor and adds ~15 us of Python overhead that would + skew warm-cache wall-clock). + """ + quantizer = _make_quantizer(rowwise, colwise, fp8_dtype, swizzle) + kind = COMBO_KIND[combo] + if kind == "plain": + return lambda: tex.quantize(x, quantizer) + if kind == "act": + op = _ACT_FNS[combo] + return lambda: op(x, quantizer) + if kind == "dact": + op = _DACT_FNS[combo] + return lambda: op(x, act_input, quantizer) # x is grad + if kind == "dbias": + return lambda: tex.bgrad_quantize(x, quantizer) # x is grad + if kind == "dbias_dact": + op = _DBIAS_DACT_FNS[combo] + return lambda: op(x, act_input, quantizer) # x is grad + raise ValueError(f"unknown combo {combo!r}") + + +# Module-level L2 evict buffer. 256 MB f32 (covers B200's ~60 MB L2 with +# headroom). Allocated lazily, reused to avoid alloc churn between bench runs. +_L2_EVICT_BUF = None + + +def _l2_evict_buf(): + global _L2_EVICT_BUF + if _L2_EVICT_BUF is None: + _L2_EVICT_BUF = torch.empty( + 256 * 1024 * 1024 // 4, dtype=torch.float32, device="cuda") + return _L2_EVICT_BUF + + +def bench_once(name, fn, warmup, iters, evict_l2=False, single=False): + """Time `fn()` under one of three modes (see module docstring). + + Returns the average per-iter time in milliseconds. `single` takes precedence + over `evict_l2`. + """ + if single: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + + evict = _l2_evict_buf() + nvtx.range_push(f"{name}_single") + evict.zero_() # async L2 flush + torch.cuda.synchronize() # drain evict; L2 cold AND GPU idle + t0 = time.perf_counter_ns() + fn() + torch.cuda.synchronize() + t1 = time.perf_counter_ns() + nvtx.range_pop() + return (t1 - t0) / 1e6 + + if evict_l2: + # GPU time, cold cache: flush L2, sync, then time the kernel with CUDA + # events (the event pair measures pure on-GPU kernel time per iter). + for _ in range(warmup): + fn() + torch.cuda.synchronize() + + evict = _l2_evict_buf() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + total_ms = 0.0 + nvtx.range_push(f"{name}_measure") + for i in range(iters): + evict.zero_() # async L2 flush + torch.cuda.synchronize() # drain evict; L2 cold, GPU idle + nvtx.range_push(f"{name}_iter_{i}") + start.record() + fn() # kernel launch + end.record() + torch.cuda.synchronize() # wait for kernel + nvtx.range_pop() + total_ms += start.elapsed_time(end) # GPU time (ms) + nvtx.range_pop() + return total_ms / iters + + # Warm cache, CPU time: tight launch loop, NO sync, NO flush. Kernels queue + # asynchronously, so this is host dispatch cost / warm-cache throughput. + for _ in range(warmup): + fn() + torch.cuda.synchronize() + + nvtx.range_push(f"{name}_warm") + t0 = time.perf_counter_ns() + for _ in range(iters): + fn() + t1 = time.perf_counter_ns() + nvtx.range_pop() + return (t1 - t0) / 1e6 / iters + + +def _workload_bytes(M, N, in_bytes_per_elt, need_act_input, rowwise, colwise): + """Bytes the quantize moves through HBM: input read + FP8 out + e8m0 scales. + + A lower bound on DRAM traffic (assumes each byte touched once; ignores the + dbias workspace round-trip). Used to turn a measured kernel time into GB/s. + """ + bytes_in = M * N * in_bytes_per_elt * (2 if need_act_input else 1) + bytes_out = bytes_scale = 0 + if rowwise: + bytes_out += M * N # rowwise FP8 data (uint8) + bytes_scale += M * (N // 32) # rowwise e8m0 scales + if colwise: + bytes_out += M * N # colwise FP8 data + bytes_scale += (M // 32) * N # colwise e8m0 scales + return bytes_in + bytes_out + bytes_scale + + +def bench_shape(M, N, rowwise, colwise, warmup, iters, combo="plain", + in_dtype="bf16", fp8_dtype="e4m3", swizzle=False, + evict_l2=False, single=False, gpu_nsys=False): + torch.manual_seed(0) + torch.cuda.manual_seed(0) + in_dt = _TORCH_IN_DTYPES[in_dtype] + x = torch.randn(M, N, dtype=in_dt, device="cuda") + backend = backend_label() + # dact / dbias_dact read a second input (the fwd activation input). + need_act_input = COMBO_KIND[combo] in ("dact", "dbias_dact") + act_input = (torch.randn(M, N, dtype=in_dt, device="cuda") + if need_act_input else None) + + dir_label = "both" if (rowwise and colwise) else ("row" if rowwise else "col") + tag = f"{combo}_{in_dtype}_{fp8_dtype}" + if swizzle: + tag += "_sw" + + nvtx.range_push(f"shape_{M}x{N}_{dir_label}_{backend}_{tag}") + + fn = make_fn(combo, x, act_input, rowwise, colwise, + fp8_dtype=fp8_dtype, swizzle=swizzle) + total_bytes = _workload_bytes(M, N, x.element_size(), need_act_input, + rowwise, colwise) + + if gpu_nsys: + # The driver profiles the WHOLE matrix in one nsys run per backend and + # attributes per-workload kernel time via NVTX ranges (nsys nvtx_kern_sum): + # the kernel NAME encodes the config but not the shape, so we tag each + # workload with a same-named QBENCH range (which carries M/N) that wraps + # ONLY fn(). nvtx_kern_sum reports the real CUPTI kernel durations bucketed + # per range; warmup + the L2-evict sit OUTSIDE the range (blank range) and + # are ignored. Emit the byte count so the driver can compute GB/s. + print(f"NSYS_BYTES backend={backend} tag={tag} combo={combo} " + f"M={M} N={N} dir={dir_label} bytes={total_bytes} iters={iters}", + flush=True) + for _ in range(max(warmup, 1)): # warmup OUTSIDE any range + fn() + torch.cuda.synchronize() + + evict = _l2_evict_buf() + rng = f"QBENCH|{tag}|{M}x{N}|{dir_label}" + for i in range(iters): + evict.zero_() # cold L2, OUTSIDE the range + torch.cuda.synchronize() + nvtx.range_push(rng) # same name every iter -> + fn() # nvtx_kern_sum aggregates + torch.cuda.synchronize() + nvtx.range_pop() + ms = None + else: + # Warm caches / trigger the CuTeDSL JIT compile once (not counted). + nvtx.range_push("warm") + fn() + torch.cuda.synchronize() + nvtx.range_pop() + ms = bench_once( + f"{backend}_{M}x{N}_{dir_label}_{tag}", + fn, warmup, iters, evict_l2=evict_l2, single=single, + ) + + nvtx.range_pop() # close shape_ range + + gbps = (total_bytes / (ms * 1e-3) / 1e9) if ms is not None else None + + return { + "backend": backend, + "shape": (M, N), + "dir": dir_label, + "combo": combo, + "tag": tag, + "ms": ms, + "gbps": gbps, + "bytes": total_bytes, + } + + +def main(): + parser = argparse.ArgumentParser( + formatter_class=argparse.RawDescriptionHelpFormatter, + description=__doc__, + epilog=( + "Shape selection (pick one of --preset / --shapes):\n" + " --preset PRESET named sweep (see --list-presets)\n" + " --shapes 'M,N;...' custom list\n" + ), + ) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iters", type=int, default=100) + parser.add_argument("--direction", choices=["row", "col", "both", "all"], + default="all", + help="Which direction(s) to benchmark") + parser.add_argument("--directions", type=str, default=None, + help="Comma-separated subset of row,col,both " + "(overrides --direction; lets one process cover all).") + parser.add_argument("--combo", type=str, default="plain", + choices=sorted(COMBO_KIND), + help=f"Operation: one of {sorted(COMBO_KIND)}") + parser.add_argument("--combos", type=str, default=None, + help="Comma-separated list of combos (overrides --combo)") + parser.add_argument("--in-dtype", type=str, default="bf16", + choices=sorted(_TORCH_IN_DTYPES)) + parser.add_argument("--in-dtypes", type=str, default=None, + help="Comma-separated list of input dtypes") + parser.add_argument("--fp8", type=str, default="e4m3", + choices=sorted(_FP8_DTYPES)) + parser.add_argument("--fp8s", type=str, default=None, + help="Comma-separated list of fp8 output dtypes") + parser.add_argument("--swizzle", action="store_true", + help="Enable GEMM-swizzled scales (optimize_for_gemm)") + parser.add_argument("--swizzles", type=str, default=None, + help="Comma-separated subset of off,on (overrides " + "--swizzle; lets one process cover both).") + parser.add_argument("--evict-l2", dest="evict_l2", action="store_true", + default=True, + help="GPU time, cold cache (default): flush L2 before " + "each iter, time the kernel with CUDA events.") + parser.add_argument("--no-evict-l2", dest="evict_l2", action="store_false", + help="CPU time, warm cache: tight launch loop, no sync, " + "no L2 flush (host dispatch throughput).") + parser.add_argument("--single", action="store_true", + help="One-shot cold-cache wall-clock; overrides --iters " + "and takes precedence over --evict-l2.") + parser.add_argument("--gpu-nsys", dest="gpu_nsys", action="store_true", + help="GPU time via nsys: run cold-L2 iters with each fn() " + "wrapped in a same-named QBENCH NVTX range and print " + "the byte count (no in-process timing). The driver " + "reads per-range kernel time from nvtx_kern_sum. Used " + "by run_mxfp8_benchmark.py.") + parser.add_argument("--preset", type=str, default=None, + choices=sorted(SHAPE_PRESETS), + help=f"Shape preset: one of {sorted(SHAPE_PRESETS)}") + parser.add_argument("--shapes", type=str, default=None, + help="Custom shapes: 'M,N;M,N;...' Overrides --preset") + parser.add_argument("--list-presets", action="store_true", + help="Print all presets and exit") + parser.add_argument("--csv", type=str, default=None, + help="Write results as CSV to this file") + args = parser.parse_args() + + if args.list_presets: + for name, shapes in SHAPE_PRESETS.items(): + shapes_str = ", ".join(f"{m}x{n}" for m, n in shapes) + print(f" {name:8s} {shapes_str}") + return 0 + + if args.shapes: + shapes = parse_shapes(args.shapes) + elif args.preset: + shapes = SHAPE_PRESETS[args.preset] + else: + shapes = SHAPE_PRESETS["default"] + + _DIR_RCW = {"row": (True, False), "col": (False, True), "both": (True, True)} + if args.directions: + dir_names = [d.strip() for d in args.directions.split(",") if d.strip()] + elif args.direction == "all": + dir_names = ["row", "col", "both"] + else: + dir_names = [args.direction] + for d in dir_names: + if d not in _DIR_RCW: + print(f"unknown direction: {d}", file=sys.stderr) + return 1 + dirs = [(d, *_DIR_RCW[d]) for d in dir_names] + + # Swizzle list: --swizzles 'off,on' overrides the single --swizzle flag so one + # process can cover both scale layouts. + if args.swizzles: + swizzles = [s.strip().lower() == "on" + for s in args.swizzles.split(",") if s.strip()] + else: + swizzles = [args.swizzle] + + if args.combos: + combos = [c.strip() for c in args.combos.split(",")] + for c in combos: + if c not in COMBO_KIND: + print(f"unknown combo: {c}", file=sys.stderr) + return 1 + else: + combos = [args.combo] + + in_dtypes = ([d.strip() for d in args.in_dtypes.split(",")] + if args.in_dtypes else [args.in_dtype]) + fp8s = ([d.strip() for d in args.fp8s.split(",")] + if args.fp8s else [args.fp8]) + + backend = backend_label() + mode = ("gpu-nsys (kernel time from nsys summary)" if args.gpu_nsys else + "single (one-shot cold wall-clock)" if args.single else + "evict-l2 (GPU time, cold cache)" if args.evict_l2 else + "warm cache (CPU dispatch time, no sync)") + print(f"Backend: {backend} " + f"(NVTE_ENABLE_CUTEDSL_QUANT_BACKEND=" + f"{os.environ.get('NVTE_ENABLE_CUTEDSL_QUANT_BACKEND', '')})") + print(f"Benchmarking {len(shapes)} shape(s) × {len(dirs)} direction(s) × " + f"{len(combos)} combo(s) × {len(in_dtypes)} in-dtype × {len(fp8s)} fp8 " + f"× {len(swizzles)} swizzle") + print(f" mode: {mode}") + print(f" warmup={args.warmup} iters={args.iters}") + print(f" combos: {combos} in_dtypes: {in_dtypes} fp8: {fp8s} " + f"dirs: {dir_names} swizzles: {swizzles}") + for m, n in shapes: + print(f" - {m}x{n}") + print() + + # --gpu-nsys uses NVTX ranges (no cudaProfiler); otherwise wrap the whole sweep + # so the nsys timeline is annotated when run under nsys manually. + if not args.gpu_nsys: + torch.cuda.profiler.start() # used with --capture-range=cudaProfilerApi + + results = [] + for combo in combos: + for in_dtype in in_dtypes: + for fp8 in fp8s: + for M, N in shapes: + for _, rw, cw in dirs: + for sw in swizzles: + r = bench_shape(M, N, rw, cw, args.warmup, args.iters, + combo, in_dtype=in_dtype, fp8_dtype=fp8, + swizzle=sw, evict_l2=args.evict_l2, + single=args.single, + gpu_nsys=args.gpu_nsys) + results.append(r) + + if not args.gpu_nsys: + torch.cuda.profiler.stop() + + if args.gpu_nsys: + # Kernel time is measured by nsys (the driver parses its summary) and the + # NSYS_BYTES line is emitted by bench_shape up front (before the capture + # range, since stop-shutdown may end the process). Nothing to print here. + return 0 + + # Summary. Columns are positional so run_nsys_profile.sh can parse them: + # backend tag shape dir us GB/s + print() + print(f" ({mode})") + print(f"{'backend':>7} {'tag':>30} {'shape':>12} {'dir':>4} " + f"{'us':>9} {'GB/s':>9}") + print("-" * 84) + for r in results: + M, N = r["shape"] + mxn = f"{M}x{N}" + print(f"{r['backend']:>7} {r['tag']:>30} {mxn:>12} {r['dir']:>4} " + f"{r['ms']*1000:9.2f} {r['gbps']:9.1f}") + + if args.csv: + with open(args.csv, "w", newline="") as f: + w = csv.writer(f) + w.writerow(["backend", "tag", "combo", "M", "N", "dir", + "us", "gbps", "bytes"]) + for r in results: + M, N = r["shape"] + w.writerow([r["backend"], r["tag"], r["combo"], M, N, r["dir"], + f"{r['ms']*1000:.3f}", f"{r['gbps']:.2f}", + r["bytes"]]) + print(f"\nCSV written to {args.csv}") + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/pytorch/mxfp8/run_mxfp8_benchmark.py b/tests/pytorch/mxfp8/run_mxfp8_benchmark.py new file mode 100644 index 0000000000..5e961d490e --- /dev/null +++ b/tests/pytorch/mxfp8/run_mxfp8_benchmark.py @@ -0,0 +1,477 @@ +#!/usr/bin/env python3 +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""One-shot driver for the MXFP8 quantize benchmark. + +`bench_mxfp8_cutedsl.py` measures ONE backend in ONE timing mode per process +(the backend is latched by NVTE_ENABLE_CUTEDSL_QUANT_BACKEND at import). This +driver runs every {backend} x {timing-mode} combination in its own subprocess +and prints a single merged table, so one command gives you the full picture: + + backend: cpp = CUDA C++ kernels (NVTE_ENABLE_CUTEDSL_QUANT_BACKEND=0) + dsl = CuTeDSL kernels (NVTE_ENABLE_CUTEDSL_QUANT_BACKEND=1) + mode: GPU = kernel time, cold L2 (nsys NVTX Range Kernel Summary; the + whole matrix in ONE nsys run per + backend, per-shape via NVTX, needs nsys) + CPU = host dispatch time (in-process wall clock, --no-evict-l2) + +Default (curated): combos {plain, dbias, gelu, dgelu} x directions {row, col, +both} (both = bidirectional, rowwise+colwise in one pass) x swizzle {off, on} x +bf16 x e4m3 x an LLM-representative shape set (hidden 4096-14336, a few thousand +tokens), for both backends and both modes. Override any axis (--preset/--shapes +for sizes), or use --all for the full matrix. + +Usage: + python run_mxfp8_benchmark.py # curated default + python run_mxfp8_benchmark.py --preset llm --modes gpu + python run_mxfp8_benchmark.py --combos plain --directions row --swizzle off + python run_mxfp8_benchmark.py --backends dsl # CuTeDSL only + python run_mxfp8_benchmark.py --all --preset tiny # everything +""" + +import argparse +import csv +import os +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path + +BENCH = Path(__file__).with_name("bench_mxfp8_cutedsl.py") +# Repo root (…/tests/pytorch/mxfp8/run_mxfp8_benchmark.py -> 3 levels up). The bench +# is launched as a script, so its own dir lands on sys.path and `import +# transformer_engine` would resolve to the INSTALLED package — not this checkout — +# meaning the CuTeDSL backend silently runs the installed (often stale / CUDA- +# fallback) kernels. Putting the repo root on PYTHONPATH forces the local TE so the +# `dsl` backend actually exercises the local CuTeDSL kernels. +_REPO_ROOT = Path(__file__).resolve().parents[3] +# CPU mode times host dispatch in-process; GPU mode reads kernel time from nsys. +_CPU_FLAG = "--no-evict-l2" + +# Kernel-summary rows whose name matches any of these are the L2-evict op (a +# torch fill / memset), not the quantize — excluded from the measured GPU time. +_EVICT_NAME_PATTERNS = ("memset", "memcpy", "fill", "elementwise_kernel", + "vectorized_elementwise") + +# Full axes (mirror bench_mxfp8_cutedsl.py) — used to expand "all" / --all. +_ACTS = ["gelu", "silu", "relu", "qgelu", "srelu"] +_ALL_COMBOS = (["plain", "dbias"] + _ACTS + + ["d" + a for a in _ACTS] + ["dbias_d" + a for a in _ACTS]) +_ALL_DTYPES = ["bf16", "fp16", "fp32"] +_ALL_FP8 = ["e4m3", "e5m2"] + +# Default shapes (tokens M x hidden N): LLM-representative — hidden dims 4096- +# 14336 (7B/70B hidden + Llama-3 MLP intermediate), a few thousand tokens. All +# multiples of 128 so the swizzled-scale layout applies. Override with --preset +# / --shapes. +_DEFAULT_SHAPES = "4096,4096;4096,8192;8192,8192;4096,14336" + + +def _expand(val, full): + """Expand a comma list, turning the literal 'all' into the full axis.""" + if val is None: + return None + items = [v.strip() for v in val.split(",") if v.strip()] + return ",".join(full) if items == ["all"] else val + + +def _detect_cute_dsl_arch(): + """sm_[a] for the current device (CuTeDSL compile target).""" + try: + import torch + + major, minor = torch.cuda.get_device_capability() + return f"sm_{major}{minor}{'a' if major >= 9 else ''}" + except Exception: + return None + + +def _backend_env(backend): + """Process env that latches the backend (and CuTeDSL arch) for a bench run.""" + env = dict(os.environ) + env["NVTE_ENABLE_CUTEDSL_QUANT_BACKEND"] = "1" if backend == "dsl" else "0" + if backend == "dsl": + if "CUTE_DSL_ARCH" not in env: + arch = _detect_cute_dsl_arch() + if arch: + env["CUTE_DSL_ARCH"] = arch + # We EXPECT the CuTeDSL backend to handle every quantize in a `dsl` run. + # If the dispatcher falls back to CUDA (e.g. the wrong transformer_engine + # got imported, or the kernel didn't register/compile), this makes it warn + # loudly instead of silently producing CUDA numbers labelled "dsl". + env.setdefault("NVTE_WARN_IF_CUTEDSL_BACKEND_NOT_CHOSEN", "1") + # Force the bench subprocess to import THIS checkout's transformer_engine (so + # `dsl` runs the local CuTeDSL kernels), not the installed package. Without + # this, `python bench.py` imports installed TE and `dsl` silently falls back. + existing = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = str(_REPO_ROOT) + (os.pathsep + existing if existing else "") + return env + + +def _run_cpu(backend, passthrough): + """CPU mode: bench loops shapes/combos in-process, times host dispatch with a + wall clock, writes CSV. Returns {(tag, M, N, dir): (us, gbps)}.""" + env = _backend_env(backend) + with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as f: + csv_path = f.name + cmd = [sys.executable, str(BENCH), _CPU_FLAG, "--csv", csv_path] + passthrough + print(f"[run] backend={backend:3s} mode=cpu: {' '.join(cmd)}", file=sys.stderr) + try: + proc = subprocess.run(env=env, args=cmd, stdout=subprocess.DEVNULL) + if proc.returncode != 0: + print(f"[warn] backend={backend} mode=cpu exited {proc.returncode}; " + "skipping (is this backend available?)", file=sys.stderr) + return None + rows = {} + with open(csv_path) as fh: + for r in csv.DictReader(fh): + rows[(r["tag"], int(r["M"]), int(r["N"]), r["dir"])] = ( + float(r["us"]), float(r["gbps"])) + return rows + finally: + if os.path.exists(csv_path): + os.remove(csv_path) + + +def _parse_nsys_bytes_all(stdout): + """All NSYS_BYTES lines -> {(tag, M, N, dir): bytes}. The bench emits one per + workload before its timed loop.""" + out = {} + for line in stdout.splitlines(): + if line.startswith("NSYS_BYTES"): + kv = dict(tok.split("=", 1) for tok in line.split()[1:] if "=" in tok) + try: + out[(kv["tag"], int(kv["M"]), int(kv["N"]), kv["dir"])] = int(kv["bytes"]) + except (KeyError, ValueError): + continue + return out + + +def _parse_nvtx_kern_sum(stats_csv, iters): + """Per-workload per-iter kernel time (ns) from `nsys stats --report + nvtx_kern_sum --format csv`. Each row is (NVTX range, kernel) with the real + CUPTI 'Total Time (ns)'. We sum the kernel Total over each QBENCH range and + divide by its range-instance count (== iters) -> per-iter kernel time, the + same precision as cuda_gpu_kern_sum but bucketed by workload. The L2-evict + sits in the blank range (outside QBENCH) and is skipped. Returns + {(tag, M, N, dir): per_iter_ns}.""" + lines = stats_csv.splitlines() + hdr = next((i for i, ln in enumerate(lines) + if "NVTX Range" in ln and "Total Time" in ln), None) + if hdr is None: + return {} + reader = csv.DictReader(lines[hdr:]) + fields = reader.fieldnames or [] + rng_col = next((c for c in fields if c.strip() == "NVTX Range"), None) + tot_col = next((c for c in fields if "Total Time" in c), None) + inst_col = next((c for c in fields if c.strip() == "NVTX Inst"), None) + name_col = next((c for c in fields if "Kernel Name" in c), None) + if not (rng_col and tot_col and inst_col): + return {} + # range key -> [summed kernel Total ns, range instances] + acc = {} + for row in reader: + name = (row.get(rng_col) or "").lstrip(":").strip() + if not name.startswith("QBENCH|"): + continue + kname = (row.get(name_col) or "").lower() if name_col else "" + if any(p in kname for p in _EVICT_NAME_PATTERNS): + continue # defensive; evict is in the blank range anyway + parts = name.split("|") # QBENCH | tag | MxN | dir + if len(parts) != 4: + continue + _, tag, mxn, d = parts + try: + M, N = (int(v) for v in mxn.lower().split("x")) + tot = float((row[tot_col] or "0").replace(",", "")) + inst = int(float((row[inst_col] or "0").replace(",", ""))) + except (TypeError, ValueError): + continue + key = (tag, M, N, d) + cur = acc.setdefault(key, [0.0, inst]) + cur[0] += tot + cur[1] = inst or cur[1] + return {k: (tot / (inst or iters)) for k, (tot, inst) in acc.items() if tot > 0} + + +def _run_gpu_nsys_backend(nsys, env, backend, combos, in_dtypes, fp8s, shapes_str, + directions, swizzles, warmup, iters, out_dir=None): + """Profile the WHOLE matrix for one backend in a SINGLE nsys run; attribute + per-workload kernel time via NVTX ranges (nvtx_kern_sum). Returns + {(tag, M, N, dir): (us, gbps)}. + + One process (not one per workload) because the kernel NAME can't separate + shapes — the QBENCH NVTX range carries M/N. Whole-process profile flushes + reliably at exit (no capture range needed). If out_dir is set, the raw report, + the nvtx_kern_sum CSV and the run log are kept there. + + Note CuTeDSL JIT-compiles each distinct config on first use; here that all + happens once, in this one process, during each workload's warmup.""" + tmp = None if out_dir else tempfile.TemporaryDirectory() + rep = (os.path.join(out_dir, f"{backend}_matrix") if out_dir + else os.path.join(tmp.name, "rep")) + try: + sw_arg = ",".join("on" if s else "off" for s in swizzles) + bench_cmd = [sys.executable, str(BENCH), "--gpu-nsys", + "--combos", combos, "--in-dtypes", in_dtypes, "--fp8s", fp8s, + "--directions", ",".join(directions), "--swizzles", sw_arg, + "--shapes", shapes_str, "--warmup", str(warmup), + "--iters", str(iters)] + cmd = [nsys, "profile", "-o", rep, "-f", "true"] + bench_cmd + print(f"[run] backend={backend} nsys (full matrix in 1 process): " + f"combos={combos} dirs={','.join(directions)} sw={sw_arg} " + f"shapes={shapes_str}", file=sys.stderr) + proc = subprocess.run(cmd, env=env, capture_output=True, text=True) + if out_dir: + with open(rep + ".log", "w") as fh: + fh.write(f"$ {' '.join(cmd)}\n[exit {proc.returncode}]\n" + f"--- stdout ---\n{proc.stdout}\n--- stderr ---\n{proc.stderr}") + if proc.returncode != 0: + print(f"[warn] nsys profile failed ({proc.returncode}) for " + f"backend={backend}; skipping.\n{proc.stderr[-800:]}", + file=sys.stderr) + return {} + bytes_map = _parse_nsys_bytes_all(proc.stdout) + # --force-export=true: always re-derive the SQLite from the freshly + # captured .nsys-rep. Otherwise `nsys stats` reuses a stale .sqlite left + # next to a reused output path (e.g. --nsys-out), silently reporting old data. + stats = subprocess.run( + [nsys, "stats", "--force-export=true", "--report", "nvtx_kern_sum", + "--format", "csv", rep + ".nsys-rep"], capture_output=True, text=True) + if out_dir: + with open(rep + ".nvtx_kern_sum.csv", "w") as fh: + fh.write(stats.stdout) + if stats.returncode != 0: + print(f"[warn] nsys stats failed for backend={backend}; skipping." + f"\n{stats.stderr[-500:]}", file=sys.stderr) + return {} + per_iter = _parse_nvtx_kern_sum(stats.stdout, iters) + rows = {} + for key, per_iter_ns in per_iter.items(): + b = bytes_map.get(key) + if b is None or per_iter_ns <= 0: + continue + rows[key] = (per_iter_ns / 1e3, b / per_iter_ns) # us, GB/s + if not rows: + print(f"[warn] no QBENCH ranges parsed from nsys for backend={backend}.", + file=sys.stderr) + elif out_dir: + print(f"[nsys] saved {rep}.nsys-rep ({len(rows)} workloads)", + file=sys.stderr) + return rows + finally: + if tmp is not None: + # Temp path: the whole dir (report + SQLite) is removed. + tmp.cleanup() + else: + # --nsys-out path: keep the raw .nsys-rep, but delete the generated + # SQLite so a later run can never read stale data from it. + sqlite = rep + ".sqlite" + if os.path.exists(sqlite): + os.remove(sqlite) + + +def _parse_shapes(shapes_str): + """'M,N;M,N;...' -> [(M, N), ...].""" + out = [] + for tok in shapes_str.split(";"): + tok = tok.strip() + if tok: + m, n = tok.split(",") + out.append((int(m), int(n))) + return out + + +def _preset_shapes(preset): + """Resolve a bench shape preset to [(M, N), ...] by asking the bench.""" + out = subprocess.run([sys.executable, str(BENCH), "--list-presets"], + capture_output=True, text=True) + for line in out.stdout.splitlines(): + parts = line.split() + if parts and parts[0] == preset: + shapes = [] + for tok in " ".join(parts[1:]).split(","): + tok = tok.strip().lower() + if "x" in tok: + m, n = tok.split("x") + shapes.append((int(m), int(n))) + return shapes + return [] + + +def _fwd(args_ns, passthrough_keys): + """Rebuild the forwarded bench CLI flags from parsed args.""" + out = [] + for flag, val in passthrough_keys.items(): + if val is None or val is False: + continue + if val is True: + out.append(flag) + else: + out += [flag, str(val)] + return out + + +def main(): + ap = argparse.ArgumentParser( + formatter_class=argparse.RawDescriptionHelpFormatter, description=__doc__) + ap.add_argument("--backends", default="cpp,dsl", + help="Comma-separated: cpp (CUDA), dsl (CuTeDSL). Default both.") + ap.add_argument("--modes", default="gpu,cpu", + help="Comma-separated: gpu (kernel time), cpu (dispatch time). Default both.") + ap.add_argument("--all", action="store_true", + help="Override the curated defaults with EVERY case: all 17 " + "combos x row/col/both x all 3 input dtypes x both fp8 " + "formats x swizzle on+off. Very heavy — pair with a small " + "--preset and modest --iters.") + # Curated defaults: plain + one act + one dact (+ plain dbias), rowwise, + # columnwise, and bidirectional (both), swizzle on+off. Override any axis + # explicitly; 'all' expands an axis. + ap.add_argument("--combos", default="plain,dbias,gelu,dgelu") + ap.add_argument("--directions", default="row,col,both", + help="Comma-separated subset of row,col,both " + "(both = bidirectional, rowwise+colwise in one pass).") + ap.add_argument("--swizzle", choices=["off", "on", "both"], default="both", + help="Swizzled scale layout: off / on / both. Default both.") + # Shapes: default to an LLM-representative set; --preset / --shapes override. + ap.add_argument("--preset", default=None) + # Forwarded to bench_mxfp8_cutedsl.py (see its --help for semantics). + ap.add_argument("--shapes") + ap.add_argument("--in-dtypes", dest="in_dtypes") + ap.add_argument("--fp8s") + ap.add_argument("--warmup", type=int) + ap.add_argument("--iters", type=int) + ap.add_argument("--nsys-out", dest="nsys_out", default=None, + help="Keep raw nsys artifacts in this dir (GPU mode): per " + "workload a