diff --git a/benchmark/cutedsl_fusion_kernels/README.md b/benchmark/cutedsl_fusion_kernels/README.md new file mode 100644 index 000000000..ff64aaae5 --- /dev/null +++ b/benchmark/cutedsl_fusion_kernels/README.md @@ -0,0 +1,94 @@ +# cuteDSL Fusion Benchmarks + +Microbenchmark comparing **fused cuDNN Frontend** grouped-GEMM + activation +kernels against the equivalent **unfused Transformer Engine** path +(grouped GEMM + activation) for Mixture-of-Experts activations. Timed with CUDA +graph replay; reports the per-activation speedup (`te/cudnn`). + +Activations covered: `swiglu`, `dswiglu`, `srelu`, `dsrelu`, `geglu`, `dgeglu` +(the `d*` variants are the corresponding backward / GLU-bprop patterns). + +For GLU-family kernels, the functional test also accounts for cuDNN Frontend's +32-column interleaved GLU layout. Forward GLU weights are packed from logical +`[gate, up]` order into `[gate32, up32, ...]` order before the fused path. For +backward GLU kernels, the saved activation input is packed the same way and the +fused gradient output is unpacked before comparison. + +For `geglu` and `dgeglu`, cuDNN Frontend's cuteDSL implementation uses +Megatron-style QuickGeGLU: `quick_gelu(gate) * (up + linear_offset)`, with +`quick_gelu(gate) = gate * sigmoid(1.702 * gate)`. The benchmark also applies +the cuDNN FE clamp settings used by this kernel path: gate is clamped above and +the linear half is clamped symmetrically. TE's plain `tex.geglu`/`tex.dgeglu` +use tanh-GELU instead, so the functional oracle uses the cuDNN FE QuickGeGLU +formula for those two tests. This mirrors Megatron-LM's +`megatron/core/fusions/fused_bias_geglu.py` QuickGeGLU helper. By default this +oracle is an eager PyTorch expression; pass +`--functional-reference-backend torch_compile` to run the same expression through +`torch.compile`. + +For `dsrelu`, cuDNN Frontend's cuteDSL backward epilogue treats the grouped GEMM +output as the saved SReLU input and `c_tensor` as the upstream gradient. The +functional oracle follows that operand order. + +## How to run + +In the latest PyTorch container (e.g. `nvcr.io/nvidia/pytorch:26.05-py3`), with +`cutedsl_fusion_benchmarks.py` available on the host: + +```bash +docker run -it --gpus=all --workdir /workspace \ + -v "$(pwd):/host:ro" \ + nvcr.io/nvidia/pytorch:26.05-py3 \ + bash -c "cp /host/cutedsl_fusion_benchmarks.py /workspace && bash" +``` + +Then, inside the container: + +```bash +PIP_CONSTRAINT="" pip install nvidia-cudnn-frontend[cutedsl] transformer-engine + +python cutedsl_fusion_benchmarks.py +``` + +## Running the functional test + +Run only the functional (correctness) subset: + +```bash +python cutedsl_fusion_benchmarks.py --activation all --functional-only +``` + +## Running custom sizes + +```bash +python cutedsl_fusion_benchmarks.py --activation all --experts 8 --tokens 4096 --glu-bprop-tokens 2048 --k 8192 --n 4096 +``` + +## Profiling with nsys + +```bash +nsys profile --force-overwrite=true --trace=cuda,nvtx,osrt --cuda-graph-trace=node --sample=none --cpuctxsw=none \ + python cutedsl_fusion_benchmarks.py --activation all --experts 8 --tokens 4096 --glu-bprop-tokens 2048 --k 8192 --n 4096 +``` + +## Results + +Pass `--output-dir` to write `output.txt` (device/versions/summary), +`output_verbose.txt` (per-activation detail), and `results.png` (speedup plot): + +```bash +python cutedsl_fusion_benchmarks.py --output-dir results +``` + +The results below were generated with the default shape +(`experts=8, tokens/expert=4096, K=8192, N=4096`, `bf16`) using +`torch 2.12.0a0+nv26.05`, `transformer-engine 2.15.0`, +`nvidia-cudnn-frontend 1.23.0`, and `nvidia-cutlass-dsl 4.4.1`. + +### B200 + +![cuDNN fused-epilogue speedups vs Transformer Engine on NVIDIA GB200](results/B200/results.png) + +### B300 + +![cuDNN fused-epilogue speedups vs Transformer Engine on NVIDIA GB300](results/B300/results.png) diff --git a/benchmark/cutedsl_fusion_kernels/cutedsl_fusion_benchmarks.py b/benchmark/cutedsl_fusion_kernels/cutedsl_fusion_benchmarks.py new file mode 100644 index 000000000..9e2c9c4c2 --- /dev/null +++ b/benchmark/cutedsl_fusion_kernels/cutedsl_fusion_benchmarks.py @@ -0,0 +1,1613 @@ +#!/usr/bin/env python3 +"""Benchmark fused cuDNN Frontend activations vs unfused Transformer Engine kernels. + +Supported activations: + swiglu, dswiglu, srelu, dsrelu, geglu, dgeglu + +Fused path: + cuDNN Frontend grouped GEMM + activation wrappers + +Unfused TE path: + TE MXFP8 quantize -> TE grouped GEMM -> optional pointwise scale + -> TE activation + MXFP8 quantize + +Both paths are warmed up, captured with CUDA Graphs, and timed with graph replay. +Functional tests can also be run with --functional-only. They use shared +TE-quantized MXFP8 inputs for the fused and reference paths and compare +dequantized outputs. For GeGLU/DGeGLU, the functional oracle uses an eager +PyTorch reference by default and matches cuDNN Frontend's Megatron-style +QuickGeGLU, clamp, and linear-offset semantics; TE's tex.geglu uses tanh-GELU +and is intentionally different. For dSReLU, the oracle matches cuDNN Frontend's operand order where +the grouped GEMM output is the saved SReLU input and c_tensor is the upstream +gradient. + +Shape convention: + E = number of experts + T = tokens per expert + K = input / gradient hidden size + N = activation output hidden size + +Forward GLU activations use a grouped GEMM with output width 2N and activation +output width N. Backward GLU activations use grouped GEMM output width N and +activation-gradient output width 2N. When comparing against a forward GLU shape, +pass --glu-bprop-tokens as half of --tokens for dswiglu/dgeglu. SReLU/dSReLU use +width N. + +Example: + python cutedsl_fusion_benchmarks.py --activation all --experts 8 --tokens 4096 --glu-bprop-tokens 2048 --k 8192 --n 4096 + python cutedsl_fusion_benchmarks.py --activation all --functional-only +""" + +from __future__ import annotations + +import argparse +import importlib.metadata as metadata +import inspect +import io +import os +import statistics +import sys +from collections.abc import Callable + +import torch +import torch.nn.functional as F + +ACTIVATIONS = ("swiglu", "dswiglu", "srelu", "dsrelu", "geglu", "dgeglu") +FORWARD_GLU = {"swiglu", "geglu"} +BACKWARD_GLU = {"dswiglu", "dgeglu"} +FORWARD_RELU = {"srelu"} +BACKWARD_RELU = {"dsrelu"} + +MXFP8_BLOCK = 32 +M_ALIGNED = 256 +GEGLU_ALPHA = 1.702 +GEGLU_LINEAR_OFFSET = 1.0 +GEGLU_CLAMP_MAX = 7.0 +GEGLU_CLAMP_MIN = -7.0 + +_COMPILED_GEGLU_REFERENCE: Callable[[torch.Tensor], torch.Tensor] | None = None +_COMPILED_DGEGLU_REFERENCE: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor] | None = None + + +def ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b + + +def parse_dtype(name: str) -> torch.dtype: + if name in ("bf16", "bfloat16"): + return torch.bfloat16 + if name in ("fp16", "float16"): + return torch.float16 + raise argparse.ArgumentTypeError("dtype must be bf16 or fp16") + + +def package_version(name: str) -> str: + try: + return metadata.version(name) + except metadata.PackageNotFoundError: + return "not-installed" + + +def check_dims(tokens: int, k: int, n: int) -> None: + dims = {"tokens": tokens, "k": k, "n": n, "2*n": 2 * n} + bad = [name for name, value in dims.items() if value % MXFP8_BLOCK != 0] + if bad: + raise ValueError(f"MXFP8 dimensions must be divisible by {MXFP8_BLOCK}: {bad}") + if tokens % M_ALIGNED: + raise ValueError(f"tokens must be a multiple of {M_ALIGNED} for this benchmark") + if k % 64 or n % 64: + raise ValueError("cuDNN grouped-GEMM fusion expects k and n to be multiples of 64") + + +def activation_shape(activation: str, n: int) -> tuple[int, int | None, int]: + """Return GEMM N, optional fused C input width, and activation output width.""" + if activation in FORWARD_GLU: + return 2 * n, None, n + if activation in BACKWARD_GLU: + return n, 2 * n, 2 * n + if activation in FORWARD_RELU: + return n, None, n + if activation in BACKWARD_RELU: + return n, n, n + raise ValueError(f"Unsupported activation: {activation}") + + +def benchmark_tokens_for_activation(args: argparse.Namespace, activation: str) -> int: + if activation in BACKWARD_GLU and args.glu_bprop_tokens is not None: + return args.glu_bprop_tokens + return args.tokens + + +def random_layout_tensor( + *, + groups: int, + mode0: int, + mode1: int, + mode0_major: bool, + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """Create a tensor in cuDNN Frontend grouped-GEMM layout.""" + shape = (groups, mode1, mode0) if mode0_major else (groups, mode0, mode1) + order = (2, 1, 0) if mode0_major else (1, 2, 0) + src = torch.empty(shape, dtype=torch.float32, device=device).uniform_(-2.0, 2.0) + return src.permute(order).to(dtype) + + +def random_scale_tensor( + *, + groups: int, + mn: int, + k: int, + sf_vec_size: int, + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """Create MXFP8 scale factors in cuDNN Frontend's M32x4xK4 layout.""" + sf_k = ceil_div(k, sf_vec_size) + base_shape = (groups, ceil_div(mn, 128), ceil_div(sf_k, 4), 32, 4, 4) + scale = torch.empty(base_shape, dtype=torch.float32, device=device).uniform_(1.0, 3.0) + return scale.permute(3, 4, 1, 5, 2, 0).to(torch.int8).to(dtype) + + +def view_as_fp8_e4m3(tensor: torch.Tensor) -> torch.Tensor: + return tensor.view(dtype=torch.float8_e4m3fn) + + +def view_as_e8m0(tensor: torch.Tensor) -> torch.Tensor: + return tensor.view(dtype=torch.float8_e8m0fnu) + + +def k_major_3d(tensor: torch.Tensor) -> torch.Tensor: + rows, cols = tensor.shape + return tensor.as_strided((rows, cols, 1), (cols, 1, rows * cols)) + + +def compact_sf_to_cudnn_layout( + scale: torch.Tensor, + *, + mn: int, + nk: int, + groups: int, +) -> torch.Tensor: + """Convert compact MXFP8 scales to cuDNN Frontend's M32x4xK4 layout.""" + sf_k = ceil_div(nk, MXFP8_BLOCK) + m_blocks = ceil_div(mn, 128) + k_blocks = ceil_div(sf_k, 4) + padded = torch.zeros( + (m_blocks * 128, k_blocks * 4, groups), + dtype=scale.dtype, + device=scale.device, + ) + padded[:mn, :sf_k, :] = scale.reshape(mn, sf_k, groups) + base_data = padded.view(m_blocks, 4, 32, k_blocks, 4, groups).permute(5, 0, 3, 2, 1, 4) + + # Keep the same strided tensor contract as cuDNN Frontend test utilities. + base = torch.empty( + (groups, m_blocks, k_blocks, 32, 4, 4), + dtype=scale.dtype, + device=scale.device, + ) + base.copy_(base_data) + return base.permute(3, 4, 1, 5, 2, 0) + + +def cudnn_layout_to_compact_sf( + scale: torch.Tensor, + *, + mn: int, + nk: int, + groups: int, +) -> torch.Tensor: + """Convert cuDNN Frontend's M32x4xK4 scale layout back to compact scales.""" + sf_k = ceil_div(nk, MXFP8_BLOCK) + m_blocks = ceil_div(mn, 128) + k_blocks = ceil_div(sf_k, 4) + base = scale.permute(5, 2, 4, 0, 1, 3) + padded = ( + base.permute(1, 4, 3, 2, 5, 0) + .contiguous() + .view( + m_blocks * 128, + k_blocks * 4, + groups, + ) + ) + return padded[:mn, :sf_k, :] + + +def dequantize_cudnn_mxfp8_output( + tensor: torch.Tensor, + scale: torch.Tensor, + *, + out_n: int, + norm_const: float, +) -> torch.Tensor: + """Dequantize cuDNN FP8 output using the generated row scale factors.""" + m = tensor.shape[0] + compact = cudnn_layout_to_compact_sf(scale, mn=m, nk=out_n, groups=1).to(torch.float32) + expanded = compact.unsqueeze(2).expand(m, compact.shape[1], MXFP8_BLOCK, 1) + expanded = expanded.reshape(m, compact.shape[1] * MXFP8_BLOCK, 1)[:, :out_n, :] + return tensor.to(torch.float32) * expanded / norm_const + + +def glu_interleaved_indices(n: int, *, device: torch.device) -> torch.Tensor: + """Column order expected by cuDNN GLU kernels for a logical [gate, up] tensor.""" + if n % MXFP8_BLOCK: + raise ValueError(f"GLU width must be divisible by {MXFP8_BLOCK}, got {n}") + gate = torch.arange(n, dtype=torch.long, device=device).view(-1, MXFP8_BLOCK) + up = torch.arange(n, 2 * n, dtype=torch.long, device=device).view(-1, MXFP8_BLOCK) + return torch.stack((gate, up), dim=1).reshape(-1) + + +def pack_glu_interleaved(tensor: torch.Tensor) -> torch.Tensor: + """Pack logical [gate, up] columns into [gate32, up32, ...] GLU layout.""" + n2 = tensor.shape[1] + if n2 % 2: + raise ValueError(f"GLU tensor width must be even, got {n2}") + indices = glu_interleaved_indices(n2 // 2, device=tensor.device) + return tensor.index_select(1, indices).contiguous() + + +def unpack_glu_interleaved(tensor: torch.Tensor) -> torch.Tensor: + """Unpack [gate32, up32, ...] GLU columns back to logical [gate, up].""" + n2 = tensor.shape[1] + if n2 % 2: + raise ValueError(f"GLU tensor width must be even, got {n2}") + indices = glu_interleaved_indices(n2 // 2, device=tensor.device) + out = torch.empty_like(tensor) + out.index_copy_(1, indices, tensor) + return out + + +def require_torch_compile() -> None: + if not hasattr(torch, "compile"): + raise RuntimeError("torch.compile is required for --functional-reference-backend torch_compile") + + +def quick_gelu(gate: torch.Tensor) -> torch.Tensor: + """Megatron QuickGELU approximation used by cuDNN FE's GeGLU wrappers.""" + return gate * torch.sigmoid(GEGLU_ALPHA * gate) + + +def quick_geglu_formula(input_tensor: torch.Tensor) -> torch.Tensor: + """Megatron-style QuickGeGLU with cuDNN FE clamp and linear-offset settings.""" + gate, up = input_tensor.to(torch.float32).chunk(2, dim=-1) + gate = gate.clamp(max=GEGLU_CLAMP_MAX) + up = up.clamp(min=GEGLU_CLAMP_MIN, max=GEGLU_CLAMP_MAX) + out = quick_gelu(gate) * (up + GEGLU_LINEAR_OFFSET) + return out.to(input_tensor.dtype) + + +def compiled_geglu_formula() -> Callable[[torch.Tensor], torch.Tensor]: + global _COMPILED_GEGLU_REFERENCE + require_torch_compile() + if _COMPILED_GEGLU_REFERENCE is None: + _COMPILED_GEGLU_REFERENCE = torch.compile(quick_geglu_formula, fullgraph=True, dynamic=False) + return _COMPILED_GEGLU_REFERENCE + + +def cudnn_geglu_reference(input_tensor: torch.Tensor, quantizer, backend: str) -> object: + if backend == "torch_compile": + return quantizer(compiled_geglu_formula()(input_tensor)) + return quantizer(quick_geglu_formula(input_tensor)) + + +def cudnn_dgeglu_formula( + acc_tensor: torch.Tensor, + activation_input: torch.Tensor, + prob_tensor: torch.Tensor, +) -> torch.Tensor: + """Derivative for cuDNN FE's clamped, offset Megatron-style QuickGeGLU.""" + grad = acc_tensor.to(torch.float32) * prob_tensor.to(torch.float32) + gate_raw, up_raw = activation_input.to(torch.float32).chunk(2, dim=-1) + gate = gate_raw.clamp(max=GEGLU_CLAMP_MAX) + up = up_raw.clamp(min=GEGLU_CLAMP_MIN, max=GEGLU_CLAMP_MAX) + sigmoid = torch.sigmoid(GEGLU_ALPHA * gate) + + dgate = grad * sigmoid * (1.0 + GEGLU_ALPHA * gate * (1.0 - sigmoid)) + dgate = dgate * (up + GEGLU_LINEAR_OFFSET) + dup = grad * gate * sigmoid + + dgate = torch.where(gate_raw <= GEGLU_CLAMP_MAX, dgate, torch.zeros_like(dgate)) + up_mask = (up_raw >= GEGLU_CLAMP_MIN) & (up_raw <= GEGLU_CLAMP_MAX) + dup = torch.where(up_mask, dup, torch.zeros_like(dup)) + return torch.cat((dgate, dup), dim=-1).to(activation_input.dtype) + + +def compiled_dgeglu_formula() -> Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]: + global _COMPILED_DGEGLU_REFERENCE + require_torch_compile() + if _COMPILED_DGEGLU_REFERENCE is None: + _COMPILED_DGEGLU_REFERENCE = torch.compile(cudnn_dgeglu_formula, fullgraph=True, dynamic=False) + return _COMPILED_DGEGLU_REFERENCE + + +def cudnn_dgeglu_reference( + acc_tensor: torch.Tensor, + activation_input: torch.Tensor, + prob_tensor: torch.Tensor, + quantizer, + backend: str, +) -> object: + if backend == "torch_compile": + return quantizer(compiled_dgeglu_formula()(acc_tensor, activation_input, prob_tensor)) + return quantizer(cudnn_dgeglu_formula(acc_tensor, activation_input, prob_tensor)) + + +def cudnn_dsrelu_formula(acc_tensor: torch.Tensor, grad_tensor: torch.Tensor) -> torch.Tensor: + """cuDNN FE dSReLU math: 2 * relu(acc) * upstream_grad for alpha=prob=1.""" + acc_relu = torch.relu(acc_tensor.to(torch.float32)) + out = 2.0 * acc_relu * grad_tensor.to(torch.float32) + return out.to(grad_tensor.dtype) + + +def cudnn_dsrelu_reference(acc_tensor: torch.Tensor, grad_tensor: torch.Tensor, quantizer) -> object: + return quantizer(cudnn_dsrelu_formula(acc_tensor, grad_tensor)) + + +def pytorch_activation_formula( + *, + activation: str, + tensor: torch.Tensor, + activation_input: torch.Tensor | None, + prob: torch.Tensor, +) -> torch.Tensor: + """PyTorch pointwise baseline with cuDNN FE semantics for GeGLU/dSReLU cases.""" + if activation == "swiglu": + gate, up = tensor.to(torch.float32).chunk(2, dim=-1) + return (F.silu(gate) * up).to(tensor.dtype) + if activation == "geglu": + return quick_geglu_formula(tensor) + if activation == "srelu": + return torch.relu(tensor.to(torch.float32)).square().to(tensor.dtype) + + if activation_input is None: + raise ValueError(f"{activation} requires activation_input") + + if activation == "dswiglu": + grad = tensor.to(torch.float32) * prob.to(torch.float32) + gate, up = activation_input.to(torch.float32).chunk(2, dim=-1) + sigmoid = torch.sigmoid(gate) + silu = gate * sigmoid + dsilu = sigmoid * (1.0 + gate * (1.0 - sigmoid)) + return torch.cat((grad * up * dsilu, grad * silu), dim=-1).to(activation_input.dtype) + if activation == "dgeglu": + return cudnn_dgeglu_formula(tensor, activation_input, prob) + if activation == "dsrelu": + return cudnn_dsrelu_formula(tensor, activation_input) + + raise ValueError(f"Unsupported activation: {activation}") + + +def make_cudnn_inputs( + *, + activation: str, + experts: int, + tokens: int, + k: int, + n: int, + c_dtype: torch.dtype, + device: torch.device, +) -> dict[str, torch.Tensor | None]: + m = experts * tokens + gemm_n, c_cols, _ = activation_shape(activation, n) + sf_vec_size = 32 + + out: dict[str, torch.Tensor | None] = { + "a_tensor": random_layout_tensor( + groups=1, + mode0=m, + mode1=k, + mode0_major=False, + dtype=torch.float8_e4m3fn, + device=device, + ), + "b_tensor": random_layout_tensor( + groups=experts, + mode0=gemm_n, + mode1=k, + mode0_major=False, + dtype=torch.float8_e4m3fn, + device=device, + ), + "sfa_tensor": random_scale_tensor( + groups=1, + mn=m, + k=k, + sf_vec_size=sf_vec_size, + dtype=torch.float8_e8m0fnu, + device=device, + ), + "sfb_tensor": random_scale_tensor( + groups=experts, + mn=gemm_n, + k=k, + sf_vec_size=sf_vec_size, + dtype=torch.float8_e8m0fnu, + device=device, + ), + "padded_offsets": torch.arange( + tokens, + m + 1, + tokens, + dtype=torch.int32, + device=device, + ), + "alpha_tensor": torch.ones(experts, dtype=torch.float32, device=device), + "beta_tensor": torch.ones(experts, dtype=torch.float32, device=device), + "prob_tensor": torch.ones(m, 1, 1, dtype=torch.float32, device=device), + "dprob_tensor": torch.zeros(m, 1, 1, dtype=torch.float32, device=device), + "norm_const_tensor": torch.tensor([0.01], dtype=torch.float32, device=device), + "c_tensor": None, + } + + if c_cols is not None: + out["c_tensor"] = random_layout_tensor( + groups=1, + mode0=m, + mode1=c_cols, + mode0_major=False, + dtype=c_dtype, + device=device, + ) + return out + + +def make_mxfp8_quantizer(*, rowwise: bool, columnwise: bool): + import transformer_engine.pytorch # noqa: F401 + import transformer_engine_torch as tex + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + return MXFP8Quantizer(tex.DType.kFloat8E4M3, rowwise=rowwise, columnwise=columnwise) + + +def capture_cuda_graph(fn: Callable[[], object], *, warmup: int) -> tuple[torch.cuda.CUDAGraph, object]: + warmup_stream = torch.cuda.Stream() + warmup_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup_stream): + for _ in range(warmup): + fn() + torch.cuda.current_stream().wait_stream(warmup_stream) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_outputs = fn() + torch.cuda.synchronize() + return graph, graph_outputs + + +def cuda_graph_time_ms( + graph: torch.cuda.CUDAGraph, + *, + iters: int, + samples: int, +) -> list[float]: + out = [] + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + for _ in range(samples): + start.record() + for _ in range(iters): + graph.replay() + end.record() + torch.cuda.synchronize() + out.append(start.elapsed_time(end) / iters) + return out + + +def output_get(outputs: object, names: tuple[str, ...], indices: tuple[int, ...]) -> torch.Tensor: + """Fetch a named tensor from cuDNN TupleDict-like outputs with tuple fallback.""" + for name in names: + try: + value = outputs[name] # type: ignore[index] + if isinstance(value, torch.Tensor): + return value + except (KeyError, IndexError, TypeError, AttributeError): + pass + for index in indices: + try: + value = outputs[index] # type: ignore[index] + if isinstance(value, torch.Tensor): + return value + except (KeyError, IndexError, TypeError, AttributeError): + pass + keys = list(outputs.keys()) if hasattr(outputs, "keys") else type(outputs).__name__ # type: ignore[attr-defined] + raise AssertionError(f"could not find output tensor {names}; available={keys}") + + +def check_tensor_contract( + *, + name: str, + tensor: torch.Tensor, + shape: tuple[int, ...], + dtype: torch.dtype, +) -> None: + if tuple(tensor.shape) != shape: + raise AssertionError(f"{name} shape mismatch: got {tuple(tensor.shape)}, expected {shape}") + if tensor.dtype != dtype: + raise AssertionError(f"{name} dtype mismatch: got {tensor.dtype}, expected {dtype}") + if not tensor.is_cuda: + raise AssertionError(f"{name} must be a CUDA tensor") + + +def check_mxfp8_payload(*, tensor: object, rowwise_data: torch.Tensor, shape: tuple[int, ...]) -> None: + if tuple(rowwise_data.shape) != shape: + raise AssertionError(f"unfused MXFP8 rowwise payload shape mismatch: got {tuple(rowwise_data.shape)}, expected {shape}") + if rowwise_data.dtype not in (torch.uint8, torch.float8_e4m3fn): + raise AssertionError("unfused MXFP8 rowwise payload dtype mismatch: " f"got {rowwise_data.dtype}, expected packed uint8 or float8_e4m3fn") + if not rowwise_data.is_cuda: + raise AssertionError("unfused MXFP8 rowwise payload must be a CUDA tensor") + + if getattr(tensor, "_fp8_dtype", None) is None: + raise AssertionError("unfused output is missing MXFP8 dtype metadata") + if getattr(tensor, "_rowwise_scale_inv", None) is None: + raise AssertionError("unfused output is missing rowwise MXFP8 scales") + + +def check_functional_outputs( + *, + activation: str, + fused_outputs: object, + unfused_output: object, + m: int, + out_n: int, + dtype: torch.dtype, +) -> None: + if activation in FORWARD_GLU or activation in FORWARD_RELU: + fused_primary = output_get(fused_outputs, ("d_tensor", "D"), (1, 0)) + fused_name = "fused d_tensor" + else: + fused_primary = output_get(fused_outputs, ("d_row_tensor", "d_tensor", "D"), (0,)) + fused_name = "fused d_row_tensor" + + check_tensor_contract( + name=fused_name, + tensor=fused_primary, + shape=(m, out_n, 1), + dtype=torch.float8_e4m3fn, + ) + + if not hasattr(unfused_output, "shape"): + raise AssertionError(f"unfused output has no shape attribute: {type(unfused_output).__name__}") + if tuple(unfused_output.shape) != (m, out_n): # type: ignore[attr-defined] + raise AssertionError(f"unfused output shape mismatch: got {tuple(unfused_output.shape)}, expected {(m, out_n)}") # type: ignore[attr-defined] + if getattr(unfused_output, "dtype", None) != dtype: + raise AssertionError(f"unfused output dtype mismatch: got {getattr(unfused_output, 'dtype', None)}, expected {dtype}") + + rowwise_data = getattr(unfused_output, "_rowwise_data", None) + if rowwise_data is None: + raise AssertionError("unfused output is not rowwise MXFP8-quantized") + check_mxfp8_payload(tensor=unfused_output, rowwise_data=rowwise_data, shape=(m, out_n)) + + +def make_shared_functional_case( + *, + activation: str, + tex, + run_grouped_gemm, + experts: int, + tokens: int, + m: int, + k: int, + n: int, + dtype: torch.dtype, + device: torch.device, + wrappers, + use_dynamic_sched: bool, + reference_backend: str, +) -> tuple[Callable[[], object], Callable[[], object]]: + gemm_n, c_cols, _ = activation_shape(activation, n) + norm_const = 0.01 + + a = torch.randn(m, k, dtype=dtype, device=device) * 0.1 + weight = torch.randn(experts, k, gemm_n, dtype=dtype, device=device) * 0.1 + + q_a = make_mxfp8_quantizer(rowwise=True, columnwise=False) + q_w = make_mxfp8_quantizer(rowwise=True, columnwise=True) + q_out = make_mxfp8_quantizer(rowwise=True, columnwise=True) + + a_parts = [q_a(a[i * tokens : (i + 1) * tokens]) for i in range(experts)] + w_parts = [q_w(weight[i]) for i in range(experts)] + fused_weight = weight + if activation in FORWARD_GLU: + fused_weight = torch.stack( + [pack_glu_interleaved(weight[i]) for i in range(experts)], + dim=0, + ) + fused_w_parts = [q_w(fused_weight[i]) for i in range(experts)] + + a_payload = torch.cat([view_as_fp8_e4m3(part._rowwise_data) for part in a_parts], dim=0).contiguous() + a_scale = torch.cat([view_as_e8m0(part._rowwise_scale_inv) for part in a_parts], dim=0) + b_payload = torch.stack( + [view_as_fp8_e4m3(part._columnwise_data).T.contiguous() for part in fused_w_parts], + dim=0, + ).permute(1, 2, 0) + b_scale = torch.stack( + [view_as_e8m0(part._columnwise_scale_inv).T.contiguous() for part in fused_w_parts], + dim=2, + ) + + act_input = None + c_tensor = None + if c_cols is not None: + act_input = torch.randn(m, c_cols, dtype=dtype, device=device) * 0.1 + fused_act_input = pack_glu_interleaved(act_input) if activation in BACKWARD_GLU else act_input + c_tensor = k_major_3d(fused_act_input) + + fused_inputs: dict[str, torch.Tensor | None] = { + "a_tensor": k_major_3d(a_payload), + "b_tensor": b_payload, + "sfa_tensor": compact_sf_to_cudnn_layout( + a_scale.reshape(m, k // MXFP8_BLOCK, 1), + mn=m, + nk=k, + groups=1, + ), + "sfb_tensor": compact_sf_to_cudnn_layout( + b_scale, + mn=gemm_n, + nk=k, + groups=experts, + ), + "padded_offsets": torch.arange( + tokens, + m + 1, + tokens, + dtype=torch.int32, + device=device, + ), + "alpha_tensor": torch.ones(experts, dtype=torch.float32, device=device), + "beta_tensor": torch.ones(experts, dtype=torch.float32, device=device), + "prob_tensor": torch.ones(m, 1, 1, dtype=torch.float32, device=device), + "dprob_tensor": torch.zeros(m, 1, 1, dtype=torch.float32, device=device), + "norm_const_tensor": torch.tensor([norm_const], dtype=torch.float32, device=device), + "c_tensor": c_tensor, + } + + fused = make_fused_fn( + activation=activation, + inputs=fused_inputs, + dtype=dtype, + use_dynamic_sched=use_dynamic_sched, + wrappers=wrappers, + ) + + gemm_out = torch.empty(m, gemm_n, dtype=dtype, device=device) + prob = torch.ones(m, 1, dtype=dtype, device=device) + dact_dy = torch.empty_like(gemm_out) + + if activation in FORWARD_GLU or activation in FORWARD_RELU: + + def unfused() -> object: + run_grouped_gemm(w_parts, a_parts, gemm_out) + if activation == "geglu": + return cudnn_geglu_reference(gemm_out, q_out, reference_backend) + return getattr(tex, activation)(gemm_out, q_out) + + return fused, unfused + + if activation in BACKWARD_GLU or activation in BACKWARD_RELU: + + def unfused() -> object: + run_grouped_gemm(w_parts, a_parts, gemm_out) + if activation == "dgeglu": + return cudnn_dgeglu_reference(gemm_out, act_input, prob, q_out, reference_backend) + if activation == "dsrelu": + return cudnn_dsrelu_reference(gemm_out, act_input, q_out) + torch.mul(gemm_out, prob, out=dact_dy) + return getattr(tex, activation)(dact_dy, act_input, q_out) + + return fused, unfused + + raise ValueError(f"Unsupported activation: {activation}") + + +def compare_functional_outputs( + *, + activation: str, + fused_outputs: object, + unfused_output: object, + m: int, + out_n: int, + dtype: torch.dtype, + atol: float, + rtol: float, +) -> tuple[float, float, float]: + check_functional_outputs( + activation=activation, + fused_outputs=fused_outputs, + unfused_output=unfused_output, + m=m, + out_n=out_n, + dtype=dtype, + ) + + if activation in FORWARD_GLU or activation in FORWARD_RELU: + fused_primary = output_get(fused_outputs, ("d_tensor", "D"), (1, 0)) + else: + fused_primary = output_get(fused_outputs, ("d_row_tensor", "d_tensor", "D"), (0,)) + + sfd_row = output_get(fused_outputs, ("sfd_row_tensor",), (4, 5)) + fused_deq = dequantize_cudnn_mxfp8_output( + fused_primary, + sfd_row, + out_n=out_n, + norm_const=0.01, + ).reshape(m, out_n) + if activation in BACKWARD_GLU: + fused_deq = unpack_glu_interleaved(fused_deq) + unfused_deq = unfused_output.dequantize(dtype=torch.float32).reshape(m, out_n) # type: ignore[attr-defined] + + diff = (fused_deq - unfused_deq).abs() + max_abs = diff.max().item() + mean_abs = diff.mean().item() + p99_abs = torch.quantile(diff.flatten(), 0.99).item() + + try: + torch.testing.assert_close(fused_deq, unfused_deq, atol=atol, rtol=rtol) + except AssertionError as exc: + raise AssertionError( + f"{activation} output mismatch: max_abs={max_abs:.6g}, " f"p99_abs={p99_abs:.6g}, mean_abs={mean_abs:.6g}, " f"atol={atol}, rtol={rtol}" + ) from exc + + return max_abs, p99_abs, mean_abs + + +def run_functional_tests( + *, + activations: tuple[str, ...], + args, + tex, + general_grouped_gemm, + wrappers, + device: torch.device, +) -> None: + e = args.functional_experts + t = args.functional_tokens + k = args.functional_k + n = args.functional_n + m = e * t + check_dims(t, k, n) + + print() + print("functional tests:") + print(" scope: shared-input dequantized output comparison") + print(f" shape: experts={e}, tokens/expert={t}, M={m}, K={k}, N={n}") + print(f" tolerance: atol={args.functional_atol}, rtol={args.functional_rtol}") + if any(activation in {"geglu", "dgeglu"} for activation in activations): + print(" geglu reference: " f"{args.functional_reference_backend} Megatron-style QuickGeGLU + clamp + linear-offset semantics") + if "dsrelu" in activations: + print(" dsrelu reference: cuDNN FE operand order, GEMM output is saved activation input") + + run_grouped_gemm = make_grouped_gemm_runner( + general_grouped_gemm, + experts=e, + tokens=t, + dtype=args.dtype, + ) + + with torch.no_grad(): + for index, activation in enumerate(activations): + torch.manual_seed(args.seed + 1000 + index) + _, _, out_n = activation_shape(activation, n) + fused, unfused = make_shared_functional_case( + activation=activation, + tex=tex, + run_grouped_gemm=run_grouped_gemm, + experts=e, + tokens=t, + m=m, + k=k, + n=n, + dtype=args.dtype, + device=device, + wrappers=wrappers, + use_dynamic_sched=not args.no_dynamic_sched, + reference_backend=args.functional_reference_backend, + ) + + fused_outputs = fused() + unfused_output = unfused() + torch.cuda.synchronize() + max_abs, p99_abs, mean_abs = compare_functional_outputs( + activation=activation, + fused_outputs=fused_outputs, + unfused_output=unfused_output, + m=m, + out_n=out_n, + dtype=args.dtype, + atol=args.functional_atol, + rtol=args.functional_rtol, + ) + print(f" {activation:<8} PASS " f"max_abs={max_abs:.4g} p99_abs={p99_abs:.4g} mean_abs={mean_abs:.4g}") + + +def make_grouped_gemm_runner(general_grouped_gemm, *, experts: int, tokens: int, dtype: torch.dtype): + """Handle TE grouped-GEMM signature differences across NGC and TE main.""" + params = inspect.signature(general_grouped_gemm).parameters + if "quantization_params" in params: + + def run(w_parts, a_parts, gemm_out) -> None: + general_grouped_gemm( + w_parts, + a_parts, + [gemm_out], + [None] * experts, + dtype, + layout="NN", + m_splits=[tokens] * experts, + use_split_accumulator=True, + single_output=True, + ) + + return run + + from transformer_engine.pytorch.module.base import get_multi_stream_cublas_workspace + + workspace = get_multi_stream_cublas_workspace() + + def run(w_parts, a_parts, gemm_out) -> None: + general_grouped_gemm( + w_parts, + a_parts, + [gemm_out], + dtype, + workspace, + layout="NN", + m_splits=[tokens] * experts, + use_split_accumulator=True, + single_output=True, + ) + + return run + + +def make_fused_fn( + *, + activation: str, + inputs: dict[str, torch.Tensor | None], + dtype: torch.dtype, + use_dynamic_sched: bool, + wrappers, +): + common = { + "a_tensor": inputs["a_tensor"], + "b_tensor": inputs["b_tensor"], + "sfa_tensor": inputs["sfa_tensor"], + "sfb_tensor": inputs["sfb_tensor"], + "padded_offsets": inputs["padded_offsets"], + "alpha_tensor": inputs["alpha_tensor"], + "norm_const_tensor": inputs["norm_const_tensor"], + "acc_dtype": torch.float32, + "d_dtype": torch.float8_e4m3fn, + "cd_major": "n", + "mma_tiler_mn": (256, 256), + "cluster_shape_mn": (2, 1), + "sf_vec_size": 32, + "vector_f32": False, + "m_aligned": M_ALIGNED, + "discrete_col_sfd": False, + } + + def current_stream(): + return wrappers["cuda"].CUstream(torch.cuda.current_stream().cuda_stream) + + if activation in FORWARD_GLU: + + def fused() -> object: + return wrappers["glu"]( + **common, + prob_tensor=inputs["prob_tensor"], + c_dtype=dtype, + act_func=activation, + use_dynamic_sched=use_dynamic_sched, + current_stream=current_stream(), + ) + + return fused + + if activation in FORWARD_RELU: + + def fused() -> object: + return wrappers["srelu"]( + **common, + prob_tensor=inputs["prob_tensor"], + c_dtype=dtype, + use_dynamic_sched=use_dynamic_sched, + current_stream=current_stream(), + ) + + return fused + + if activation in BACKWARD_GLU: + + def fused() -> object: + inputs["dprob_tensor"].zero_() + return wrappers["dglu"]( + **common, + c_tensor=inputs["c_tensor"], + beta_tensor=inputs["beta_tensor"], + prob_tensor=inputs["prob_tensor"], + dprob_tensor=inputs["dprob_tensor"], + generate_dbias=False, + act_func=activation, + use_dynamic_sched=use_dynamic_sched, + current_stream=current_stream(), + ) + + return fused + + if activation in BACKWARD_RELU: + + def fused() -> object: + inputs["dprob_tensor"].zero_() + return wrappers["dsrelu"]( + **common, + c_tensor=inputs["c_tensor"], + prob_tensor=inputs["prob_tensor"], + dprob_tensor=inputs["dprob_tensor"], + generate_dbias=False, + use_dynamic_sched=use_dynamic_sched, + current_stream=current_stream(), + ) + + return fused + + raise ValueError(f"Unsupported activation: {activation}") + + +def make_unfused_te_fn( + *, + activation: str, + tex, + run_grouped_gemm, + experts: int, + tokens: int, + m: int, + k: int, + n: int, + dtype: torch.dtype, + device: torch.device, +): + gemm_n, c_cols, _ = activation_shape(activation, n) + + a = torch.randn(m, k, dtype=dtype, device=device) * 0.1 + weight = torch.randn(experts, k, gemm_n, dtype=dtype, device=device) * 0.1 + + q_a = make_mxfp8_quantizer(rowwise=True, columnwise=False) + q_w = make_mxfp8_quantizer(rowwise=True, columnwise=True) + q_out = make_mxfp8_quantizer(rowwise=True, columnwise=True) + + a_parts = [q_a(a[i * tokens : (i + 1) * tokens]) for i in range(experts)] + w_parts = [q_w(weight[i]) for i in range(experts)] + gemm_out = torch.empty(m, gemm_n, dtype=dtype, device=device) + + prob = torch.ones(m, 1, dtype=dtype, device=device) + dact_dy = torch.empty_like(gemm_out) + act_input = None + if c_cols is not None: + act_input = torch.randn(m, c_cols, dtype=dtype, device=device) * 0.1 + + if activation in FORWARD_GLU or activation in FORWARD_RELU: + + def unfused() -> object: + run_grouped_gemm(w_parts, a_parts, gemm_out) + return getattr(tex, activation)(gemm_out, q_out) + + return unfused + + if activation in BACKWARD_GLU or activation in BACKWARD_RELU: + + def unfused() -> object: + run_grouped_gemm(w_parts, a_parts, gemm_out) + torch.mul(gemm_out, prob, out=dact_dy) + return getattr(tex, activation)(dact_dy, act_input, q_out) + + return unfused + + raise ValueError(f"Unsupported activation: {activation}") + + +def make_pytorch_baseline_fn( + *, + activation: str, + run_grouped_gemm, + experts: int, + tokens: int, + m: int, + k: int, + n: int, + dtype: torch.dtype, + device: torch.device, +): + """TE MXFP8 grouped GEMM followed by PyTorch pointwise activation and TE quantize.""" + gemm_n, c_cols, _ = activation_shape(activation, n) + + a = torch.randn(m, k, dtype=dtype, device=device) * 0.1 + weight = torch.randn(experts, k, gemm_n, dtype=dtype, device=device) * 0.1 + + q_a = make_mxfp8_quantizer(rowwise=True, columnwise=False) + q_w = make_mxfp8_quantizer(rowwise=True, columnwise=True) + q_out = make_mxfp8_quantizer(rowwise=True, columnwise=True) + + a_parts = [q_a(a[i * tokens : (i + 1) * tokens]) for i in range(experts)] + w_parts = [q_w(weight[i]) for i in range(experts)] + gemm_out = torch.empty(m, gemm_n, dtype=dtype, device=device) + + prob = torch.ones(m, 1, dtype=dtype, device=device) + act_input = None + if c_cols is not None: + act_input = torch.randn(m, c_cols, dtype=dtype, device=device) * 0.1 + + def pytorch_baseline() -> object: + run_grouped_gemm(w_parts, a_parts, gemm_out) + out = pytorch_activation_formula( + activation=activation, + tensor=gemm_out, + activation_input=act_input, + prob=prob, + ) + return q_out(out) + + return pytorch_baseline + + +def benchmark_activation( + *, + activation: str, + args, + tex, + general_grouped_gemm, + wrappers, + device: torch.device, +) -> tuple[str, float, float, float | None, list[float], list[float], list[float] | None]: + e, k, n = args.experts, args.k, args.n + t = benchmark_tokens_for_activation(args, activation) + m = e * t + gemm_n, c_cols, out_n = activation_shape(activation, n) + + print() + print(f"activation: {activation}") + print(f"shape: experts={e}, tokens/expert={t}, M={m}, K={k}, N={n}") + print(f"fused GEMM N={gemm_n}, fused C input N={c_cols}, activation output N={out_n}") + + fused_inputs = make_cudnn_inputs( + activation=activation, + experts=e, + tokens=t, + k=k, + n=n, + c_dtype=args.dtype, + device=device, + ) + fused = make_fused_fn( + activation=activation, + inputs=fused_inputs, + dtype=args.dtype, + use_dynamic_sched=not args.no_dynamic_sched, + wrappers=wrappers, + ) + + run_grouped_gemm = make_grouped_gemm_runner( + general_grouped_gemm, + experts=e, + tokens=t, + dtype=args.dtype, + ) + unfused = make_unfused_te_fn( + activation=activation, + tex=tex, + run_grouped_gemm=run_grouped_gemm, + experts=e, + tokens=t, + m=m, + k=k, + n=n, + dtype=args.dtype, + device=device, + ) + fused_graph, fused_outputs = capture_cuda_graph(fused, warmup=args.warmup) + unfused_graph, unfused_outputs = capture_cuda_graph(unfused, warmup=args.warmup) + graph_outputs = [fused_outputs, unfused_outputs] + + pytorch_graph = None + pytorch_ms = None + pytorch_avg = None + if args.include_pytorch_baseline: + pytorch_baseline = make_pytorch_baseline_fn( + activation=activation, + run_grouped_gemm=run_grouped_gemm, + experts=e, + tokens=t, + m=m, + k=k, + n=n, + dtype=args.dtype, + device=device, + ) + pytorch_graph, pytorch_outputs = capture_cuda_graph(pytorch_baseline, warmup=args.warmup) + graph_outputs.append(pytorch_outputs) + + fused_ms = cuda_graph_time_ms(fused_graph, iters=args.iters, samples=args.samples) + unfused_ms = cuda_graph_time_ms(unfused_graph, iters=args.iters, samples=args.samples) + if pytorch_graph is not None: + pytorch_ms = cuda_graph_time_ms(pytorch_graph, iters=args.iters, samples=args.samples) + + fused_avg = statistics.mean(fused_ms) + unfused_avg = statistics.mean(unfused_ms) + if pytorch_ms is not None: + pytorch_avg = statistics.mean(pytorch_ms) + print(f"fused cuDNN avg: {fused_avg:.3f} ms") + print(f"unfused TE avg: {unfused_avg:.3f} ms") + print(f"speedup vs TE: {unfused_avg / fused_avg:.2f}x") + if pytorch_avg is not None: + print(f"PyTorch avg: {pytorch_avg:.3f} ms") + print(f"speedup vs torch:{pytorch_avg / fused_avg:6.2f}x") + print(f"fused samples: {[round(x, 3) for x in fused_ms]}") + print(f"TE samples: {[round(x, 3) for x in unfused_ms]}") + if pytorch_ms is not None: + print(f"PyTorch samples: {[round(x, 3) for x in pytorch_ms]}") + del graph_outputs + return activation, fused_avg, unfused_avg, pytorch_avg, fused_ms, unfused_ms, pytorch_ms + + +class _Tee: + """Write to several streams at once (e.g. real stdout and a capture buffer).""" + + def __init__(self, *streams) -> None: + self.streams = streams + + def write(self, data: str) -> int: + for stream in self.streams: + stream.write(data) + return len(data) + + def flush(self) -> None: + for stream in self.streams: + stream.flush() + + +# Lines kept in output.txt: device, dtype, library versions, and shapes. +OUTPUT_TXT_PREFIXES = ( + "device:", + "dtype:", + "torch:", + "transformer_engine:", + "nvidia-cudnn-frontend:", + "nvidia-cutlass-dsl:", + "shape:", + "base shape:", + "GLU bprop shape:", +) + +# Per-activation shape lines move next to each bar in the chart, so they are +# dropped from the chart's header block (but kept in output.txt). +SHAPE_PREFIXES = ("shape:", "base shape:", "GLU bprop shape:") + +DISPLAY_NAMES = { + "swiglu": "SwiGLU", + "dswiglu": "dSwiGLU", + "srelu": "SReLU", + "dsrelu": "dSReLU", + "geglu": "QuickGeGLU", + "dgeglu": "dQuickGeGLU", +} + + +# Preferred header fonts; matplotlib falls back to the first one available. +HEADER_FONT = ["Helvetica Neue", "Helvetica", "Arial", "DejaVu Sans"] +HEADER_ACCENT = "#00866a" + +# Prettier labels for the environment metadata lines in the chart header. +LABEL_PRETTY = { + "device": "Device", + "dtype": "Dtype", + "torch": "PyTorch", + "transformer_engine": "Transformer Engine", + "nvidia-cudnn-frontend": "cuDNN Frontend", + "nvidia-cutlass-dsl": "cuTLASS DSL", +} + + +def display_name(activation: str) -> str: + return DISPLAY_NAMES.get(activation, activation) + + +def family_color(activation: str) -> str: + """Bar color grouped by activation family (GLU=green, ReLU=blue, GeGLU=amber).""" + if activation in ("swiglu", "dswiglu"): + return "#0db58f" + if activation in ("srelu", "dsrelu"): + return "#6c8ff0" + return "#f2a51f" + + +def format_summary(results: list, *, include_pytorch: bool) -> list[str]: + """Render the benchmark summary table as a list of lines.""" + lines = ["summary:"] + if include_pytorch: + lines.append(f"{'activation':<10} {'cudnn_ms':>10} {'te_ms':>10} {'pytorch_ms':>12} " f"{'te/cudnn':>10} {'torch/cudnn':>12}") + for activation, fused_avg, unfused_avg, pytorch_avg, _, _, _ in results: + assert pytorch_avg is not None + lines.append( + f"{activation:<10} {fused_avg:10.3f} {unfused_avg:10.3f} {pytorch_avg:12.3f} " + f"{unfused_avg / fused_avg:9.2f}x {pytorch_avg / fused_avg:11.2f}x" + ) + else: + lines.append(f"{'activation':<10} {'cudnn_ms':>10} {'te_ms':>10} {'te/cudnn':>10}") + for activation, fused_avg, unfused_avg, _, _, _, _ in results: + lines.append(f"{activation:<10} {fused_avg:10.3f} {unfused_avg:10.3f} {unfused_avg / fused_avg:9.2f}x") + return lines + + +def write_results_png( + results: list, + *, + include_pytorch: bool, + shapes: dict[str, str], + info_lines: list[str], + path: str, +) -> None: + """Render speedups as a horizontal bar chart with a metadata header block.""" + try: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + from matplotlib import font_manager + from matplotlib.patches import Patch + except ImportError: + print(f"matplotlib not available; skipping {path}") + return + + # Only request fonts that exist so matplotlib does not warn on every glyph. + available = {f.name for f in font_manager.fontManager.ttflist} + header_font = [f for f in HEADER_FONT if f in available] or ["sans-serif"] + + activations = [r[0] for r in results] + te_speedups = [unfused / fused for _, fused, unfused, _, _, _, _ in results] + torch_speedups = [(pytorch / fused) if pytorch is not None else None for _, fused, _, pytorch, _, _, _ in results] + have_torch = include_pytorch and all(s is not None for s in torch_speedups) + + # Header highlights: prominent speedup stats plus the environment metadata. + avg_te = sum(te_speedups) / len(te_speedups) + best_idx = max(range(len(te_speedups)), key=lambda i: te_speedups[i]) + stat_groups = [ + (f"{avg_te:.2f}x", "average speedup vs TE"), + (f"{te_speedups[best_idx]:.2f}x", f"best: {display_name(activations[best_idx])}"), + ] + if have_torch: + avg_torch = sum(torch_speedups) / len(torch_speedups) + stat_groups.append((f"{avg_torch:.2f}x", "average speedup vs PyTorch")) + stat_groups.append((str(len(results)), "activations benchmarked")) + + env_pairs = [] + for line in info_lines: + key, _, value = line.partition(":") + env_pairs.append((LABEL_PRETTY.get(key.strip(), key.strip()), value.strip())) + env_rows = (len(env_pairs) + 1) // 2 + + # Round the x-axis up to a clean 0.5 multiple that clears every bar. + raw_max = max(list(te_speedups) + [s for s in torch_speedups if s is not None] + [1.0]) + tick_max = 0.5 + while tick_max < raw_max + 1e-9: + tick_max += 0.5 + tick_max = max(tick_max, 1.5) + gutter = 0.55 * tick_max # left label column, in data units + + n = len(results) + head_h = 1.35 + 0.42 * env_rows + chart_h = 0.95 * n + 0.7 + fig = plt.figure(figsize=(14.0, head_h + chart_h + 0.5)) + fig.suptitle( + "cuDNN fused-epilogue speedups vs Transformer Engine", + fontsize=20, + fontweight="bold", + color="#0f172a", + family=header_font, + x=0.07, + ha="left", + ) + grid = fig.add_gridspec(2, 1, height_ratios=[head_h, chart_h], hspace=0.08) + + ax_info = fig.add_subplot(grid[0]) + ax_info.axis("off") + + # Stat highlights across the top: big accent number + muted caption. + n_stat = len(stat_groups) + for i, (value, caption) in enumerate(stat_groups): + x = i / n_stat + ax_info.text(x, 0.98, value, transform=ax_info.transAxes, va="top", ha="left", fontsize=26, fontweight="bold", color=HEADER_ACCENT, family=header_font) + ax_info.text(x, 0.60, caption, transform=ax_info.transAxes, va="top", ha="left", fontsize=12, color="#64748b", family=header_font) + + # Thin divider between the highlights and the environment metadata. + ax_info.plot([0.0, 1.0], [0.50, 0.50], transform=ax_info.transAxes, color="#e2e8f0", lw=1.2) + + # Environment metadata as a two-column label/value grid. + label_x = (0.0, 0.50) + value_x = (0.16, 0.74) + row_top = 0.40 + row_step = 0.40 / max(env_rows, 1) + for idx, (label, value) in enumerate(env_pairs): + col = idx % 2 + y = row_top - (idx // 2) * row_step + ax_info.text( + label_x[col], y, label, transform=ax_info.transAxes, va="top", ha="left", fontsize=12, fontweight="bold", color="#334155", family=header_font + ) + ax_info.text(value_x[col], y, value, transform=ax_info.transAxes, va="top", ha="left", fontsize=12, color="#0f172a", family=header_font) + + ax = fig.add_subplot(grid[1]) + for spine in ("top", "right", "left"): + ax.spines[spine].set_visible(False) + ax.set_yticks([]) + ax.set_xlim(-gutter, tick_max * 1.12) + ax.set_ylim(-0.7, n - 0.3) + ax.invert_yaxis() + + xticks = [] + tick = 0.0 + while tick <= tick_max + 1e-9: + xticks.append(round(tick, 1)) + tick += 0.5 + ax.set_xticks(xticks) + ax.set_xticklabels([f"{x:g}x" if x > 0 else "0" for x in xticks], color="#64748b") + ax.tick_params(axis="x", length=0) + for x in xticks: + ax.axvline(x, color="#eef2f7", lw=1.0, zorder=0) + ax.axvline(0.0, color="#cbd5e1", lw=1.2, zorder=1) + ax.axvline(1.0, color="#94a3b8", lw=1.4, ls=(0, (5, 7)), zorder=1) + ax.text(1.0, -0.62, "TE parity", color="#64748b", fontsize=11, ha="center", va="bottom") + + label_pad = tick_max * 0.02 + for i, activation in enumerate(activations): + color = family_color(activation) + shape = shapes.get(activation, "") + ax.text(-gutter * 0.96, i - 0.14, display_name(activation), ha="left", va="center", fontweight="bold", fontsize=13, color="#0f172a") + if shape: + ax.text(-gutter * 0.96, i + 0.18, shape, ha="left", va="center", fontsize=9.5, color="#64748b") + + if have_torch: + ax.barh(i - 0.17, te_speedups[i], height=0.30, color=color, zorder=3) + ax.barh(i + 0.17, torch_speedups[i], height=0.30, color=color, alpha=0.5, zorder=3) + ax.text(te_speedups[i] + label_pad, i - 0.17, f"{te_speedups[i]:.2f}x", ha="left", va="center", fontsize=11, fontweight="bold", color="#0f172a") + ax.text(torch_speedups[i] + label_pad, i + 0.17, f"{torch_speedups[i]:.2f}x", ha="left", va="center", fontsize=10, color="#475569") + else: + ax.barh(i, te_speedups[i], height=0.5, color=color, zorder=3) + ax.text(te_speedups[i] + label_pad, i, f"{te_speedups[i]:.2f}x", ha="left", va="center", fontsize=12, fontweight="bold", color="#0f172a") + + if have_torch: + handles = [ + Patch(facecolor="#475569", label="cuDNN vs TE"), + Patch(facecolor="#475569", alpha=0.5, label="cuDNN vs PyTorch"), + ] + ax.legend( + handles=handles, + loc="lower right", + bbox_to_anchor=(1.0, 1.0), + ncol=2, + frameon=False, + fontsize=10, + ) + + ax.set_xlabel("speedup (higher is better)", color="#475569", fontsize=11) + + fig.savefig(path, dpi=150, bbox_inches="tight", facecolor="white") + plt.close(fig) + print(f"wrote {path}") + + +def write_output_files( + *, + output_dir: str, + verbose_text: str, + results: list, + args: argparse.Namespace, +) -> None: + """Write output_verbose.txt, output.txt, and results.png into output_dir.""" + include_pytorch = args.include_pytorch_baseline + verbose_path = os.path.join(output_dir, "output_verbose.txt") + with open(verbose_path, "w") as handle: + handle.write(verbose_text) + print(f"wrote {verbose_path}") + + summary_lines = format_summary(results, include_pytorch=include_pytorch) if results else [] + # Preserve order but drop duplicates (per-activation "shape:" lines repeat the top-level one). + key_lines: list[str] = [] + seen: set[str] = set() + for line in verbose_text.splitlines(): + if line.startswith(OUTPUT_TXT_PREFIXES) and line not in seen: + seen.add(line) + key_lines.append(line) + output_path = os.path.join(output_dir, "output.txt") + with open(output_path, "w") as handle: + for line in key_lines: + handle.write(line + "\n") + if summary_lines: + handle.write("\n") + for line in summary_lines: + handle.write(line + "\n") + print(f"wrote {output_path}") + + if results: + # Per-activation GEMM shape: experts x tokens/expert x GEMM-N x K. + shapes = {} + for activation, *_ in results: + tokens = benchmark_tokens_for_activation(args, activation) + gemm_n, _, _ = activation_shape(activation, args.n) + shapes[activation] = f"GEMM: {args.experts} x {tokens} x {gemm_n} x {args.k}" + # The chart header keeps device/dtype/version info; shapes move next to each bar. + chart_info_lines = [line for line in key_lines if not line.startswith(SHAPE_PREFIXES)] + write_results_png( + results, + include_pytorch=include_pytorch, + shapes=shapes, + info_lines=chart_info_lines, + path=os.path.join(output_dir, "results.png"), + ) + + +def run_benchmarks(args: argparse.Namespace) -> list: + """Run functional tests and/or benchmarks, printing progress. Returns timing results.""" + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + torch.cuda.set_device(args.device) + device = torch.device("cuda", args.device) + + import transformer_engine.pytorch # noqa: F401 + import transformer_engine_torch as tex + from cuda.bindings import driver as cuda + from cudnn import ( + grouped_gemm_dglu_wrapper_sm100, + grouped_gemm_dsrelu_wrapper_sm100, + grouped_gemm_glu_wrapper_sm100, + grouped_gemm_srelu_wrapper_sm100, + ) + from transformer_engine.pytorch.cpp_extensions import general_grouped_gemm + + torch.manual_seed(args.seed) + + print(f"device: {torch.cuda.get_device_name(args.device)}") + if not args.functional_only: + if args.glu_bprop_tokens is None: + print(f"shape: experts={args.experts}, tokens/expert={args.tokens}, " f"M={args.experts * args.tokens}, K={args.k}, N={args.n}") + else: + print(f"base shape: experts={args.experts}, tokens/expert={args.tokens}, " f"M={args.experts * args.tokens}, K={args.k}, N={args.n}") + print( + f"GLU bprop shape: experts={args.experts}, " + f"tokens/expert={args.glu_bprop_tokens}, " + f"M={args.experts * args.glu_bprop_tokens}, K={args.k}, N={args.n}" + ) + print(f"dtype: {args.dtype}") + print(f"torch: {torch.__version__}") + print(f"transformer_engine: {package_version('transformer-engine')}") + print(f"nvidia-cudnn-frontend: {package_version('nvidia-cudnn-frontend')}") + print(f"nvidia-cutlass-dsl: {package_version('nvidia-cutlass-dsl')}") + + wrappers = { + "cuda": cuda, + "glu": grouped_gemm_glu_wrapper_sm100, + "srelu": grouped_gemm_srelu_wrapper_sm100, + "dsrelu": grouped_gemm_dsrelu_wrapper_sm100, + "dglu": grouped_gemm_dglu_wrapper_sm100, + } + activations = ACTIVATIONS if args.activation == "all" else (args.activation,) + + if args.functional_test or args.functional_only: + run_functional_tests( + activations=activations, + args=args, + tex=tex, + general_grouped_gemm=general_grouped_gemm, + wrappers=wrappers, + device=device, + ) + if args.functional_only: + return [] + + for tokens in sorted({benchmark_tokens_for_activation(args, activation) for activation in activations}): + check_dims(tokens, args.k, args.n) + print("timing mode: CUDA graph replay") + + results = [ + benchmark_activation( + activation=activation, + args=args, + tex=tex, + general_grouped_gemm=general_grouped_gemm, + wrappers=wrappers, + device=device, + ) + for activation in activations + ] + + if len(results) > 1: + print() + for line in format_summary(results, include_pytorch=args.include_pytorch_baseline): + print(line) + + return results + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--activation", choices=("all",) + ACTIVATIONS, default="all") + parser.add_argument("--experts", type=int, default=8) + parser.add_argument("--tokens", type=int, default=4096, help="tokens per expert") + parser.add_argument( + "--glu-bprop-tokens", + type=int, + default=None, + help="tokens per expert for dswiglu/dgeglu; use half of --tokens for paired GLU bprop", + ) + parser.add_argument("--k", type=int, default=8192, help="input / gradient hidden size") + parser.add_argument("--n", type=int, default=4096, help="activation output hidden size") + parser.add_argument("--dtype", type=parse_dtype, default=torch.bfloat16) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iters", type=int, default=20) + parser.add_argument("--samples", type=int, default=5) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument("--device", type=int, default=0) + parser.add_argument("--no-dynamic-sched", action="store_true") + parser.add_argument( + "--include-pytorch-baseline", + action="store_true", + help="also time the TE GEMM + PyTorch pointwise baseline", + ) + parser.add_argument( + "--functional-test", + action="store_true", + help="run small execution/shape/dtype tests before benchmarking", + ) + parser.add_argument( + "--functional-only", + action="store_true", + help="run only the functional tests and skip timing", + ) + parser.add_argument("--functional-experts", type=int, default=2) + parser.add_argument("--functional-tokens", type=int, default=256) + parser.add_argument("--functional-k", type=int, default=128) + parser.add_argument("--functional-n", type=int, default=256) + parser.add_argument("--functional-atol", type=float, default=0.5) + parser.add_argument("--functional-rtol", type=float, default=0.5) + parser.add_argument( + "--functional-reference-backend", + choices=("torch_compile", "torch"), + default="torch", + help="reference backend for cuDNN-specific GeGLU/DGeGLU functional checks", + ) + parser.add_argument( + "--output-dir", + type=str, + default=None, + help=( + "if set, write output_verbose.txt (full stdout), output.txt (device/version/" + "dtype/shape/summary), and results.png (summary table) to this directory" + ), + ) + args = parser.parse_args() + + if args.output_dir is None: + run_benchmarks(args) + return + + os.makedirs(args.output_dir, exist_ok=True) + buffer = io.StringIO() + original_stdout = sys.stdout + sys.stdout = _Tee(original_stdout, buffer) + try: + results = run_benchmarks(args) + finally: + sys.stdout = original_stdout + + write_output_files( + output_dir=args.output_dir, + verbose_text=buffer.getvalue(), + results=results, + args=args, + ) + + +if __name__ == "__main__": + main() diff --git a/benchmark/cutedsl_fusion_kernels/results/B200/output.txt b/benchmark/cutedsl_fusion_kernels/results/B200/output.txt new file mode 100644 index 000000000..eec7c8195 --- /dev/null +++ b/benchmark/cutedsl_fusion_kernels/results/B200/output.txt @@ -0,0 +1,16 @@ +device: NVIDIA GB200 +shape: experts=8, tokens/expert=4096, M=32768, K=8192, N=4096 +dtype: torch.bfloat16 +torch: 2.12.0a0+5aff3928d8.nv26.05 +transformer_engine: 2.15.0+42b84005 +nvidia-cudnn-frontend: 1.23.0 +nvidia-cutlass-dsl: 4.4.1 + +summary: +activation cudnn_ms te_ms te/cudnn +swiglu 1.550 1.934 1.25x +dswiglu 0.870 1.734 1.99x +srelu 0.739 1.041 1.41x +dsrelu 0.781 1.341 1.72x +geglu 1.544 1.890 1.22x +dgeglu 0.889 1.820 2.05x diff --git a/benchmark/cutedsl_fusion_kernels/results/B200/output_verbose.txt b/benchmark/cutedsl_fusion_kernels/results/B200/output_verbose.txt new file mode 100644 index 000000000..f3c7abdec --- /dev/null +++ b/benchmark/cutedsl_fusion_kernels/results/B200/output_verbose.txt @@ -0,0 +1,71 @@ +device: NVIDIA GB200 +shape: experts=8, tokens/expert=4096, M=32768, K=8192, N=4096 +dtype: torch.bfloat16 +torch: 2.12.0a0+5aff3928d8.nv26.05 +transformer_engine: 2.15.0+42b84005 +nvidia-cudnn-frontend: 1.23.0 +nvidia-cutlass-dsl: 4.4.1 +timing mode: CUDA graph replay + +activation: swiglu +shape: experts=8, tokens/expert=4096, M=32768, K=8192, N=4096 +fused GEMM N=8192, fused C input N=None, activation output N=4096 +fused cuDNN avg: 1.550 ms +unfused TE avg: 1.934 ms +speedup vs TE: 1.25x +fused samples: [1.357, 1.425, 1.64, 1.667, 1.663] +TE samples: [2.212, 1.937, 1.814, 1.839, 1.869] + +activation: dswiglu +shape: experts=8, tokens/expert=4096, M=32768, K=8192, N=4096 +fused GEMM N=4096, fused C input N=8192, activation output N=8192 +fused cuDNN avg: 0.870 ms +unfused TE avg: 1.734 ms +speedup vs TE: 1.99x +fused samples: [0.806, 0.806, 0.793, 0.963, 0.98] +TE samples: [2.254, 1.639, 1.626, 1.579, 1.571] + +activation: srelu +shape: experts=8, tokens/expert=4096, M=32768, K=8192, N=4096 +fused GEMM N=4096, fused C input N=None, activation output N=4096 +fused cuDNN avg: 0.739 ms +unfused TE avg: 1.041 ms +speedup vs TE: 1.41x +fused samples: [0.721, 0.718, 0.718, 0.734, 0.803] +TE samples: [1.1, 1.107, 1.023, 1.008, 0.97] + +activation: dsrelu +shape: experts=8, tokens/expert=4096, M=32768, K=8192, N=4096 +fused GEMM N=4096, fused C input N=4096, activation output N=4096 +fused cuDNN avg: 0.781 ms +unfused TE avg: 1.341 ms +speedup vs TE: 1.72x +fused samples: [0.731, 0.729, 0.73, 0.848, 0.866] +TE samples: [1.539, 1.386, 1.27, 1.265, 1.243] + +activation: geglu +shape: experts=8, tokens/expert=4096, M=32768, K=8192, N=4096 +fused GEMM N=8192, fused C input N=None, activation output N=4096 +fused cuDNN avg: 1.544 ms +unfused TE avg: 1.890 ms +speedup vs TE: 1.22x +fused samples: [1.358, 1.523, 1.6, 1.647, 1.595] +TE samples: [2.036, 1.827, 1.842, 1.87, 1.874] + +activation: dgeglu +shape: experts=8, tokens/expert=4096, M=32768, K=8192, N=4096 +fused GEMM N=4096, fused C input N=8192, activation output N=8192 +fused cuDNN avg: 0.889 ms +unfused TE avg: 1.820 ms +speedup vs TE: 2.05x +fused samples: [0.8, 0.795, 0.896, 0.951, 1.003] +TE samples: [2.369, 1.717, 1.668, 1.702, 1.644] + +summary: +activation cudnn_ms te_ms te/cudnn +swiglu 1.550 1.934 1.25x +dswiglu 0.870 1.734 1.99x +srelu 0.739 1.041 1.41x +dsrelu 0.781 1.341 1.72x +geglu 1.544 1.890 1.22x +dgeglu 0.889 1.820 2.05x diff --git a/benchmark/cutedsl_fusion_kernels/results/B200/results.png b/benchmark/cutedsl_fusion_kernels/results/B200/results.png new file mode 100644 index 000000000..36795201b Binary files /dev/null and b/benchmark/cutedsl_fusion_kernels/results/B200/results.png differ diff --git a/benchmark/cutedsl_fusion_kernels/results/B300/output.txt b/benchmark/cutedsl_fusion_kernels/results/B300/output.txt new file mode 100644 index 000000000..4ebfb6cab --- /dev/null +++ b/benchmark/cutedsl_fusion_kernels/results/B300/output.txt @@ -0,0 +1,16 @@ +device: NVIDIA GB300 +shape: experts=8, tokens/expert=4096, M=32768, K=8192, N=4096 +dtype: torch.bfloat16 +torch: 2.12.0a0+5aff3928d8.nv26.05 +transformer_engine: 2.15.0+42b84005 +nvidia-cudnn-frontend: 1.23.0 +nvidia-cutlass-dsl: 4.4.1 + +summary: +activation cudnn_ms te_ms te/cudnn +swiglu 1.319 1.695 1.28x +dswiglu 0.740 1.549 2.09x +srelu 0.675 0.873 1.29x +dsrelu 0.686 1.163 1.70x +geglu 1.317 1.631 1.24x +dgeglu 0.751 1.666 2.22x diff --git a/benchmark/cutedsl_fusion_kernels/results/B300/output_verbose.txt b/benchmark/cutedsl_fusion_kernels/results/B300/output_verbose.txt new file mode 100644 index 000000000..b19d83aab --- /dev/null +++ b/benchmark/cutedsl_fusion_kernels/results/B300/output_verbose.txt @@ -0,0 +1,71 @@ +device: NVIDIA GB300 +shape: experts=8, tokens/expert=4096, M=32768, K=8192, N=4096 +dtype: torch.bfloat16 +torch: 2.12.0a0+5aff3928d8.nv26.05 +transformer_engine: 2.15.0+42b84005 +nvidia-cudnn-frontend: 1.23.0 +nvidia-cutlass-dsl: 4.4.1 +timing mode: CUDA graph replay + +activation: swiglu +shape: experts=8, tokens/expert=4096, M=32768, K=8192, N=4096 +fused GEMM N=8192, fused C input N=None, activation output N=4096 +fused cuDNN avg: 1.319 ms +unfused TE avg: 1.695 ms +speedup vs TE: 1.28x +fused samples: [1.233, 1.243, 1.374, 1.378, 1.369] +TE samples: [1.82, 1.612, 1.603, 1.703, 1.738] + +activation: dswiglu +shape: experts=8, tokens/expert=4096, M=32768, K=8192, N=4096 +fused GEMM N=4096, fused C input N=8192, activation output N=8192 +fused cuDNN avg: 0.740 ms +unfused TE avg: 1.549 ms +speedup vs TE: 2.09x +fused samples: [0.71, 0.716, 0.718, 0.763, 0.793] +TE samples: [1.856, 1.512, 1.464, 1.457, 1.454] + +activation: srelu +shape: experts=8, tokens/expert=4096, M=32768, K=8192, N=4096 +fused GEMM N=4096, fused C input N=None, activation output N=4096 +fused cuDNN avg: 0.675 ms +unfused TE avg: 0.873 ms +speedup vs TE: 1.29x +fused samples: [0.653, 0.654, 0.686, 0.677, 0.705] +TE samples: [0.942, 0.921, 0.85, 0.842, 0.811] + +activation: dsrelu +shape: experts=8, tokens/expert=4096, M=32768, K=8192, N=4096 +fused GEMM N=4096, fused C input N=4096, activation output N=4096 +fused cuDNN avg: 0.686 ms +unfused TE avg: 1.163 ms +speedup vs TE: 1.70x +fused samples: [0.659, 0.666, 0.677, 0.688, 0.743] +TE samples: [1.343, 1.167, 1.116, 1.096, 1.096] + +activation: geglu +shape: experts=8, tokens/expert=4096, M=32768, K=8192, N=4096 +fused GEMM N=8192, fused C input N=None, activation output N=4096 +fused cuDNN avg: 1.317 ms +unfused TE avg: 1.631 ms +speedup vs TE: 1.24x +fused samples: [1.231, 1.245, 1.344, 1.376, 1.389] +TE samples: [1.735, 1.626, 1.58, 1.65, 1.565] + +activation: dgeglu +shape: experts=8, tokens/expert=4096, M=32768, K=8192, N=4096 +fused GEMM N=4096, fused C input N=8192, activation output N=8192 +fused cuDNN avg: 0.751 ms +unfused TE avg: 1.666 ms +speedup vs TE: 2.22x +fused samples: [0.716, 0.717, 0.75, 0.774, 0.795] +TE samples: [2.058, 1.595, 1.569, 1.557, 1.553] + +summary: +activation cudnn_ms te_ms te/cudnn +swiglu 1.319 1.695 1.28x +dswiglu 0.740 1.549 2.09x +srelu 0.675 0.873 1.29x +dsrelu 0.686 1.163 1.70x +geglu 1.317 1.631 1.24x +dgeglu 0.751 1.666 2.22x diff --git a/benchmark/cutedsl_fusion_kernels/results/B300/results.png b/benchmark/cutedsl_fusion_kernels/results/B300/results.png new file mode 100644 index 000000000..178841447 Binary files /dev/null and b/benchmark/cutedsl_fusion_kernels/results/B300/results.png differ