diff --git a/.github/scripts/aiter_test.sh b/.github/scripts/aiter_test.sh index 56ece7eed5..5fae2fa2c7 100755 --- a/.github/scripts/aiter_test.sh +++ b/.github/scripts/aiter_test.sh @@ -84,9 +84,9 @@ for file in "${sharded_files[@]}"; do # batch gate so they exercise the persistent kernel at every batch size. test_cmd=(timeout 60m python3 "$file") case "$file" in - op_tests/multigpu_tests/test_mega_moe_gfx1250.py) + op_tests/multigpu_tests/bench_mega_moe.py) { - echo "Running gfx1250 MegaMoE fused-scatter accuracy on 8 GPUs when supported" + echo "Running MegaMoE fused-scatter accuracy on 8 GPUs when supported" } | tee -a latest_test.log test_cmd=( timeout 60m diff --git a/aiter/benchmark_data_init.py b/aiter/benchmark_data_init.py new file mode 100644 index 0000000000..78dfdb994c --- /dev/null +++ b/aiter/benchmark_data_init.py @@ -0,0 +1,348 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +"""Shared benchmark data and scale initialization helpers.""" + +import torch + +# --------------------------------------------------------------------------- # +# DATA / SCALE init. +# gen = make_generator(seed) +# x = fill(shape, dist, gen, dtype=...) # bf16 / fp32 / float8 +# s = fill_scale(shape, dist, gen) # float32 block scale +# xq = fill_fp4(shape, dist, gen) # MXFP4 packed e2m1 +# x8 = fill_fp8(shape, dist, gen) # MXFP8 e4m3 +# s8 = fill_scale_e8m0(shape, dist, gen) # MX E8M0 on-wire +# s4 = fill_scale_e4m3(shape, dist, gen) # NVFP4 E4M3 on-wire +# +# OCP (Open Compute Project) published the MX microscaling formats: e2m1/e4m3 +# data plus a tiny E8M0 (or E4M3) scale per block. fill_fp* emit those on-wire +# buffers. Large 2-D tensors are filled in row chunks (~1 GiB f32 staging). +# --------------------------------------------------------------------------- # +DATA_DISTS = ("zero", "constant", "uniform", "norm") +SCALE_DISTS = DATA_DISTS +SCALE_UNIFORM = (0.5, 2.0) +SCALE_NORM_MEAN, SCALE_NORM_STD = 1.0, 0.25 +FP8_E4M3 = torch.float8_e4m3fn +FP4_UNIFORM = (-3.0, 3.0) # e2m1 max is 6.0; keep headroom +FP8_UNIFORM = (-6.0, 6.0) +E8M0_BIAS = 127 +E8M0_NEUTRAL = 0x7F # 2^0 = 1.0 +E4M3_NEUTRAL = 0x38 # e4m3 exp bias -> 1.0 +E4M3_SCALE_MEAN, E4M3_SCALE_STD = 0.34375, 0.08 +POW2_BINOMIAL_N = 10 +E8M0_SCALE_DISTS = ("zero", "constant", "uniform", "norm", "auto", "pow2_binomial") +E4M3_SCALE_DISTS = ("zero", "constant", "uniform", "norm", "auto") +_STAGE_ELEMS = 1 << 28 # 256M f32 = 1 GiB per chunk + + +def make_generator(seed, device="cuda"): + """Seeded ``torch.Generator`` -- same seed => bit-identical buffers.""" + return torch.Generator(device=device).manual_seed(int(seed)) + + +def add_data_init_args( + parser, + *, + default_dist="uniform", + default_scale="constant", + default_seed=0, + include_scale=True, +): + """Attach shared data initialization arguments to a benchmark parser.""" + parser.add_argument( + "--data-init", + dest="data_init", + nargs="+", + choices=list(DATA_DISTS), + default=[default_dist], + help="DATA init: zero | constant | uniform | norm (N(0,1)). " + "e.g.: --data-init uniform norm", + ) + if include_scale: + parser.add_argument( + "--scale-init", + dest="scale_init", + nargs="+", + choices=list(SCALE_DISTS), + default=[default_scale], + help="SCALE init (non-negative float): zero | constant(=1) | " + "uniform U(0.5,2) | norm N(1,0.25). Independent of --data-init.", + ) + parser.add_argument( + "--seed", + type=int, + default=default_seed, + help="RNG seed; same seed -> bit-identical uniform/norm buffers", + ) + return parser + + +def _row_chunks(rows, cols): + """Row slices whose f32 staging stays around _STAGE_ELEMS elements.""" + step = max(_STAGE_ELEMS // max(cols, 1), 1) + for start in range(0, rows, step): + yield start, min(start + step, rows) + + +def _canon_dist(dist, allowed): + if dist == "gaussian": + dist = "norm" + if dist not in allowed: + raise ValueError(f"dist {dist!r}; choose from {allowed}") + return dist + + +def _sample_data_f32(shape, dist, gen, *, lo, hi, device): + if dist == "uniform": + return torch.empty(shape, dtype=torch.float32, device=device).uniform_( + lo, hi, generator=gen + ) + if dist == "norm": + return torch.empty(shape, dtype=torch.float32, device=device).normal_( + 0.0, 1.0, generator=gen + ) + raise ValueError(f"data dist {dist!r} is not continuous; use fill dispatch") + + +def _sample_scale_f32(shape, dist, gen, *, lo, hi, device): + if dist == "uniform": + v = torch.empty(shape, dtype=torch.float32, device=device).uniform_( + lo, hi, generator=gen + ) + elif dist == "norm": + v = torch.empty(shape, dtype=torch.float32, device=device).normal_( + SCALE_NORM_MEAN, SCALE_NORM_STD, generator=gen + ) + else: + raise ValueError(f"scale dist {dist!r} is not continuous; use fill_scale") + v.clamp_(min=0.0) + return v + + +def _fill_sampled(shape, dist, gen, *, dtype, device, uniform, constant, sample_fn): + if dist == "zero": + return torch.zeros(shape, dtype=dtype, device=device) + if dist == "constant": + return torch.full(shape, constant, dtype=dtype, device=device) + lo, hi = uniform + if len(shape) != 2: + return sample_fn(shape, dist, gen, lo=lo, hi=hi, device=device).to(dtype) + rows, cols = shape + out = torch.empty(shape, dtype=dtype, device=device) + for r0, r1 in _row_chunks(rows, cols): + v = sample_fn((r1 - r0, cols), dist, gen, lo=lo, hi=hi, device=device) + out[r0:r1] = v.to(dtype) + del v + return out + + +def fill( + shape, + dist, + gen, + *, + dtype=torch.float32, + device="cuda", + uniform=(-1.0, 1.0), + constant=1.0, +): + """Return a ``dtype`` DATA tensor of ``shape``. + + ``dist`` in {zero, constant, uniform, norm}. ``uniform`` is U(lo, hi); + ``norm`` / ``gaussian`` is N(0, 1). ``zero`` / ``constant`` ignore ``gen``. + """ + dist = _canon_dist(dist, DATA_DISTS) + return _fill_sampled( + shape, + dist, + gen, + dtype=dtype, + device=device, + uniform=uniform, + constant=constant, + sample_fn=_sample_data_f32, + ) + + +def fill_scale( + shape, + dist, + gen, + *, + dtype=torch.float32, + device="cuda", + uniform=SCALE_UNIFORM, + constant=1.0, +): + """Return a non-negative float SCALE tensor of ``shape``. + + Same dist names as ``fill``, sampled independently. ``constant`` defaults + to 1.0 (neutral). ``norm`` is N(1, 0.25) clamped >= 0 -- not DATA's N(0,1). + For MX on-wire scales use ``fill_scale_e8m0`` / ``fill_scale_e4m3``. + """ + dist = _canon_dist(dist, SCALE_DISTS) + return _fill_sampled( + shape, + dist, + gen, + dtype=dtype, + device=device, + uniform=uniform, + constant=constant, + sample_fn=_sample_scale_f32, + ) + + +def _f32_to_e8m0(v: torch.Tensor) -> torch.Tensor: + """Round positive floats to the nearest E8M0 on-wire byte (bias 127).""" + e = torch.zeros_like(v, dtype=torch.int32) + pos = v > 0 + e[pos] = v[pos].log2().round().to(torch.int32) + E8M0_BIAS + return e.clamp_(0, 255).to(torch.uint8) + + +def _popcount64(x: torch.Tensor) -> torch.Tensor: + """Population count for a non-negative int64 tensor (SWAR bit-hack).""" + x = x - ((x >> 1) & 0x5555555555555555) + x = (x & 0x3333333333333333) + ((x >> 2) & 0x3333333333333333) + x = (x + (x >> 4)) & 0x0F0F0F0F0F0F0F0F + return (x * 0x0101010101010101) >> 56 + + +def fill_fp4(shape, dist, gen, *, uniform=FP4_UNIFORM, device="cuda", constant=0): + """MXFP4 on-wire: packed e2m1 ``uint8`` of shape ``(rows, cols // 2)``. + + Samples with the same DATA dists as ``fill``, then round-to-nearest e2m1. + ``shape`` is the logical ``(rows, cols)``; ``cols`` must be even. + """ + dist = _canon_dist(dist, DATA_DISTS) + rows, cols = shape + assert cols % 2 == 0, f"FP4 needs even columns, got {cols}" + packed = (rows, cols // 2) + if dist == "zero": + return torch.zeros(packed, dtype=torch.uint8, device=device) + if dist == "constant": + return torch.full(packed, int(constant), dtype=torch.uint8, device=device) + + from aiter.utility import fp4_utils # local: fp4_utils pulls in triton + + out = torch.empty(packed, dtype=torch.uint8, device=device) + for r0, r1 in _row_chunks(rows, cols): + v = _sample_data_f32( + (r1 - r0, cols), + dist, + gen, + lo=uniform[0], + hi=uniform[1], + device=device, + ) + out[r0:r1] = fp4_utils.f32_to_mxfp4(v).view(torch.uint8) + del v + return out + + +def fill_fp8( + shape, + dist, + gen, + *, + dtype=FP8_E4M3, + uniform=FP8_UNIFORM, + device="cuda", + constant=0.5, +): + """MXFP8 on-wire: e4m3 tensor of ``shape``.""" + dist = _canon_dist(dist, DATA_DISTS) + if dist == "zero": + return torch.zeros(shape, dtype=dtype, device=device) + if dist == "constant": + return torch.full( + shape, float(constant), dtype=torch.float32, device=device + ).to(dtype) + if len(shape) != 2: + v = _sample_data_f32( + shape, + dist, + gen, + lo=uniform[0], + hi=uniform[1], + device=device, + ) + return v.to(dtype) + rows, cols = shape + out = torch.empty(shape, dtype=dtype, device=device) + for r0, r1 in _row_chunks(rows, cols): + v = _sample_data_f32( + (r1 - r0, cols), + dist, + gen, + lo=uniform[0], + hi=uniform[1], + device=device, + ) + out[r0:r1] = v.to(dtype) + del v + return out + + +def fill_scale_e8m0( + shape, + dist="auto", + gen=None, + *, + device="cuda", + n=POW2_BINOMIAL_N, + constant=E8M0_NEUTRAL, +): + """MX E8M0 on-wire ``uint8`` (biased exponent, bias 127). + + ``zero`` / ``constant`` / ``uniform`` / ``norm`` map from our SCALE dists + (float then round to nearest power-of-two byte). ``auto`` / + ``pow2_binomial`` match the MX GEMM default: 2^(Binomial(21,0.5)-11). + """ + if dist == "gaussian": + dist = "norm" + if dist not in E8M0_SCALE_DISTS: + raise ValueError(f"E8M0 scale dist {dist!r}; choose from {E8M0_SCALE_DISTS}") + if dist == "zero": + return torch.zeros(shape, dtype=torch.uint8, device=device) + if dist == "constant": + return torch.full(shape, int(constant), dtype=torch.uint8, device=device) + if dist in ("uniform", "norm"): + v = fill_scale(shape, dist, gen, device=device) + return _f32_to_e8m0(v) + # auto / pow2_binomial: Binomial(k, 0.5) == popcount of a uniform k-bit int + trials = 2 * n + 1 + assert trials <= 24, "pow2_binomial popcount path assumes <= 24 trials" + bits = torch.randint( + 0, 1 << trials, shape, dtype=torch.int64, device=device, generator=gen + ) + e = _popcount64(bits).to(torch.int32) - (n + 1) + return (e + E8M0_BIAS).clamp_(0, 255).to(torch.uint8) + + +def fill_scale_e4m3( + shape, dist="auto", gen=None, *, device="cuda", constant=E4M3_NEUTRAL +): + """NVFP4 / E4M3 on-wire ``uint8``. + + ``auto`` -> N(0.34375, 0.08) clamped >= 0, then cast e4m3 (MX GEMM default). + ``uniform`` / ``norm`` use ``fill_scale`` then cast. ``constant`` is 0x38 + (1.0). + """ + if dist == "gaussian": + dist = "auto" + if dist not in E4M3_SCALE_DISTS: + raise ValueError(f"E4M3 scale dist {dist!r}; choose from {E4M3_SCALE_DISTS}") + if dist == "zero": + return torch.zeros(shape, dtype=torch.uint8, device=device) + if dist == "constant": + return torch.full(shape, int(constant), dtype=torch.uint8, device=device) + if dist in ("uniform", "norm"): + v = fill_scale(shape, dist, gen, device=device) + return v.to(FP8_E4M3).view(torch.uint8) + v = torch.empty(shape, dtype=torch.float32, device=device).normal_( + E4M3_SCALE_MEAN, E4M3_SCALE_STD, generator=gen + ) + v.clamp_(min=0.0) + return v.to(FP8_E4M3).view(torch.uint8) diff --git a/aiter/benchmark_reporting.py b/aiter/benchmark_reporting.py new file mode 100644 index 0000000000..69d22c65a3 --- /dev/null +++ b/aiter/benchmark_reporting.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +"""Structured output helpers for benchmark drivers.""" + +import json + +import pandas as pd + + +def _json_default(value): + """Render non-native JSON scalars without expanding their attributes.""" + return str(value).removeprefix("torch.") + + +def print_json_table(name, rows, keep=None): + """Print benchmark rows as one record-oriented JSON object. + + A single-line object is intentional: parent benchmark drivers can validate + and forward it without parsing pandas' human-readable table formats. + """ + if isinstance(rows, pd.DataFrame): + df = rows.copy() + else: + df = pd.DataFrame([row for row in rows if row is not None]) + if not df.empty: + df = df.replace("", pd.NA).dropna(axis=1, how="all") + if keep is not None: + cols = [column for column in keep if column in df.columns] + cols += [ + column + for column in df.columns + if "err_msg" in column and column not in cols + ] + df = df[cols] + records = json.loads(df.to_json(orient="records", default_handler=_json_default)) + print(json.dumps({"name": name, "rows": records}), flush=True) diff --git a/aiter/smi_monitor.py b/aiter/smi_monitor.py new file mode 100644 index 0000000000..0de583f61c --- /dev/null +++ b/aiter/smi_monitor.py @@ -0,0 +1,564 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +# ruff: noqa: BLE001, PYI034, S110, UP035, UP037 +"""AMD GPU metrics monitor using amdsmi. + +ROCm ships the binding without a setup.py. If the normal import fails, this +module temporarily searches ``/opt/rocm/share/amd_smi`` while importing it. +Neither ``PYTHONPATH`` nor the caller's lasting ``sys.path`` is changed. +The amdsmi session and HIP-device handle mapping are initialized lazily once +per process and reused by every monitor instance. + +Usage (context manager): + with GpuMonitor(device_index=0, interval_s=0.05) as mon: + run_workload() + samples = mon.samples # list[dict] + +Usage (explicit start/stop): + mon = GpuMonitor(device_index=0, interval_s=0.05) + mon.start() + run_workload() + mon.stop() + samples = mon.samples +""" + +from __future__ import annotations + +import atexit +import ctypes +import importlib +import inspect +import json +import os +import sys +import threading +import time +from contextlib import contextmanager +from contextvars import ContextVar +from enum import Enum +from functools import cache +from typing import Generator + +import numpy as np +import torch + +SMI_RESULT_PREFIX = "AITER_SMI_RESULT " +_ROCM_AMDSMI_PATH = "/opt/rocm/share/amd_smi" +_SMI_LABEL_COUNTS = {} +_SMI_CALL_LABEL = ContextVar("aiter_smi_call_label", default=None) + + +def _smi_label_value(value): + """Return a compact, stable label value, or None for opaque arguments.""" + if isinstance(value, torch.Tensor): + shape = "x".join(map(str, value.shape)) or "scalar" + return f"{shape}:{str(value.dtype).removeprefix('torch.')}" + if isinstance(value, torch.dtype): + return str(value).removeprefix("torch.") + if isinstance(value, Enum): + return str(value.value) + if value is None or isinstance(value, (str, bool, int, float, np.generic)): + return str(value) + if isinstance(value, (tuple, list)): + items = [_smi_label_value(item) for item in value] + if all(item is not None for item in items): + return ",".join(items) + return None + + +def _smi_call_tag(func, callargs): + """Build a call-local SMI label from @benchmark's named arguments.""" + source = os.path.splitext(os.path.basename(func.__code__.co_filename))[0] + parts = [f"{source}.{func.__name__}"] + aliases = {"m": "M", "n": "N", "k": "K", "t": "T", "h": "H", "d": "D"} + for name, value in callargs.items(): + formatted = _smi_label_value(value) + if formatted is None: + continue + formatted = formatted.replace("/", "_").replace("\n", "") + parts.append(f"{aliases.get(name, name)}={formatted}") + return "/".join(parts) + + +def _smi_perftest_tag(func, args, kwargs): + """Return scalar call details that distinguish one perftest invocation.""" + try: + signature = inspect.signature(func) + callargs = signature.bind(*args, **kwargs) + callargs.apply_defaults() + except (TypeError, ValueError): + return None + + parts = [] + for name, parameter in signature.parameters.items(): + if parameter.kind in ( + inspect.Parameter.VAR_POSITIONAL, + inspect.Parameter.VAR_KEYWORD, + ): + continue + value = callargs.arguments.get(name) + if value is None or isinstance(value, torch.Tensor) or callable(value): + continue + formatted = _smi_label_value(value) + if formatted is None: + continue + formatted = formatted.replace("/", "_").replace("\n", "") + parts.append(f"{name}={formatted}") + return "/".join(parts) or None + + +@contextmanager +def benchmark_call_context(func, callargs): + """Expose one benchmark call's metadata to an inner SMI replay.""" + token = _SMI_CALL_LABEL.set(_smi_call_tag(func, callargs)) + try: + yield + finally: + _SMI_CALL_LABEL.reset(token) + + +def replay_with_smi_metadata( + func, + args, + kwargs, + replay, + *, + synchronize, + estimated_us: float | None = None, +): + """Replay one callable with stable benchmark metadata and skip filtering.""" + if not smi_replay_enabled(): + return None + + fn_name = getattr(func, "__name__", "kernel") + skipped = { + name.strip() + for name in os.environ.get("AITER_SMI_SKIP_FUNCTIONS", "").split(",") + if name.strip() + } + if fn_name in skipped: + return None + + case_label = _SMI_CALL_LABEL.get() or os.environ.get( + "AITER_SMI_LABEL", "benchmark_case" + ) + perftest_tag = _smi_perftest_tag(func, args, kwargs) + if perftest_tag: + case_label = f"{case_label}/{perftest_tag}" + label_key = (case_label, fn_name) + occurrence = _SMI_LABEL_COUNTS.get(label_key, 0) + 1 + _SMI_LABEL_COUNTS[label_key] = occurrence + return replay_with_smi( + replay, + label=f"{case_label}/{fn_name}#{occurrence}", + synchronize=synchronize, + estimated_us=estimated_us, + ) + + +def _import_amdsmi(): + """Import amdsmi, temporarily searching ROCm's unpackaged binding.""" + try: + return importlib.import_module("amdsmi") + except ImportError: + if not os.path.isdir(_ROCM_AMDSMI_PATH): + raise + + added_path = _ROCM_AMDSMI_PATH not in sys.path + if added_path: + sys.path.insert(0, _ROCM_AMDSMI_PATH) + try: + return importlib.import_module("amdsmi") + finally: + if added_path: + sys.path.remove(_ROCM_AMDSMI_PATH) + + +try: + amdsmi = _import_amdsmi() + + _AMDSMI_AVAILABLE = True +except ImportError: + _AMDSMI_AVAILABLE = False + + +_AMDSMI_LOCK = threading.Lock() +_AMDSMI_INITIALIZED = False +_AMDSMI_HANDLES_BY_BDF = {} +_AMDSMI_HANDLES_BY_HIP_DEVICE = {} + + +# ------------------------------------------------------------------ +# HIP device -> amdsmi handle via PCIe BDF +# ------------------------------------------------------------------ + + +@cache +def _hip_runtime_library(): + """Load and cache the HIP runtime used for device-to-BDF lookup.""" + import glob as _glob + + candidates = ["libamdhip64.so"] + sorted( + _glob.glob("/opt/rocm/lib/libamdhip64.so.*"), reverse=True + ) + for name in candidates: + try: + return ctypes.CDLL(name) + except OSError: + continue + raise RuntimeError( + "libamdhip64.so not found (tried unversioned and /opt/rocm/lib/libamdhip64.so.*); " + "is ROCm installed?" + ) + + +@cache +def _hip_device_bdf(hip_device: int) -> str: + """Return the PCIe BDF string for a HIP device index, e.g. '0000:03:00.0'. + + Calls ``hipDeviceGetPCIBusId`` via ctypes so there is no hard dependency on + PyTorch or the hip-python package. + """ + libhip = _hip_runtime_library() + buf = ctypes.create_string_buffer(64) + ret = libhip.hipDeviceGetPCIBusId(buf, ctypes.c_int(64), ctypes.c_int(hip_device)) + if ret != 0: + raise RuntimeError(f"hipDeviceGetPCIBusId failed with error code {ret}") + return buf.value.decode().lower().strip() + + +def _amdsmi_bdf_str(handle) -> str: + """Normalise the BDF returned by amdsmi into 'dddd:bb:dd.f' lowercase.""" + raw = amdsmi.amdsmi_get_gpu_device_bdf(handle) + if isinstance(raw, str): + return raw.lower().strip() + # Some amdsmi versions return a dict: {'domain': 0, 'bus': 3, 'device': 0, 'function': 0} + return ( + f"{raw['domain']:04x}:{raw['bus']:02x}:{raw['device']:02x}.{raw['function']:x}" + ) + + +def _ensure_amdsmi_initialized() -> None: + """Initialize amdsmi and enumerate processor handles once per process.""" + global _AMDSMI_INITIALIZED, _AMDSMI_HANDLES_BY_BDF + + if not _AMDSMI_AVAILABLE: + raise ImportError("amdsmi is not installed or not importable") + with _AMDSMI_LOCK: + if _AMDSMI_INITIALIZED: + return + amdsmi.amdsmi_init() + try: + _AMDSMI_HANDLES_BY_BDF = { + _amdsmi_bdf_str(handle): handle + for handle in amdsmi.amdsmi_get_processor_handles() + } + except BaseException: + amdsmi.amdsmi_shut_down() + raise + _AMDSMI_INITIALIZED = True + + +def _shutdown_amdsmi() -> None: + """Release the process-wide amdsmi session during interpreter shutdown.""" + global _AMDSMI_INITIALIZED + + with _AMDSMI_LOCK: + if not _AMDSMI_INITIALIZED: + return + try: + amdsmi.amdsmi_shut_down() + finally: + _AMDSMI_INITIALIZED = False + _AMDSMI_HANDLES_BY_BDF.clear() + _AMDSMI_HANDLES_BY_HIP_DEVICE.clear() + + +atexit.register(_shutdown_amdsmi) + + +def hip_device_to_amdsmi_handle(hip_device: int): + """Return the amdsmi processor handle that corresponds to a HIP device index. + + Uses PCIe BDF as the stable identifier linking the two numbering schemes. + + Args: + hip_device: HIP device ordinal (as used by ``torch.cuda`` / HIP runtime). + + Returns: + The amdsmi processor handle for that GPU. + + Raises: + RuntimeError: if no amdsmi handle matches the HIP device's BDF. + ImportError: if amdsmi is not available. + """ + _ensure_amdsmi_initialized() + with _AMDSMI_LOCK: + cached = _AMDSMI_HANDLES_BY_HIP_DEVICE.get(hip_device) + if cached is not None: + return cached + target_bdf = _hip_device_bdf(hip_device) + handle = _AMDSMI_HANDLES_BY_BDF.get(target_bdf) + if handle is None: + raise RuntimeError( + f"No amdsmi handle found with BDF {target_bdf!r} " + f"(HIP device {hip_device})" + ) + _AMDSMI_HANDLES_BY_HIP_DEVICE[hip_device] = handle + return handle + + +def _collect_sample(handle) -> dict: + """Collect one snapshot from a single GPU handle.""" + sample: dict = {"timestamp_s": time.perf_counter()} + try: + metrics = amdsmi.amdsmi_get_gpu_metrics_info(handle) + sample["gfx_clk_mhz"] = metrics.get("current_gfxclk", None) + sample["soc_clk_mhz"] = metrics.get("current_socclk", None) + sample["power_w"] = metrics.get("current_socket_power", None) + sample["temp_hotspot_c"] = metrics.get("temperature_hotspot", None) + except Exception: + pass + try: + info = amdsmi.amdsmi_get_gpu_activity(handle) + sample["gfx_activity_pct"] = info.get("gfx_activity", None) + sample["umc_activity_pct"] = info.get("umc_activity", None) + except Exception: + pass + try: + mem = amdsmi.amdsmi_get_gpu_memory_usage(handle, amdsmi.AmdSmiMemoryType.VRAM) + sample["vram_used_mb"] = mem / 1024 / 1024 + except Exception: + pass + return sample + + +class GpuMonitor: + """Poll AMD GPU metrics on a background thread. + + Args: + device_index: Integer ordinal of the GPU to monitor (default 0), or a + pre-resolved amdsmi processor handle (e.g. from + ``hip_device_to_amdsmi_handle``). + interval_s: Polling interval in seconds (default 0.05 = 50 ms). + """ + + def __init__(self, device_index: int = 0, interval_s: float = 0.05) -> None: + if not _AMDSMI_AVAILABLE: + raise ImportError("amdsmi is not installed or not importable") + self._device_index = device_index + self._interval_s = interval_s + self._samples: list[dict] = [] + self._thread: threading.Thread | None = None + self._stop_event = threading.Event() + self._ready_event = threading.Event() + self._error: BaseException | None = None + self._handle = None + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def start(self) -> None: + """Begin background polling. Safe to call only once per instance.""" + if self._thread is not None and self._thread.is_alive(): + raise RuntimeError("GpuMonitor is already running") + self._samples = [] + self._error = None + self._stop_event.clear() + self._ready_event.clear() + try: + if isinstance(self._device_index, int): + self._handle = hip_device_to_amdsmi_handle(self._device_index) + else: + _ensure_amdsmi_initialized() + self._handle = self._device_index + except BaseException as error: + raise RuntimeError( + f"failed to initialize amdsmi monitor: {error}" + ) from error + self._thread = threading.Thread(target=self._poll_loop, daemon=True) + self._thread.start() + if not self._ready_event.wait(timeout=10.0): + self._stop_event.set() + raise RuntimeError("timed out while initializing amdsmi monitor") + if self._error is not None: + error = self._error + self.stop() + raise RuntimeError(f"failed to initialize amdsmi monitor: {error}") + + def stop(self) -> None: + """Stop background polling and wait for the thread to finish.""" + self._stop_event.set() + if self._thread is not None: + self._thread.join() + self._thread = None + + @property + def samples(self) -> list[dict]: + """Collected samples; each is a dict with 'timestamp_s' plus metric keys.""" + return list(self._samples) + + def summary( + self, *, start_s: float | None = None, end_s: float | None = None + ) -> dict: + """Return metric summaries, optionally restricted to a timestamp window.""" + samples = [ + sample + for sample in self._samples + if (start_s is None or sample["timestamp_s"] >= start_s) + and (end_s is None or sample["timestamp_s"] <= end_s) + ] + if not samples: + return {} + keys = {key for sample in samples for key in sample if key != "timestamp_s"} + result: dict = {} + for key in sorted(keys): + vals = sorted( + s[key] for s in samples if s.get(key) is not None and s[key] != "N/A" + ) + if not vals: + continue + n = len(vals) + mid = n // 2 + median = vals[mid] if n % 2 else (vals[mid - 1] + vals[mid]) / 2 + result[key] = { + "min": vals[0], + "mean": sum(vals) / n, + "median": median, + "max": vals[-1], + "n": n, + } + return result + + # ------------------------------------------------------------------ + # Context manager + # ------------------------------------------------------------------ + + def __enter__(self) -> "GpuMonitor": + self.start() + return self + + def __exit__(self, *_) -> None: + self.stop() + + # ------------------------------------------------------------------ + # Internal + # ------------------------------------------------------------------ + + def _poll_loop(self) -> None: + try: + self._ready_event.set() + while not self._stop_event.is_set(): + t0 = time.perf_counter() + self._samples.append(_collect_sample(self._handle)) + elapsed = time.perf_counter() - t0 + remaining = self._interval_s - elapsed + if remaining > 0: + self._stop_event.wait(timeout=remaining) + except BaseException as error: + self._error = error + self._ready_event.set() + finally: + self._ready_event.set() + + +# ------------------------------------------------------------------ +# Module-level convenience functions +# ------------------------------------------------------------------ + + +@contextmanager +def monitor_gpu( + device_index: int = 0, interval_s: float = 0.05 +) -> Generator[GpuMonitor, None, None]: + """Context manager that yields a running GpuMonitor. + + Example:: + + with monitor_gpu(device_index=0) as mon: + run_workload() + print(mon.summary()) + """ + mon = GpuMonitor(device_index=device_index, interval_s=interval_s) + with mon: + yield mon + + +def smi_replay_enabled() -> bool: + """Whether the benchmark requested an isolated SMI replay window.""" + return os.environ.get("AITER_SMI_MONITOR", "0") == "1" + + +def emit_smi_result(result: dict) -> None: + """Write one structured result to the combo JSONL sink or stdout.""" + line = SMI_RESULT_PREFIX + json.dumps(result, sort_keys=True) + output_path = os.environ.get("AITER_SMI_OUTPUT_PATH") + if output_path: + with open(output_path, "a", encoding="utf-8") as output: + output.write(line + "\n") + else: + print(line, flush=True) + + +def replay_with_smi( + fn, + *, + label: str, + synchronize, + estimated_us: float | None = None, +) -> dict | None: + """Repeat one already-prepared benchmark case under the GPU monitor. + + Input creation, compilation, correctness and the latency measurement happen + before this function is called. Batching launches between synchronizations + keeps short kernels busy while still checking the wall-clock deadline often + enough for slow kernels. + """ + if not smi_replay_enabled(): + return None + + device = int(os.environ.get("AITER_SMI_DEVICE", "0")) + interval_s = float(os.environ.get("AITER_SMI_INTERVAL", "0.05")) + duration_s = float(os.environ.get("AITER_SMI_DURATION", "1.0")) + if interval_s <= 0 or duration_s <= 0: + raise ValueError("AITER_SMI_INTERVAL and AITER_SMI_DURATION must be positive") + + # Aim for roughly one synchronization per monitor tick. The cap prevents a + # near-zero/invalid latency estimate from enqueueing an unbounded amount of + # work, while a slow case is synchronized after every launch. + if estimated_us is not None and estimated_us > 0: + batch_iters = max(1, min(1024, int(interval_s * 1e6 / estimated_us))) + else: + batch_iters = 1 + + synchronize() + launches = 0 + start = time.perf_counter() + with monitor_gpu(device_index=device, interval_s=interval_s) as monitor: + while launches == 0 or time.perf_counter() - start < duration_s: + for _ in range(batch_iters): + fn() + launches += batch_iters + synchronize() + elapsed_s = time.perf_counter() - start + + result = { + "label": label, + "device": device, + "interval_s": interval_s, + "duration_s": elapsed_s, + "launches": launches, + "samples": len(monitor.samples), + "metrics": monitor.summary(), + } + expected_samples = max(1, int(duration_s / interval_s)) + result["sample_status"] = ( + "ok" + if len(monitor.samples) >= max(2, expected_samples // 2) + else "insufficient" + ) + # A shared JSONL sink survives fd silencing and child processes. Standalone + # UT runs without a sink still get a machine-readable stdout record. + emit_smi_result(result) + return result diff --git a/aiter/test_common.py b/aiter/test_common.py index 7b717d39c0..ecd8fd63ed 100644 --- a/aiter/test_common.py +++ b/aiter/test_common.py @@ -3,6 +3,7 @@ import copy import multiprocessing as mp import os +from functools import wraps import numpy as np import pandas as pd @@ -87,7 +88,6 @@ def wrapper(*args, **kwargs): end_event.record() end_event.synchronize() latencies.append(start_event.elapsed_time(end_event)) - torch.cuda.empty_cache() avg = np.mean(latencies) * 1000 logger.info(f"avg: {avg} us/iter from cuda.Event") if use_cuda_event: @@ -124,6 +124,36 @@ def wrapper(*args, **kwargs): avg = get_trace_perf(prof, num_iters) logger.info(f"avg: {avg} us/iter with hipgraph") + if os.environ.get("AITER_SMI_MONITOR", "0") == "1": + # Import lazily: normal library/test use has no amdsmi dependency. + from aiter.smi_monitor import replay_with_smi_metadata + + if testGraph: + replay = graph.replay + # One replay contains num_iters calls captured above. + replay_us = avg * num_iters + else: + replay_index = 0 + + def replay(): + nonlocal replay_index + replay_args, replay_kwargs = rotate_args[ + replay_index % len(rotate_args) + ] + replay_index += 1 + return func(*replay_args, **replay_kwargs) + + replay_us = avg + + replay_with_smi_metadata( + func, + args, + kwargs, + replay, + synchronize=torch.cuda.synchronize, + estimated_us=replay_us, + ) + return data, avg return wrapper @@ -135,7 +165,13 @@ def benchmark(): def decorator(func): def wrapper(*args, **kwargs): callargs = log_args(func, *args, **kwargs) - ret = func(*args, **kwargs) + if os.environ.get("AITER_SMI_MONITOR", "0") == "1": + from aiter.smi_monitor import benchmark_call_context + + with benchmark_call_context(func, callargs): + ret = func(*args, **kwargs) + else: + ret = func(*args, **kwargs) if ret is not None: callargs.update(ret) return callargs @@ -222,6 +258,7 @@ def run_perftest( needTrace=needTrace, use_cuda_event=use_cuda_event, ) + @wraps(func) def worker(*args, **kwargs): return func(*args, **kwargs) @@ -401,9 +438,11 @@ def get_trace_perf(prof, num_iters): df.at[avg_name, el] = df[el].sum() / actual_iters if int(os.environ.get("AITER_LOG_MORE", "0")): pd.set_option("display.expand_frame_repr", False) - pd.set_option("display.max_colwidth", 90) pd.set_option("display.float_format", "{:,.1f}".format) - logger.info(f"{df}") + # ``name`` is the only potentially long text column in this profiler + # table. Keep its full kernel symbol for downstream log parsers without + # changing pandas' process-wide column-width setting. + logger.info(df.to_string(max_colwidth=None)) return df.at[avg_name, "device_time_sum"] diff --git a/op_tests/test_flydsl_grouped_gemm_gfx1250.py b/op_tests/flydsl_tests/test_flydsl_grouped_gemm.py similarity index 77% rename from op_tests/test_flydsl_grouped_gemm_gfx1250.py rename to op_tests/flydsl_tests/test_flydsl_grouped_gemm.py index 8ba921cfbe..04f093f8f7 100644 --- a/op_tests/test_flydsl_grouped_gemm_gfx1250.py +++ b/op_tests/flydsl_tests/test_flydsl_grouped_gemm.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python3 - # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. @@ -15,7 +13,7 @@ ``AITER_USE_GROUPED_GEMM=1`` env (set automatically by the runner below). Pytest covers a small correctness case for each format. Direct execution -(``python op_tests/test_flydsl_grouped_gemm_gfx1250.py``) runs a +(``python op_tests/flydsl_tests/test_flydsl_grouped_gemm.py``) runs a DeepSeek-style perf bench (``--scenario bench``, end-to-end fused_moe), a per-kernel bench that times gemm1 and gemm2 in isolation (``--scenario kernel``), a tiny correctness check @@ -36,6 +34,7 @@ import torch from aiter import ActivationType, QuantType, logger +from aiter import benchmark_data_init as bench_init from aiter.aot.flydsl.common import run_only_env from aiter.fused_moe import ( fused_moe, @@ -239,23 +238,49 @@ def _per_1x32_fp8_dequant(x: torch.Tensor) -> torch.Tensor: # Mock data builders # --------------------------------------------------------------------------- def _pattern_packed( - experts: int, rows: int, k_pack: int, *, const_init: float | None = None + experts: int, + rows: int, + k_pack: int, + *, + data_init: str, + generator: torch.Generator, ) -> torch.Tensor: - """mxfp4 packed bytes ``(E, rows, k_pack) uint8`` from the global RNG.""" - if const_init is not None: - return torch.full((experts, rows, k_pack), int(const_init), dtype=torch.uint8) - return torch.randint(0, 256, (experts, rows, k_pack), dtype=torch.uint8) + """Build packed MXFP4 weights with the shared benchmark initializer.""" + if data_init == "constant": + return torch.full((experts, rows, k_pack), 0x11, dtype=torch.uint8) + packed = bench_init.fill_fp4((experts * rows, k_pack * 2), data_init, generator) + return packed.view(experts, rows, k_pack) def init_weight_scales( - experts: int, rows: int, n_blocks: int, *, const_init: float | None = None + experts: int, + rows: int, + n_blocks: int, + *, + scale_init: str, + generator: torch.Generator, +) -> torch.Tensor: + """Build E8M0 weight scales with the shared benchmark initializer.""" + if scale_init == "constant": + return torch.full( + (experts, rows, n_blocks), DEFAULT_SCALE_BYTE, dtype=torch.uint8 + ) + return bench_init.fill_scale_e8m0((experts, rows, n_blocks), scale_init, generator) + + +def _init_hidden( + shape: tuple[int, int], + data_format: str, + data_init: str, + generator: torch.Generator, ) -> torch.Tensor: - """Per-block e8m0 weight scale: random small scales (drawn from the global - RNG) so the n32k4 B-scale preshuffle layout is actually exercised.""" - if const_init is not None: - return torch.full((experts, rows, n_blocks), int(const_init), dtype=torch.uint8) - r = torch.randint(0, 3, (experts, rows, n_blocks), dtype=torch.int16) - return (r + (DEFAULT_SCALE_BYTE - 1)).to(torch.uint8) + """Build BF16 activations using the selected low-precision data model.""" + if data_init == "constant": + return torch.full(shape, 0.5, dtype=torch.bfloat16) + if data_format == "a4w4": + packed = bench_init.fill_fp4(shape, data_init, generator) + return fp4_utils.mxfp4_to_f32(packed).to(torch.bfloat16) + return bench_init.fill_fp8(shape, data_init, generator).to(torch.bfloat16) def _make_routing_score(tokens: int, experts: int, topk: int) -> torch.Tensor: @@ -332,7 +357,8 @@ def _run_grouped_via_fused_moe( seed: int = 0, warmup: int = 5, iters: int = 101, - const_init: float | None = None, + data_init: str = "uniform", + scale_init: str = "auto", ) -> tuple[torch.Tensor, torch.Tensor, float | None, dict | None]: """Build mxfp4 weights + routing, dispatch through ``fused_moe``. @@ -358,31 +384,50 @@ def _run_grouped_via_fused_moe( inter_pack = inter // 2 # Logical weights/scale/bias: always GGUU (gate rows then up rows). - # One global seed per case; every draw below uses the global RNG. torch.manual_seed(seed) - w1_logical = _pattern_packed(experts, 2 * inter, K_pack, const_init=const_init) - w2_logical = _pattern_packed(experts, K, inter_pack, const_init=const_init) + generator = bench_init.make_generator(seed) + w1_logical = _pattern_packed( + experts, + 2 * inter, + K_pack, + data_init=data_init, + generator=generator, + ) + w2_logical = _pattern_packed( + experts, + K, + inter_pack, + data_init=data_init, + generator=generator, + ) w1_scale_raw = init_weight_scales( - experts, 2 * inter, K // SCALE_BLOCK, const_init=const_init + experts, + 2 * inter, + K // SCALE_BLOCK, + scale_init=scale_init, + generator=generator, ) w2_scale_raw = init_weight_scales( - experts, K, inter // SCALE_BLOCK, const_init=const_init + experts, + K, + inter // SCALE_BLOCK, + scale_init=scale_init, + generator=generator, ) if use_bias: - if const_init is not None: - bias1 = torch.full((experts, 2 * inter), float(const_init)) - bias2 = torch.full((experts, K), float(const_init)) + if data_init == "constant": + bias1 = torch.full((experts, 2 * inter), 0.5) + bias2 = torch.full((experts, K), 0.5) else: - bias1 = (torch.randn((experts, 2 * inter)) * 1e-3).float() - bias2 = (torch.randn((experts, K)) * 1e-3).float() + bias1 = ( + torch.randn((experts, 2 * inter), generator=generator) * 1e-3 + ).float() + bias2 = (torch.randn((experts, K), generator=generator) * 1e-3).float() else: bias1 = torch.zeros((experts, 2 * inter)) bias2 = torch.zeros((experts, K)) # Activations: bf16; fused_moe handles the dispatched quant internally. - if const_init is not None: - hidden = torch.full((tokens, K), float(const_init), dtype=torch.bfloat16) - else: - hidden = (torch.randn((tokens, K)) * 0.5).to(torch.bfloat16) + hidden = _init_hidden((tokens, K), data_format, data_init, generator) # Routing: normal (random) by default; balanced if AITER_MOE_EXPERT_BALANCE. topk_id, topk_w = _make_topk(hidden, experts, topk) @@ -521,6 +566,57 @@ def _logits_diff(actual: torch.Tensor, expected: torch.Tensor) -> float: return float(((x - y) ** 2).sum() / denom) +def _gemm_work_metrics( + *, + experts: int, + tokens: int, + topk: int, + model_dim: int, + inter_dim: int, + data_format: str, +) -> dict[str, tuple[float, float]]: + """Return conventional GEMM FLOPs and effective bytes for both stages. + + The byte model matches the grouped-MoE tuner: logical quantized inputs and + weights plus BF16 outputs. It excludes routing, quantization, scales, bias, + and other fused-MoE auxiliary traffic, so the reported bandwidth is an + effective GEMM bandwidth rather than measured HBM transactions. + """ + input_bytes = 0.5 if data_format == "a4w4" else 1.0 + weight_bytes = 0.5 + output_bytes = 2.0 + stage1_n = 2 * inter_dim + routed_rows = tokens * topk + + gemm1_flops = routed_rows * stage1_n * model_dim * 2 + gemm1_bytes = ( + routed_rows * model_dim * input_bytes + + routed_rows * stage1_n * output_bytes + + experts * model_dim * stage1_n * weight_bytes + ) + gemm2_flops = tokens * topk * model_dim * inter_dim * 2 + gemm2_bytes = ( + tokens * topk * inter_dim * input_bytes + + tokens * model_dim * output_bytes + + experts * inter_dim * model_dim * weight_bytes + ) + return { + "gemm1": (gemm1_flops, gemm1_bytes), + "gemm2": (gemm2_flops, gemm2_bytes), + "total": (gemm1_flops + gemm2_flops, gemm1_bytes + gemm2_bytes), + } + + +def _rates( + work: tuple[float, float], us: float | None +) -> tuple[float | None, float | None]: + """Convert a (FLOPs, bytes) work estimate and microseconds to rates.""" + if us is None or us <= 0: + return None, None + flops, data_bytes = work + return flops / us / 1e6, data_bytes / us / 1e3 + + # --------------------------------------------------------------------------- # Pytest correctness suite # --------------------------------------------------------------------------- @@ -543,7 +639,9 @@ def run_moe( kernel_bench: bool = False, warmup: int = 5, iters: int = 101, - const_init: float | None = None, + seed: int = 0, + data_init: str = "uniform", + scale_init: str = "auto", check_aot_cache: bool = True, ) -> dict: """Compare grouped FlyDSL MoE vs a PyTorch fp32 ref. ``bench`` selects the @@ -579,9 +677,11 @@ def run_moe( use_bias=use_bias, bench=bench, kernel_bench=kernel_bench, + seed=seed, warmup=warmup, iters=iters, - const_init=const_init, + data_init=data_init, + scale_init=scale_init, ) mode = "kernel" if kernel_bench else ("graph" if bench else "eager") ld = _logits_diff(out, ref) @@ -606,16 +706,38 @@ def run_moe( # --- perf (bench only): timed end-to-end inside _run_grouped_via_fused_moe --- if bench: + work = _gemm_work_metrics( + experts=experts, + tokens=tokens, + topk=topk, + model_dim=model_dim, + inter_dim=inter_dim, + data_format=data_format, + ) + tflops, bandwidth = _rates(work["total"], us) print( - f"[bench {tag}] fused_moe end-to-end us = {us:.2f} (graph=True)", + f"[bench {tag}] fused_moe end-to-end us = {us:.2f}, " + f"FLOPS = {work['total'][0]:.0f}, TFLOPS = {tflops:.2f}, " + f"Bandwidth = {bandwidth:.2f} GB/s (graph=True)", flush=True, ) metrics["us"] = us + metrics["flops"] = work["total"][0] + metrics["tflops"] = tflops + metrics["bandwidth_gbs"] = bandwidth # --- perf (kernel-bench only): per-kernel gemm1/gemm2 timing (looped alone) --- if kernel_bench: kernel_us = kernel_us or {} g1 = kernel_us.get("gemm1") g2 = kernel_us.get("gemm2") + work = _gemm_work_metrics( + experts=experts, + tokens=tokens, + topk=topk, + model_dim=model_dim, + inter_dim=inter_dim, + data_format=data_format, + ) if g1 is None and g2 is None: print( f"[kernel-bench {tag}] no grouped kernels captured " @@ -623,12 +745,20 @@ def run_moe( flush=True, ) else: - g1s = "n/a" if g1 is None else f"{g1:.2f}" - g2s = "n/a" if g2 is None else f"{g2:.2f}" - print( - f"[kernel-bench {tag}] gemm1 us = {g1s} gemm2 us = {g2s}", - flush=True, - ) + for name, elapsed_us in (("gemm1", g1), ("gemm2", g2)): + if elapsed_us is None: + print(f"[kernel-bench {tag}] {name}: n/a", flush=True) + continue + tflops, bandwidth = _rates(work[name], elapsed_us) + print( + f"[kernel-bench {tag}] {name}: us = {elapsed_us:.2f}, " + f"FLOPS = {work[name][0]:.0f}, TFLOPS = {tflops:.2f}, " + f"Bandwidth = {bandwidth:.2f} GB/s", + flush=True, + ) + metrics[f"{name}_flops"] = work[name][0] + metrics[f"{name}_tflops"] = tflops + metrics[f"{name}_bandwidth_gbs"] = bandwidth metrics["gemm1_us"] = g1 metrics["gemm2_us"] = g2 return metrics @@ -840,6 +970,13 @@ def summarize(rows: list): print(f" {r}", flush=True) return rows df = pd.DataFrame(rows) + empty_perf_columns = [ + column + for column in df.columns + if ("tflops" in column.lower() or "bandwidth" in column.lower()) + and df[column].isna().all() + ] + df = df.drop(columns=empty_perf_columns) try: table = df.to_markdown(index=False) except ImportError: @@ -965,28 +1102,55 @@ def run_csv_scenario(args) -> None: # Route each row to its format (a8w4 needs AITER_FORCE_A8W4=1). set_data_format(data_format) tol = VERIFY_TOL_A8W4 if data_format == "a8w4" else VERIFY_TOL_A4W4 - try: - metrics = run_moe( - data_format, - experts=experts, - tokens=tokens, - topk=topk, - model_dim=model_dim, - inter_dim=inter_dim, - tol=tol, - activation=activation, - swiglu_limit=args.swiglu_limit, - use_bias=not args.no_bias, - check_aot_cache=not args.no_check_aot_cache, - raise_on_fail=False, - bench=True, - kernel_bench=False, - warmup=args.warmup, - iters=args.iters, - const_init=args.const_init, - ) - except Exception as exc: # noqa: BLE001 - record, keep sweeping - print(f"[csv] row {idx}: ERROR {exc!r}", flush=True) + for data_init, scale_init in args.init_pairs: + try: + metrics = run_moe( + data_format, + experts=experts, + tokens=tokens, + topk=topk, + model_dim=model_dim, + inter_dim=inter_dim, + tol=tol, + activation=activation, + swiglu_limit=args.swiglu_limit, + use_bias=not args.no_bias, + check_aot_cache=not args.no_check_aot_cache, + raise_on_fail=False, + bench=True, + kernel_bench=False, + warmup=args.warmup, + iters=args.iters, + seed=args.seed, + data_init=data_init, + scale_init=scale_init, + ) + except Exception as exc: # noqa: BLE001 - record, keep sweeping + print(f"[csv] row {idx}: ERROR {exc!r}", flush=True) + rows.append( + { + "row": idx, + "data_format": data_format, + "act": act, + "tokens": tokens, + "model_dim": model_dim, + "inter_dim": inter_dim, + "experts": experts, + "topk": topk, + "data_init": data_init, + "scale_init": scale_init, + "seed": args.seed, + "logits_diff": float("nan"), + "rel_l2": float("nan"), + "pass": False, + "error": repr(exc), + "us": None, + "gemm1_us": None, + "gemm2_us": None, + } + ) + continue + rows.append( { "row": idx, @@ -997,36 +1161,20 @@ def run_csv_scenario(args) -> None: "inter_dim": inter_dim, "experts": experts, "topk": topk, - "logits_diff": float("nan"), - "rel_l2": float("nan"), - "pass": False, - "error": repr(exc), - "us": None, - "gemm1_us": None, - "gemm2_us": None, + "data_init": data_init, + "scale_init": scale_init, + "seed": args.seed, + "logits_diff": metrics["logits_diff"], + "rel_l2": metrics["rel_l2"], + "pass": metrics["passed"], + "error": None, + "us": metrics.get("us"), + "TFLOPS": metrics.get("tflops"), + "Bandwidth (GB/s)": metrics.get("bandwidth_gbs"), + "gemm1_us": metrics.get("gemm1_us"), + "gemm2_us": metrics.get("gemm2_us"), } ) - continue - - rows.append( - { - "row": idx, - "data_format": data_format, - "act": act, - "tokens": tokens, - "model_dim": model_dim, - "inter_dim": inter_dim, - "experts": experts, - "topk": topk, - "logits_diff": metrics["logits_diff"], - "rel_l2": metrics["rel_l2"], - "pass": metrics["passed"], - "error": None, - "us": metrics.get("us"), - "gemm1_us": metrics.get("gemm1_us"), - "gemm2_us": metrics.get("gemm2_us"), - } - ) summarize(rows) failed = [r for r in rows if not r["pass"]] @@ -1083,6 +1231,21 @@ def main() -> None: parser.add_argument("--inter-dim", type=int, default=256) parser.add_argument("--warmup", type=int, default=5) parser.add_argument("--iters", type=int, default=101) + parser.add_argument( + "--data-init", + dest="data_init", + nargs="+", + choices=bench_init.DATA_DISTS, + default=None, + help="DATA initialization distribution(s), paired position-wise with " + "--scale-init (length-1 broadcasts). Default: constant uniform", + ) + parser.add_argument( + "--seed", + type=int, + default=0, + help="RNG seed for data and routing (default: 0)", + ) parser.add_argument( "--act", choices=("silu", "swiglu", "situv2"), @@ -1113,16 +1276,13 @@ def main() -> None: help="run with zero stage1/stage2 bias tensors", ) parser.add_argument( - "--const-init", - type=float, - nargs="?", - const=0.0, + "--scale-init", + dest="scale_init", + nargs="+", + choices=bench_init.E8M0_SCALE_DISTS, default=None, - metavar="VALUE", - help="initialize activations (A), weights (B), weight scales (Bs), and " - "bias to the constant VALUE instead of random values. Bare --const-init " - "uses 0.0 (zero-init). uint8 tensors (weights, scales) are filled with " - "int(VALUE).", + help="E8M0 SCALE initialization distribution(s), paired position-wise " + "with --data-init (length-1 broadcasts). Default: constant auto", ) parser.add_argument( "--real-gemm", @@ -1140,6 +1300,18 @@ def main() -> None: "miss raises. Pass this flag to allow runtime JIT compilation.", ) args = parser.parse_args() + data_init_list = args.data_init or ["constant", "uniform"] + scale_init_list = args.scale_init or ["constant", "auto"] + if len(data_init_list) == 1: + data_init_list *= len(scale_init_list) + if len(scale_init_list) == 1: + scale_init_list *= len(data_init_list) + if len(data_init_list) != len(scale_init_list): + parser.error( + "--data-init and --scale-init must have equal length " + "(or length 1 to broadcast)" + ) + args.init_pairs = list(zip(data_init_list, scale_init_list)) if not args.real_gemm: _mock_grouped_gemm() @@ -1171,48 +1343,59 @@ def main() -> None: if len(token_list) > 1: print(f"\n===== tokens={_tok} =====", flush=True) - tol = VERIFY_TOL_A8W4 if args.data_format == "a8w4" else VERIFY_TOL_A4W4 - # raise_on_fail=False so one out-of-gate token does not abort the - # sweep; the failure is recorded and reported after the table. - metrics = run_moe( - args.data_format, - experts=args.experts, - tokens=args.tokens, - topk=args.topk, - model_dim=args.model_dim, - inter_dim=args.inter_dim, - tol=tol, - activation=activation, - swiglu_limit=args.swiglu_limit, - situ_beta=args.situ_beta, - situ_linear_beta=args.situ_linear_beta, - use_bias=not args.no_bias, - check_aot_cache=not args.no_check_aot_cache, - raise_on_fail=False, - bench=args.scenario == "bench", - kernel_bench=args.scenario == "kernel", - warmup=args.warmup, - iters=args.iters, - const_init=args.const_init, - ) - rows.append( - { - "data_format": args.data_format, - "act": args.act, - "init": "random" if args.const_init is None else "const", - "experts": args.experts, - "tokens": _tok, - "topk": args.topk, - "model_dim": args.model_dim, - "inter_dim": args.inter_dim, - "logits_diff": metrics["logits_diff"], - "rel_l2": metrics["rel_l2"], - "pass": metrics["passed"], - "us": metrics.get("us"), - "gemm1_us": metrics.get("gemm1_us"), - "gemm2_us": metrics.get("gemm2_us"), - } - ) + for data_init, scale_init in args.init_pairs: + tol = VERIFY_TOL_A8W4 if args.data_format == "a8w4" else VERIFY_TOL_A4W4 + # raise_on_fail=False so one out-of-gate token does not abort the + # sweep; the failure is recorded and reported after the table. + metrics = run_moe( + args.data_format, + experts=args.experts, + tokens=args.tokens, + topk=args.topk, + model_dim=args.model_dim, + inter_dim=args.inter_dim, + tol=tol, + activation=activation, + swiglu_limit=args.swiglu_limit, + situ_beta=args.situ_beta, + situ_linear_beta=args.situ_linear_beta, + use_bias=not args.no_bias, + check_aot_cache=not args.no_check_aot_cache, + raise_on_fail=False, + bench=args.scenario == "bench", + kernel_bench=args.scenario == "kernel", + warmup=args.warmup, + iters=args.iters, + seed=args.seed, + data_init=data_init, + scale_init=scale_init, + ) + rows.append( + { + "data_format": args.data_format, + "act": args.act, + "data_init": data_init, + "scale_init": scale_init, + "seed": args.seed, + "experts": args.experts, + "tokens": _tok, + "topk": args.topk, + "model_dim": args.model_dim, + "inter_dim": args.inter_dim, + "logits_diff": metrics["logits_diff"], + "rel_l2": metrics["rel_l2"], + "pass": metrics["passed"], + "us": metrics.get("us"), + "TFLOPS": metrics.get("tflops"), + "Bandwidth (GB/s)": metrics.get("bandwidth_gbs"), + "gemm1_us": metrics.get("gemm1_us"), + "gemm2_us": metrics.get("gemm2_us"), + "gemm1_TFLOPS": metrics.get("gemm1_tflops"), + "gemm1_Bandwidth (GB/s)": metrics.get("gemm1_bandwidth_gbs"), + "gemm2_TFLOPS": metrics.get("gemm2_tflops"), + "gemm2_Bandwidth (GB/s)": metrics.get("gemm2_bandwidth_gbs"), + } + ) # Always print the summary table (verify and bench). summarize(rows) diff --git a/op_tests/flydsl_tests/test_flydsl_moe_a8w4.py b/op_tests/flydsl_tests/test_flydsl_moe.py similarity index 99% rename from op_tests/flydsl_tests/test_flydsl_moe_a8w4.py rename to op_tests/flydsl_tests/test_flydsl_moe.py index 7b203878fa..0c3d8e4208 100644 --- a/op_tests/flydsl_tests/test_flydsl_moe_a8w4.py +++ b/op_tests/flydsl_tests/test_flydsl_moe.py @@ -7,8 +7,8 @@ inter=640) and FlyDSL stage2 / E2E with GUI preshuffle on gfx950. Usage: - pytest op_tests/flydsl_tests/test_flydsl_moe_a8w4.py -q - pytest op_tests/flydsl_tests/test_flydsl_moe_a8w4.py -k tile_k + pytest op_tests/flydsl_tests/test_flydsl_moe.py -q + pytest op_tests/flydsl_tests/test_flydsl_moe.py -k tile_k """ from __future__ import annotations diff --git a/op_tests/multigpu_tests/test_mega_moe_gfx1250.py b/op_tests/multigpu_tests/bench_mega_moe.py similarity index 60% rename from op_tests/multigpu_tests/test_mega_moe_gfx1250.py rename to op_tests/multigpu_tests/bench_mega_moe.py index a183f41fa7..f6f29c16c8 100644 --- a/op_tests/multigpu_tests/test_mega_moe_gfx1250.py +++ b/op_tests/multigpu_tests/bench_mega_moe.py @@ -5,7 +5,10 @@ N (default 61, DeepSeek-V4-Pro) MoE layers are chained. The ``base`` mode uses Mori v2 dispatch -> AITER fused_moe -> Mori v2 combine. ``fused`` calls only ``MegaMoEGfx1250``, which owns AITER's dispatch -> fused_moe -> fused-combine -pipeline. The combined output plus residual feeds the next layer. +pipeline. The combined output plus residual feeds the next layer. ``both`` (the +default) walks base then fused in ONE process, so the two share the weights, the +tokens, the routings and the single fp32 reference, and land as two rows of the +same summary table -- perf and accuracy compared column by column. Two isolated paths (never touch each other's intermediates; they only share the config, the bf16 weights and the per-layer routings): @@ -22,17 +25,29 @@ Launch (4x gfx1250; every env knob below is already the script's default): cd # avoid the /app/triton namespace shadow - torchrun --standalone --nproc_per_node=4 test_mega_moe_gfx1250.py \ - -q a4w4_mxfp4 -e 384 -k 6 -hd 7168 -id 3072 --layers 61 --combine base + torchrun --standalone --nproc_per_node=4 bench_mega_moe.py \ + -q a4w4_mxfp4 -e 384 -k 6 -hd 7168 -id 3072 --layers 61 --combine both # Set MORI_CCO_BC to a prebuilt libmori_cco_device.bc to skip CCO JIT. Env / CLI: --layers --logits_tol --acc_verify --dispatch_wire --combine -tpr -hd -id -e -k --shared_E -q + --data-init --seed --warmup --iters --prof_replays + +``--data-init`` / ``--scale-init`` / ``--seed`` are the shared ubench knobs from +``aiter.benchmark_data_init.add_data_init_args``. ``--scale-init`` is accepted +for CLI +compatibility but unused here: every scale is derived by quantizing the +generated weights, never drawn independently. """ +from __future__ import annotations + import argparse +import math import os +import time +import pandas as pd import torch import torch.distributed as dist import torch.profiler as tprof @@ -46,6 +61,8 @@ get_torch_quant, pertoken_quant, ) +from aiter.benchmark_data_init import add_data_init_args, fill, make_generator +from aiter.benchmark_reporting import print_json_table from aiter.fused_moe import fused_moe from aiter.ops.flydsl.moe_common import GateMode from aiter.ops.shuffle import moe_shuffle_scale, shuffle_weight @@ -77,6 +94,9 @@ _FP8_DTYPE = dtypes.fp8 QUANT_KEYS = ["No", "per_Token", "per_128x128", "a8w4_mxfp4", "a4w4_mxfp4"] _MXFP4_KEYS = ("a8w4_mxfp4", "a4w4_mxfp4") +# add_data_init_args' --scale-init default. Kept here so main() can tell whether +# the caller asked for a scale distribution this test cannot honour. +_DEFAULT_SCALE_INIT = "constant" def _import_mori_comm(): @@ -161,6 +181,40 @@ def resolve_dispatch_wire(wire, quant_key): return wire +def resolve_data_init(data_init): + """``--data-init`` is the shared ubench list form, meant to sweep several + distributions in one invocation. Here the weights are quantized and the whole + N-layer chain is captured into a CUDA graph once per process, so a run takes + exactly one distribution -- sweep by launching the script per distribution.""" + dists = list(data_init) if isinstance(data_init, (list, tuple)) else [data_init] + if len(dists) != 1: + raise ValueError( + f"--data-init takes a single distribution here, got {dists}: the " + "weight quant and the graph capture are both per-run" + ) + return dists[0] + + +def resolve_combine_modes(combine, spec, dist_ctx): + """The combine modes one invocation benchmarks, in the order they run. + + ``both`` is the default so a plain run always produces the base-vs-fused + comparison. The fused combine is mxfp4-only, so for the other quant keys + ``both`` degrades to base alone instead of failing -- an explicit + ``--combine fused`` still raises in setup(), where the constraint belongs.""" + if combine != "both": + return [combine] + if not spec["is_mxfp4"]: + if dist_ctx.rank == 0: + print( + "# note: --combine both runs base only for this quant key -- the " + "fused combine is mxfp4-only", + flush=True, + ) + return ["base"] + return ["base", "fused"] + + # Weight quantization + shuffle (device path) / dequant (reference) def weight_per_128x128_quant(weight, quant_dtype): E, dim1, dim2 = weight.shape @@ -301,37 +355,33 @@ def moe_forward( # Shared setup (fed to BOTH reference and device path) _WEIGHT_SEED = 70000 # identical on every rank so the global expert set agrees +_WEIGHT_AMPL = 0.1 # weight amplitude, was the literal `/ 10` below -def make_shared_weights(E, hdim, idim, dtype, dev, shared_E=0, seed=_WEIGHT_SEED): +def make_shared_weights( + E, hdim, idim, dtype, dev, shared_E=0, seed=_WEIGHT_SEED, data_dist="norm" +): """One weight set reused by every layer. Same seed on all ranks so the global - expert partition is consistent. Returns bf16 (w1[E,2I,H], w2[E,H,I], sw1, sw2).""" - gen = torch.Generator(device=dev).manual_seed(seed) - w1 = ( - torch.randn((E, 2 * idim, hdim), generator=gen, device=dev, dtype=torch.float32) - / 10 - ).to(dtype) - w2 = ( - torch.randn((E, hdim, idim), generator=gen, device=dev, dtype=torch.float32) - / 10 - ).to(dtype) + expert partition is consistent. Returns bf16 (w1[E,2I,H], w2[E,H,I], sw1, sw2). + + ``data_dist`` is a ``--data-init`` distribution. Every mode is scaled down by + _WEIGHT_AMPL: at unit amplitude the narrow fp4/fp8 activation quant saturates + and the N-layer residual chain diverges to NaN after a few layers.""" + gen = make_generator(seed, device=dev) + + def _w(experts, rows, cols): + # fill() only row-chunks its fp32 staging for 2-D shapes, so ask for the + # flat [experts*rows, cols] and view it back: at E=384 that caps the + # staging near 1 GiB instead of materializing the set in fp32. + w = fill((experts * rows, cols), data_dist, gen, dtype=dtype, device=dev) + return w.mul_(_WEIGHT_AMPL).view(experts, rows, cols) + + w1 = _w(E, 2 * idim, hdim) + w2 = _w(E, hdim, idim) sw1 = sw2 = None if shared_E > 0: - sw1 = ( - torch.randn( - (shared_E, 2 * idim, hdim), - generator=gen, - device=dev, - dtype=torch.float32, - ) - / 10 - ).to(dtype) - sw2 = ( - torch.randn( - (shared_E, hdim, idim), generator=gen, device=dev, dtype=torch.float32 - ) - / 10 - ).to(dtype) + sw1 = _w(shared_E, 2 * idim, hdim) + sw2 = _w(shared_E, hdim, idim) return w1, w2, sw1, sw2 @@ -660,13 +710,19 @@ def _pipeline(self, x0): return x # ---- CUDA graph capture (all N layers in ONE graph) ---- # + # Eager passes before the capture. NOT a measurement knob, hence not on the + # CLI: they prime the fused_moe lru_cache, the JIT and the allocator, and + # without them that work would land inside the capture and fail it. 3 is what + # torch's own CUDA-graph guidance warms up with. + _CAPTURE_WARMUP = 3 + def capture(self, x0): self.x0_static = x0.clone() # warmup on a side stream: primes fused_moe lru_cache + allocator. s = torch.cuda.Stream() s.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(s): - for _ in range(3): + for _ in range(self._CAPTURE_WARMUP): self._pipeline(self.x0_static) torch.cuda.current_stream().wait_stream(s) torch.cuda.synchronize() @@ -679,18 +735,18 @@ def capture(self, x0): self.comm.barrier() # ---- perf: torch.profiler breakdown + graph-replay wall-clock ---- # - _N_WARMUP = 5 - _N_PROF_REPLAYS = 3 # graph replays captured by torch.profiler in bench() - - def bench(self): + def bench(self, warmup=5, iters=10, prof_replays=3): """Time the ONE-graph N-layer dispatch->gemm->combine chain. The graph - already contains all N layers, so a single replay IS the per-chain - measurement -- no separate replay-count knob. 5 warmup replays first. - Returns (total_us for all N layers, per_layer_us, prof_us). + already contains all N layers, so ONE replay is one full chain; `iters` of + them are timed individually after `warmup` untimed ones. + Returns (stats, prof_us) with stats = {min, median, mean, max} in us. - - total_us = host wall-clock of one graph replay (one sync after; not + - a sample is the host wall-clock of one graph replay (one sync after; not cuda.Event). For a GPU-bound MoE chain this ~= GPU time. - - torch.profiler over one EAGER pipeline pass for the per-op breakdown. + - min is the closest to an undisturbed replay, median to the steady state + of a GPU held at full load; the two drifting apart is what a straggler + rank or a throttled clock looks like, hence both are reported. + - torch.profiler over `prof_replays` replays for the per-op breakdown. NOTE: this ROCm torch build reports self_device_time_total == 0 for every event (verified even for a plain matmul), so torch.profiler cannot give a @@ -698,18 +754,30 @@ def bench(self): If a future build populates device time, prof_us below becomes > 0.""" import time - for _ in range(self._N_WARMUP): + assert iters >= 1, f"--iters must time at least one replay, got {iters}" + + for _ in range(warmup): self.graph.replay() torch.cuda.synchronize() self.comm.barrier() - # one full N-layer graph replay == the performance measurement. - t0 = time.perf_counter() - self.graph.replay() - torch.cuda.synchronize() - total_us = (time.perf_counter() - t0) * 1e6 + # each full N-layer graph replay is one measurement. + samples = [] + for _ in range(iters): + t0 = time.perf_counter() + self.graph.replay() + torch.cuda.synchronize() + samples.append((time.perf_counter() - t0) * 1e6) self.comm.barrier() + samples.sort() + stats = { + "min": samples[0], + "median": samples[len(samples) // 2], + "mean": sum(samples) / len(samples), + "max": samples[-1], + } + # torch.profiler breakdown over CUDA-graph replays: roctracer/kineto does # surface the per-kernel timeline inside the graph on this build, so we # profile the actual graph (matches the measured per-layer wall) instead of @@ -717,13 +785,13 @@ def bench(self): with tprof.profile( activities=[tprof.ProfilerActivity.CPU, tprof.ProfilerActivity.CUDA] ) as prof: - for _ in range(self._N_PROF_REPLAYS): + for _ in range(prof_replays): self.graph.replay() torch.cuda.synchronize() self.comm.barrier() self._prof = prof prof_us = sum(_event_device_us(e) for e in prof.key_averages()) - return total_us, total_us / self.n_layers, prof_us + return stats, prof_us def final_output(self): self.graph.replay() @@ -731,11 +799,20 @@ def final_output(self): return self.out_static.detach().clone() def teardown(self): + # Drop the graph and its static tensors too, not just the mori handles: + # under --combine both the next mode captures its own N-layer graph right + # after this, and the first graph's pool would otherwise stay reserved. self.graph = None + self.x0_static = None + self.out_static = None + self._prof = None if self.mega is not None: self.mega.close() + self.mega = None + self.op = None if self.comm is not None: self.comm.destroy() + self.comm = None def _event_device_us(e): @@ -748,10 +825,83 @@ def _event_device_us(e): return 0.0 +def _run_distributed_smi_replay(pipe, dist_ctx, median_us, n_layers, combine_mode): + """Replay the Mega graph while every rank monitors its local GPU. + + `combine_mode` goes into the label so a --combine both run does not file two + different pipelines under the same name.""" + if os.environ.get("AITER_SMI_MONITOR", "0") != "1": + return + + from aiter.smi_monitor import GpuMonitor, emit_smi_result + + interval_s = float(os.environ.get("AITER_SMI_INTERVAL", "0.05")) + duration_s = float(os.environ.get("AITER_SMI_DURATION", "1.0")) + if interval_s <= 0 or duration_s <= 0 or median_us <= 0: + raise ValueError( + "Mega MoE SMI interval, duration and measured median must be positive" + ) + replay_count = max(1, math.ceil(duration_s * 1e6 / median_us)) + + monitor = None + monitor_error = None + try: + monitor = GpuMonitor( + device_index=torch.cuda.current_device(), interval_s=interval_s + ) + monitor.start() + except Exception as error: # noqa: BLE001 - propagate to every rank + monitor_error = f"rank {dist_ctx.rank}: {type(error).__name__}: {error}" + + monitor_errors = dist_ctx.gather_objects(monitor_error) + failed = [error for error in monitor_errors if error is not None] + if failed: + if monitor is not None: + monitor.stop() + raise RuntimeError("Mega MoE SMI monitor failed to start: " + "; ".join(failed)) + + # Gloo barriers align CPU submission without adding a GPU collective to the + # measured Mega graph window. Every rank executes exactly the same replay + # count; a local duration loop would diverge and deadlock the collectives. + dist.barrier() + window_start = time.perf_counter() + for _ in range(replay_count): + pipe.graph.replay() + torch.cuda.synchronize() + window_end = time.perf_counter() + monitor.stop() + dist.barrier() + + samples = [ + sample + for sample in monitor.samples + if window_start <= sample["timestamp_s"] <= window_end + ] + expected_samples = max(1, int(duration_s / interval_s)) + base_label = os.environ.get("AITER_SMI_LABEL", "mega_moe") + local_result = { + "label": f"{base_label}/{combine_mode}/mega_graph_{n_layers}_layers", + "device": dist_ctx.local_rank, + "rank": dist_ctx.rank, + "interval_s": interval_s, + "duration_s": window_end - window_start, + "launches": replay_count, + "samples": len(samples), + "sample_status": ( + "ok" if len(samples) >= max(2, expected_samples // 2) else "insufficient" + ), + "metrics": monitor.summary(start_s=window_start, end_s=window_end), + } + results = dist_ctx.gather_objects(local_result) + if dist_ctx.rank == 0: + for result in results: + emit_smi_result(result) + + def _aggregate_prof_table(prof, dist_ctx, per_layer_denom=1.0, row_limit=200): """Collect the torch.profiler per-kernel table ACROSS ranks (collective; call on every rank). Each rank contributes {name: (self_device_us_total, count)}; - rank 0 returns a table of each kernel's per-call self device time with ONE + rank 0 returns rows with each kernel's per-call self device time with ONE COLUMN PER RANK plus the cross-rank mean, so a straggler (a throttled GPU, an unbalanced expert distribution) shows up as a row that disagrees across columns instead of being averaged away. `-` means the kernel never ran there. @@ -785,31 +935,122 @@ def _aggregate_prof_table(prof, dist_ctx, per_layer_denom=1.0, row_limit=200): rows.append((avg_self, name, per_call, pc_avg, avg_count)) rows.sort(key=lambda r: (-r[0], r[1])) dev_per_layer = total_self / per_layer_denom if per_layer_denom else 0.0 - # Wide enough for a full TDM GEMM name, whose tile/warp/buffer recipe and its - # `_epscatter` / `_prefetch` suffix are the whole point of reading this table - # (e.g. a8w4_tdm_fp4_t256x256x256_w2x2_b3_K3072_e96_cn4_prefetch_epscatter). - name_w = 72 - lines = [ - ( - f"# per-call self device time (us) by rank, {world} ranks " - f"(rows sorted by total self time):" - ), - f"{'Name':<{name_w}}" - + "".join(f"{f'rank{r}':>11}" for r in range(world)) - + f"{'avg':>11}{'calls':>8}", - ] - for avg_self, name, per_call, pc_avg, avg_count in rows[:row_limit]: - cells = "".join( - f"{v:>11.3f}" if v is not None else f"{'-':>11}" for v in per_call - ) - lines.append( - f"{name[:name_w]:<{name_w}}{cells}{pc_avg:>11.3f}{avg_count:>8.1f}" - ) - lines.append( - f"# TOTAL self device time over ALL {len(rows)} kernels = {total_self:.1f} us " - f"-> {dev_per_layer:.1f} us/layer (device-busy; compare to per_layer wall)" + table_rows = [] + for _avg_self, name, per_call, pc_avg, avg_count in rows[:row_limit]: + row = { + "kernel": name, + "avg_us": pc_avg, + "calls": avg_count, + } + row.update({f"rank{rank}_us": value for rank, value in enumerate(per_call)}) + table_rows.append(row) + return { + "rows": table_rows, + "summary": [ + { + "world_size": world, + "profiled_kernels": len(rows), + "total_self_device_us": total_self, + "device_us_per_layer": dev_per_layer, + } + ], + } + + +def _stage2_overlap_rate(kernel_rows, idim): + """How much of stage 2's communication the fused combine hides behind gemm2. + + Stage 2 is the second expert GEMM and everything that moves its output home. + base splits that into compute -- the K{idim} GEMM plus the gather-reduce that + lands the result -- and communication, the mori combine; the two are separate + kernels, so base pays for them back to back. fused folds the scatter into the + GEMM itself, so its stage 2 is that one (heavier) GEMM plus the small fused + combine. Whatever the sum of base's two halves loses by becoming the fused + total is time fused managed to overlap, and the most it could ever hide is + the smaller of the two halves -- hence the min() denominator, which puts a + perfect overlap at 1.0 and no overlap at 0.0. + + Both fused combine kernels count, the sync one included: it is the wait the + fused path did not manage to hide, and dropping it would book that wait as + successful overlap. + + gemm2 is matched by its K{idim} contraction, which is what separates it from + gemm1's K{hidden}. Returns None if any kernel the formula needs is absent, so + a run without --profile_table simply carries no rate.""" + + def total_us(rows, match): + hits = [r["avg_us"] for r in rows if match(r["kernel"])] + return sum(hits) if hits else None + + base, fused = kernel_rows.get("base"), kernel_rows.get("fused") + if not base or not fused: + return None + + def is_gemm2(name): + return f"_K{idim}_" in name + + parts = ( + total_us(base, is_gemm2), + total_us(base, lambda n: n.startswith("moe_gather_reduce")), + total_us(base, lambda n: n.startswith("mori_ep_combine")), + total_us(fused, is_gemm2), + total_us(fused, lambda n: n.startswith("ep_combine_fused")), ) - return "\n".join(lines) + if any(p is None for p in parts): + return None + base_gemm2, base_gather, base_comm, fused_gemm2, fused_comm = parts + compute = base_gemm2 + base_gather + comm = base_comm + if min(compute, comm) <= 0: + return None + return (compute + comm - (fused_gemm2 + fused_comm)) / min(compute, comm) + + +def _emit_table(name, rows, max_col_width=72): + """Print the rows twice: an aligned frame for whoever opens the log, then the + one machine-readable line the benchmark driver consumes. + + Both render the same DataFrame -- print_json_table builds one anyway to + serialize it -- so the readable half costs nothing but keeps a 6 KB JSON line + from being the only view of a 21-kernel table. A single row is transposed; + with several rows the columns that hold the same value everywhere are hoisted + into a one-line prefix, which is what keeps the 20+ config columns of the + summary from repeating down the table. + + max_col_width fits a full TDM GEMM name (its tile/warp/buffer recipe plus the + _prefetch / _epscatter suffix is the whole point of reading that table) while + still cutting torch's 200-char template names down to something scannable.""" + df = pd.DataFrame([row for row in rows if row is not None]) + print(f"\n# {name}", flush=True) + if df.empty: + print("# (no rows)", flush=True) + else: + # Decide what is constant BEFORE rounding: base and fused logits_diff + # agree to 4 decimals, and rounding first would hoist that difference out + # of the table as if the two modes had returned the same number. + const = [c for c in df.columns if df[c].nunique(dropna=False) == 1] + # Significant digits, not decimal places: the same table carries 555528.154 + # us and a 0.475270 logits_diff, and rounding both to 3 decimals would + # print the two modes' accuracy as an identical 0.475. + for column in df.select_dtypes(include="float").columns: + df[column] = df[column].map( + lambda v: v if pd.isna(v) else float(f"{v:.6g}") + ) + if len(df) == 1: + # astype(object) keeps each value's own type; transposing a numeric + # frame would otherwise widen the ints to float and print "4.000". + print(df.astype(object).T.to_string(header=False), flush=True) + else: + if const: + print( + "# " + " ".join(f"{c}={df[c].iloc[0]}" for c in const), flush=True + ) + df = df.drop(columns=const) + print( + df.to_string(index=False, max_colwidth=max_col_width), + flush=True, + ) + print_json_table(name, rows) def _device_shared_ffn(tokens, sw1, sw2): @@ -849,19 +1090,31 @@ def main(): E % dist_ctx.world == 0 ), f"E={E} must be divisible by world_size={dist_ctx.world}" + data_dist = resolve_data_init(args.data_init) + if dist_ctx.rank == 0: print( f"[cfg] world={dist_ctx.world} layers={n_layers} tokens/rank={ct} hidden={hdim} " f"inter={idim} E={E} topk={topk} EPR={E // dist_ctx.world} quant={args.quant_type} " f"combine={args.combine} dispatch_wire={spec['dispatch_wire']} " f"force_a8w4={os.environ['AITER_FORCE_A8W4']} " - f"gate={spec['gate_mode'].name} shared_E={args.shared_experts} gfx={get_gfx()}", + f"gate={spec['gate_mode'].name} shared_E={args.shared_experts} " + f"data_init={data_dist} seed={args.seed} gfx={get_gfx()}", flush=True, ) + if list(args.scale_init) != [_DEFAULT_SCALE_INIT]: + print( + f"# note: --scale-init {' '.join(args.scale_init)} is ignored -- " + "every scale here comes from quantizing the generated weights", + flush=True, + ) # ---- shared inputs: weights (same on all ranks) + this rank's tokens/routing. # args.seed shifts all RNG; weights stay rank-independent (identical global # experts), tokens/routing vary per rank. Default keeps runs reproducible. + # --data-init picks how the weights and the layer-0 tokens are filled; the + # routing stays random in every mode, since a zero/constant routing would + # collapse every token onto expert 0 and stop measuring dispatch/combine. w1_bf, w2_bf, sw1, sw2 = make_shared_weights( E, hdim, @@ -870,84 +1123,134 @@ def main(): dev, shared_E=args.shared_experts, seed=_WEIGHT_SEED + args.seed, + data_dist=data_dist, ) - x0 = torch.randn( - ct, - hdim, - generator=torch.Generator(device=dev).manual_seed( - 1000 + dist_ctx.rank + args.seed - ), + x0 = fill( + (ct, hdim), + data_dist, + make_generator(1000 + dist_ctx.rank + args.seed, device=dev), + dtype=dtypes.bf16, device=dev, - dtype=torch.float32, - ).to(dtypes.bf16) + ) routings = make_routings( n_layers, ct, E, topk, dev, seed=4242 + 100 * dist_ctx.rank + args.seed ) - # ---- device path (isolated): setup -> capture 61 layers in one graph -> bench. - pipe = DeviceMoEPipeline( - dist_ctx, - E, - hdim, - idim, - topk, - spec, - n_layers, - w1_bf, - w2_bf, - sw1, - sw2, - routings, - ct, - combine_mode=args.combine, - ) - pipe.setup(x0) - pipe.capture(x0) - total_us, per_layer_us, prof_us = pipe.bench() - # Aggregate perf across ranks (collective calls -> run on every rank). - total_us = dist_ctx.allreduce_avg_float(total_us) - per_layer_us = dist_ctx.allreduce_avg_float(per_layer_us) - prof_us = dist_ctx.allreduce_avg_float(prof_us) - tbl = None - if args.profile_table: + # ---- device path (isolated): setup -> capture 61 layers in one graph -> bench, + # once per combine mode. Every rank walks `modes` in the same order, so the + # collectives inside the loop stay in step. + modes = resolve_combine_modes(args.combine, spec, dist_ctx) + summary_rows = [] + outputs = {} + kernel_rows = {} # per mode, kept for the stage-2 overlap rate below + for combine_mode in modes: + if dist_ctx.rank == 0 and len(modes) > 1: + print(f"# ---- combine={combine_mode} ----", flush=True) + pipe = DeviceMoEPipeline( + dist_ctx, + E, + hdim, + idim, + topk, + spec, + n_layers, + w1_bf, + w2_bf, + sw1, + sw2, + routings, + ct, + combine_mode=combine_mode, + ) + pipe.setup(x0) + pipe.capture(x0) + stats, prof_us = pipe.bench( + warmup=args.warmup, iters=args.iters, prof_replays=args.prof_replays + ) + # Aggregate perf across ranks (collective calls -> run on every rank, and + # the dict is built in the same order everywhere so the allreduces stay in + # step). + stats = {k: dist_ctx.allreduce_avg_float(v) for k, v in stats.items()} + per_layer_us = stats["median"] / n_layers + prof_us = dist_ctx.allreduce_avg_float(prof_us) + _run_distributed_smi_replay( + pipe, dist_ctx, stats["median"], n_layers, combine_mode + ) + # Aggregate unconditionally, print only on request: bench() profiles the + # replays either way and this is one gather of a ~20-entry dict, while + # stage2_overlap_rate is a result the summary should carry whether or not + # anyone asked for the per-kernel table. Collective, so every rank calls + # it -- which is also why it must stay outside the --profile_table guard + # rather than being duplicated on both sides of it. tbl = _aggregate_prof_table( pipe._prof, dist_ctx, - per_layer_denom=pipe._N_PROF_REPLAYS * n_layers, - ) - # Save a chrome/perfetto timeline per rank so the actual kernel timeline - # (and any gaps) can be inspected directly. Opt-in (--save_trace): the - # export can stall multi-rank graph-profile runs, so it is off by default. - if args.save_trace: - _trace_path = f"/tmp/mega_trace_{args.combine}_rank{dist_ctx.rank}.json" - try: - pipe._prof.export_chrome_trace(_trace_path) - if dist_ctx.rank == 0: - print( - f"# trace saved: /tmp/mega_trace_{args.combine}_rank*.json", - flush=True, - ) - except Exception as _e: # noqa: BLE001 - if dist_ctx.rank == 0: - print(f"# trace export failed: {_e}", flush=True) - if dist_ctx.rank == 0: - prof_note = ( - f"prof_device={prof_us:.1f}us" - if prof_us > 0 - else "prof_device=n/a (this ROCm torch.profiler emits no device time)" + per_layer_denom=args.prof_replays * n_layers, ) - print( - f"# MEGA-MOE layers={n_layers} tokens/rank={ct}: " - f"total={total_us:.1f} us per_layer={per_layer_us:.1f} us " - f"(avg over {dist_ctx.world} ranks; dispatch+gemm+combine, 1 graph replay) " - f"{prof_note}", - flush=True, + if dist_ctx.rank == 0 and tbl is not None: + kernel_rows[combine_mode] = tbl["rows"] + if args.profile_table: + # Save a chrome/perfetto timeline per rank so the actual kernel + # timeline (and any gaps) can be inspected directly. Opt-in + # (--save_trace): the export can stall multi-rank graph-profile runs, + # so it is off by default. + if args.save_trace: + _trace_path = f"/tmp/mega_trace_{combine_mode}_rank{dist_ctx.rank}.json" + try: + pipe._prof.export_chrome_trace(_trace_path) + if dist_ctx.rank == 0: + print( + f"# trace saved: /tmp/mega_trace_{combine_mode}_rank*.json", + flush=True, + ) + except Exception as _e: # noqa: BLE001 + if dist_ctx.rank == 0: + print(f"# trace export failed: {_e}", flush=True) + # One table per mode, tagged with it: base and fused run a different + # kernel mix, so merging them into a single table would compare rows + # that never ran in the same pipeline. + if dist_ctx.rank == 0 and tbl is not None: + _emit_table(f"mega_moe kernel profile [{combine_mode}]", tbl["rows"]) + _emit_table( + f"mega_moe kernel profile summary [{combine_mode}]", tbl["summary"] + ) + + # Replay once more for the accuracy snapshot while the graph is still + # alive; teardown below frees it. + if args.acc_verify: + outputs[combine_mode] = pipe.final_output().float() + summary_rows.append( + { + "quant_type": args.quant_type, + "combine": combine_mode, + "data_init": data_dist, + "seed": args.seed, + "world_size": dist_ctx.world, + "tokens_per_rank": ct, + "experts": E, + "topk": topk, + "hidden": hdim, + "intermediate": idim, + "layers": n_layers, + "warmup": args.warmup, + "iters": args.iters, + "min_us": stats["min"], + "mean_us": stats["mean"], + "median_us": stats["median"], + "max_us": stats["max"], + "per_layer_us": per_layer_us, + "prof_device_us": prof_us if prof_us > 0 else None, + } ) - if tbl is not None: - print(tbl, flush=True) + pipe.teardown() + del pipe + torch.cuda.empty_cache() # ---- accuracy (isolated CPU/fp32 reference): end-to-end accumulated compare. - accuracy_failure = None + # ONE reference for every mode: the modes differ only in how combine moves the + # expert output, so they answer to the same ground truth -- and this reference + # is by far the most expensive part of the run. + failures = [] if args.acc_verify: auto_tol = args.logits_tol is None tol = ( @@ -956,29 +1259,56 @@ def main(): else args.logits_tol ) tol_desc = f"{tol:.6f}{' auto' if auto_tol else ''}" - out_dev = pipe.final_output().float() ref = RefModel(w1_bf, w2_bf, sw1, sw2, spec, dev) ref_out = ref.run(x0, routings).float() - logits_diff = _calc_diff(ref_out, out_dev) - errs = dist_ctx.allreduce_sum(0 if logits_diff < tol else 1) - avg_diff = dist_ctx.allreduce_avg_float(logits_diff) - if dist_ctx.rank == 0: - print( - f"# MEGA-CHECK layers={n_layers}: {'PASS' if errs == 0 else 'FAIL'} " - f"(avg logits_diff={avg_diff:.6f} over {dist_ctx.world} ranks, " - f"tol={tol_desc})", - flush=True, - ) - if errs != 0: - accuracy_failure = ( - f"MegaMoE accuracy check failed on {errs}/{dist_ctx.world} ranks: " - f"average logits_diff={avg_diff:.6f}, tolerance={tol_desc}" - ) + for row in summary_rows: + combine_mode = row["combine"] + logits_diff = _calc_diff(ref_out, outputs[combine_mode]) + errs = dist_ctx.allreduce_sum(0 if logits_diff < tol else 1) + avg_diff = dist_ctx.allreduce_avg_float(logits_diff) + row["logits_diff"] = avg_diff + row["logits_tol"] = tol + row["accuracy"] = "PASS" if errs == 0 else "FAIL" + if dist_ctx.rank == 0: + print( + f"# MEGA-CHECK combine={combine_mode} layers={n_layers}: " + f"{'PASS' if errs == 0 else 'FAIL'} " + f"(avg logits_diff={avg_diff:.6f} over {dist_ctx.world} ranks, " + f"tol={tol_desc})", + flush=True, + ) + if errs != 0: + failures.append( + f"combine={combine_mode} failed on {errs}/{dist_ctx.world} " + f"ranks: average logits_diff={avg_diff:.6f}, " + f"tolerance={tol_desc}" + ) + + # The summary goes last so every row carries BOTH its perf and its accuracy. + # With more than one mode the rows line up column by column, and speedup_vs_base + # spells out the one comparison the table exists for. + if len(summary_rows) > 1: + base_median = next( + (r["median_us"] for r in summary_rows if r["combine"] == "base"), None + ) + if base_median: + for row in summary_rows: + row["speedup_vs_base"] = base_median / row["median_us"] + # Needs both modes' kernel tables, so it only exists under --profile_table. + # It describes what fused did with base's stage 2, so it belongs on the + # fused row; base is the 0.0 baseline it is measured against. + overlap = _stage2_overlap_rate(kernel_rows, idim) + if overlap is not None: + for row in summary_rows: + row["stage2_overlap_rate"] = ( + 0.0 if row["combine"] == "base" else overlap + ) + if dist_ctx.rank == 0: + _emit_table("mega_moe summary", summary_rows) - pipe.teardown() dist_ctx.shutdown() - if accuracy_failure is not None: - raise AssertionError(accuracy_failure) + if failures: + raise AssertionError("MegaMoE accuracy check failed -- " + "; ".join(failures)) def _parse_args(): @@ -1006,11 +1336,28 @@ def _parse_args(): p.add_argument("-k", "--topk", type=int, default=6, help="top-k") p.add_argument("--shared_experts", type=int, default=0, help="dense shared experts") p.add_argument("--layers", type=int, default=61, help="number of MoE layers") + # The shared ubench data-init knobs: --data-init, --scale-init and --seed. + # default_dist=norm reproduces the historical N(0,1)*0.1 weights, which is + # what the _ACC_TOL budget was calibrated against. + add_data_init_args(p, default_dist="norm", default_scale=_DEFAULT_SCALE_INIT) p.add_argument( - "--seed", + "--warmup", type=int, - default=0, - help="base RNG seed for weights/tokens/routing (optional; default 0)", + default=5, + help="untimed graph replays before the timed ones", + ) + p.add_argument( + "--iters", + type=int, + default=1, + help="timed graph replays; each one is a full --layers chain, reported as " + "min/median/mean/max", + ) + p.add_argument( + "--prof_replays", + type=int, + default=3, + help="graph replays profiled for the --profile_table breakdown", ) p.add_argument( "--logits_tol", @@ -1023,7 +1370,7 @@ def _parse_args(): "--acc_verify", type=int, default=1, help="run fp32 reference accuracy check" ) p.add_argument( - "--profile_table", type=int, default=0, help="print per-kernel table" + "--profile_table", type=int, default=1, help="print per-kernel table" ) p.add_argument( "--save_trace", @@ -1045,10 +1392,12 @@ def _parse_args(): p.add_argument( "--combine", type=str, - choices=["base", "fused"], - default=os.environ.get("COMBINE", "base"), + choices=["base", "fused", "both"], + default=os.environ.get("COMBINE", "both"), help="EP combine mode: base (mori v2 dispatch/combine around fused_moe) " - "| fused (gemm2-fused P2P scatter; mxfp4 only). Falls back to $COMBINE.", + "| fused (gemm2-fused P2P scatter; mxfp4 only) | both (run base then " + "fused in one process and compare them row by row in the summary). " + "Falls back to $COMBINE.", ) return p.parse_args() diff --git a/op_tests/op_benchmarks/triton/bench_deepgemm_attention.py b/op_tests/op_benchmarks/triton/bench_deepgemm_attention.py index 7efa3fea0d..703ccd239e 100644 --- a/op_tests/op_benchmarks/triton/bench_deepgemm_attention.py +++ b/op_tests/op_benchmarks/triton/bench_deepgemm_attention.py @@ -8,6 +8,8 @@ import torch import triton +from aiter.benchmark_data_init import DATA_DISTS, fill, make_generator +from aiter.benchmark_reporting import print_json_table from aiter.ops.shuffle import shuffle_weight from aiter.ops.triton.attention.pa_mqa_logits import ( deepgemm_fp8_paged_mqa_logits, @@ -189,16 +191,18 @@ def create_paged_mqa_logits_configs(args: argparse.Namespace): return configs -def run_benchmark(args: argparse.Namespace): +def run_benchmark(args: argparse.Namespace, data_init: str = "norm"): ChunkK = 128 WavePerEU = 5 + rows = [] @triton.testing.perf_report(create_paged_mqa_logits_configs(args)) def test_deepgemm_fp8_paged_mqa_logits( batch_size, next_n, heads, index_dim, avg_kv_length, kv_storage_kind ): - torch.manual_seed(0) - random.seed(0) + torch.manual_seed(args.seed) + random.seed(args.seed) + gen = make_generator(args.seed) max_model_len = 2 * avg_kv_length blocksize = args.blocksize if args.kv_preshuffle else 1 @@ -225,19 +229,22 @@ def test_deepgemm_fp8_paged_mqa_logits( ) prefix_sum_context_lens[1:] = torch.cumsum(context_lens, dim=0) - q = torch.randn( - (batch_size, next_n, heads, index_dim), - device="cuda", + q = fill( + (batch_size * next_n * heads, index_dim), + data_init, + gen, dtype=torch.bfloat16, - ) - kv_cache = torch.randn( - (num_blocks, blocksize, 1, index_dim), - device="cuda", + ).view(batch_size, next_n, heads, index_dim) + kv_cache = fill( + (num_blocks * blocksize, index_dim), + data_init, + gen, dtype=torch.bfloat16, - ) - weights = torch.randn( + ).view(num_blocks, blocksize, 1, index_dim) + weights = fill( (batch_size * next_n, heads), - device="cuda", + data_init, + gen, dtype=torch.float32, ) @@ -273,7 +280,7 @@ def test_deepgemm_fp8_paged_mqa_logits( for i in range(batch_size): ctx_len = int(context_lens[i].item()) kv_indices[prefix_sum_context_lens[i] : prefix_sum_context_lens[i + 1]] = ( - torch.randperm(max_model_len, device="cuda")[:ctx_len] + torch.randperm(max_model_len, device="cuda", generator=gen)[:ctx_len] ) if kv_storage_kind == "non_ragged_k": @@ -369,6 +376,10 @@ def test_deepgemm_fp8_paged_mqa_logits( def calc_diff(x: torch.Tensor, y: torch.Tensor): x, y = x.double(), y.double() denominator = (x * x + y * y).sum() + # zero-init makes both logits tensors exactly zero. Treat that + # exact match as zero error instead of reporting 0/0 -> NaN. + if denominator == 0: + return torch.zeros_like(denominator) sim = 2 * (x * y).sum() / denominator return 1 - sim @@ -410,9 +421,26 @@ def calc_diff(x: torch.Tensor, y: torch.Tensor): os.system("zip -r paged_mqa_logits_aot_kernel paged_mqa_logits") + rows.append( + { + "data_init": data_init, + "seed": args.seed, + "batch": batch_size, + "next_n": next_n, + "heads": heads, + "index_dim": index_dim, + "avg_kv_len": avg_kv_length, + "kv_storage": kv_storage_kind, + "blocksize": blocksize, + "latency_us": elapsed_us, + "TFLOPS": flops, + "logits_diff": float(logits_diff), + } + ) return flops - test_deepgemm_fp8_paged_mqa_logits.run(print_data=True) + test_deepgemm_fp8_paged_mqa_logits.run(print_data=False) + print_json_table("paged_mqa_logits summary", rows) if __name__ == "__main__": @@ -474,6 +502,21 @@ def calc_diff(x: torch.Tensor, y: torch.Tensor): action="store_true", help="Disable varctx schedule (only applies with --kv_preshuffle)", ) + parser.add_argument( + "--data-init", + nargs="+", + choices=list(DATA_DISTS), + default=["norm"], + help="DATA initialization distribution(s) for Q, KV and weights", + ) + parser.add_argument( + "--seed", + type=int, + default=0, + help="RNG seed for input data and generated index tables (default: 0)", + ) args = parser.parse_args() - run_benchmark(args) + for data_init in args.data_init: + print(f"data_init={data_init} seed={args.seed}") + run_benchmark(args, data_init=data_init) diff --git a/op_tests/test_f4gemm.py b/op_tests/test_f4gemm.py index eb2a45589a..bfa792e70a 100644 --- a/op_tests/test_f4gemm.py +++ b/op_tests/test_f4gemm.py @@ -30,19 +30,22 @@ import aiter from aiter import dtypes +from aiter.benchmark_data_init import ( + fill_fp4, + fill_scale_e4m3, + fill_scale_e8m0, + make_generator, +) from aiter.jit.utils.chip_info import get_gfx_runtime as get_gfx from aiter.ops.gemm_op_a4w4 import MXFP8_OUT_SCALE_BLOCK, unpack_mxfp8_out_scale from aiter.ops.shuffle import shuffle_scale_f4, shuffle_weight_f4 -from aiter.test_common import benchmark, checkAllclose, run_perftest +from aiter.test_common import ( + benchmark, + checkAllclose, + run_perftest, +) from aiter.utility import fp4_utils -try: - import bench_init -except ImportError as e: - if e.name != "bench_init": - raise - from op_tests import bench_init - torch.set_default_device("cuda") torch.set_printoptions(sci_mode=False) pd.set_option("display.max_columns", 30) @@ -249,21 +252,14 @@ def run_torch_nvfp4(xq, wq, xs, ws, gA, gB): def _prep_mxfp4(M, N, K, apre, data_init, scale_init, gen): # DATA (fp4 e2m1, packed 2/byte). data & scale are sampled *independently*. - if data_init == "constant": - # f4gemm.cpp data_init=0: A=0x22, B=0x33 (fixed representable e2m1). - xq = torch.full((M, K // 2), 0x22, dtype=torch.uint8) - wq = torch.full((N, K // 2), 0x33, dtype=torch.uint8) - else: # uniform / gaussian / trig / random - xq = bench_init.fill_fp4((M, K), data_init, gen) - wq = bench_init.fill_fp4((N, K), data_init, gen) - # SCALE (e8m0 per-32). auto -> pow2_binomial for E8M0. - if scale_init == "constant": - # neutral e8m0 scale 0x7F (exp 0 -> 2^0 = 1.0). - xs = torch.full((M, K // MXFP4_SCALE_BLOCK), 0x7F, dtype=torch.uint8) - ws = torch.full((N, K // MXFP4_SCALE_BLOCK), 0x7F, dtype=torch.uint8) - else: # auto / pow2_binomial / random - xs = bench_init.fill_scale_e8m0((M, K // MXFP4_SCALE_BLOCK), scale_init, gen) - ws = bench_init.fill_scale_e8m0((N, K // MXFP4_SCALE_BLOCK), scale_init, gen) + # constant is built inside fill_fp4 via constant=; f4gemm.cpp data_init=0 + # keeps A=0x22, B=0x33 (fixed representable e2m1). + xq = fill_fp4((M, K), data_init, gen, constant=0x22) + wq = fill_fp4((N, K), data_init, gen, constant=0x33) + # SCALE (e8m0 per-32). auto -> pow2_binomial for E8M0; constant default is + # the neutral e8m0 byte 0x7F (exp 0 -> 2^0 = 1.0). + xs = fill_scale_e8m0((M, K // MXFP4_SCALE_BLOCK), scale_init, gen) + ws = fill_scale_e8m0((N, K // MXFP4_SCALE_BLOCK), scale_init, gen) ref = run_torch_mxfp4(xq, wq, xs, ws) inp = { "A": shuffle_weight_f4(xq) if apre else xq, @@ -277,23 +273,16 @@ def _prep_mxfp4(M, N, K, apre, data_init, scale_init, gen): def _prep_nvfp4(M, N, K, apre, data_init, scale_init, gen): - # DATA (fp4 e2m1). data & scale sampled independently (bench_init). - if data_init == "constant": - # f4gemm.cpp data_init=0: A=0x22, B=0x33 (fixed representable e2m1). - xq = torch.full((M, K // 2), 0x22, dtype=torch.uint8) - wq = torch.full((N, K // 2), 0x33, dtype=torch.uint8) - else: # uniform / gaussian / trig / random - xq = bench_init.fill_fp4((M, K), data_init, gen) - wq = bench_init.fill_fp4((N, K), data_init, gen) - # SCALE (e4m3 per-16). auto -> gaussian(0.34375,0.08) for E4M3. - if scale_init == "constant": - # neutral e4m3 scale 0x38 (exp 7 = bias -> 1.0). - xs = torch.full((M, K // NVFP4_SCALE_BLOCK), 0x38, dtype=torch.uint8) - ws = torch.full((N, K // NVFP4_SCALE_BLOCK), 0x38, dtype=torch.uint8) - else: # auto / gaussian / random - xs = bench_init.fill_scale_e4m3((M, K // NVFP4_SCALE_BLOCK), scale_init, gen) - ws = bench_init.fill_scale_e4m3((N, K // NVFP4_SCALE_BLOCK), scale_init, gen) - # Per-tensor global scale is NOT part of bench_init: keep neutral. + # DATA (fp4 e2m1). data & scale sampled independently (test_common). + # constant built inside fill_fp4 via constant=; f4gemm.cpp data_init=0 keeps + # A=0x22, B=0x33 (fixed representable e2m1). + xq = fill_fp4((M, K), data_init, gen, constant=0x22) + wq = fill_fp4((N, K), data_init, gen, constant=0x33) + # SCALE (e4m3 per-16). auto -> gaussian(0.34375,0.08) for E4M3; constant + # default is the neutral e4m3 byte 0x38 (exp 7 = bias -> 1.0). + xs = fill_scale_e4m3((M, K // NVFP4_SCALE_BLOCK), scale_init, gen) + ws = fill_scale_e4m3((N, K // NVFP4_SCALE_BLOCK), scale_init, gen) + # Per-tensor global scale is NOT part of the init helpers: keep neutral. gA = gB = 1.0 ref = run_torch_nvfp4(xq, wq, xs, ws, gA, gB) inp = { @@ -320,6 +309,10 @@ def test_gemm( seed=0, mode="perf", knl_name=None, + num_warmup=2, + num_iters=None, + test_graph=False, + num_rotate=0, ): # Skip unsupported combos up front (before prep/shuffle) so they show as # "not support" rather than crashing on a shape assert. @@ -341,6 +334,8 @@ def test_gemm( return { "gfx": get_gfx(), "knl_name": actual_knl, + "tile": "256x256", + "cluster": "4x4", "asm us": float("nan"), "asm TFLOPS": float("nan"), "asm TB/s": float("nan"), @@ -352,7 +347,7 @@ def test_gemm( assert K % block == 0, f"K must be a multiple of {block}" out_fp8 = outtype == "fp8" out_dtype = _OUT_DTYPE[outtype] - gen = bench_init.make_generator(seed) # fixed seed -> bit-identical buffers + gen = make_generator(seed) # fixed seed -> bit-identical buffers prep = _prep_mxfp4 if intype == "mxfp4" else _prep_nvfp4 inp, ref_f32 = prep(M, N, K, apre, data_init, scale_init, gen) # Reference in the kernel's output form: block-scaled (fp8 e4m3 data + e8m0 @@ -362,7 +357,8 @@ def test_gemm( else: ref = ref_f32.to(out_dtype) needTrace = mode == "profile" - num_iters = 5 if mode == "func" else 101 + # --iters overrides; unset keeps the mode default (func=5, perf/profile=101). + num_iters = num_iters if num_iters is not None else (5 if mode == "func" else 101) # Kernel/.co base name for this config (used for logging, and to derive the # mangled knl_name when an explicit dispatch is requested). See @@ -432,7 +428,11 @@ def run_asm(A, B, sA, sB): # the verbatim knl_name otherwise (kept in the table, see main()). actual_knl = knl_name if (knl_name and knl_name != "auto") else base ret = {"gfx": get_gfx(), "knl_name": actual_knl} - # F4GEMM tiles are always 256x256 (see f4gemm.csv). Report TG occupancy. + # Structured algo details (f4gemm.csv columns): F4GEMM tiles are always + # 256x256 with a 4x4 (cluster_x x cluster_y) cluster; no splitk/unroll axis. + ret["tile"] = "256x256" + ret["cluster"] = "4x4" + # Report TG occupancy for the 256x256 tile. _report_active_tg(M, N, 256, 256, base) # Only a missing .co is reported as "not support"; any other failure (OOM, # memory fault, shape assert, ...) must propagate, not show as a green cell. @@ -444,7 +444,13 @@ def run_asm(A, B, sA, sB): for name, (fn, fn_args) in candidates.items(): try: out, us = run_perftest( - fn, *fn_args, num_iters=num_iters, needTrace=needTrace + fn, + *fn_args, + num_iters=num_iters, + num_warmup=num_warmup, + testGraph=test_graph, + num_rotate_args=num_rotate, + needTrace=needTrace, ) except Exception as e: if not any(m in str(e) for m in _NOT_SUPPORTED_MARKERS): @@ -586,34 +592,34 @@ def main(): parser.add_argument( "--data-init", dest="data_init", - nargs="*", - choices=["constant", "uniform", "gaussian", "trig", "random"], + nargs="+", + choices=["zero", "constant", "uniform", "norm"], default=None, - help="DATA init distribution(s) (mblas-style; sampled independently of scale).\n" + help="DATA init distribution(s) (sampled independently of scale).\n" "Paired position-wise with --scale-init (length-1 broadcasts).\n" "Default (unset): perf/profile = 'constant uniform', func = 'uniform'\n" "(func drops constant: its exact-boundary values trigger e8m0/e4m3\n" "edge rounding that shows as spurious warnings).\n" + " zero = all-zero e2m1 codes\n" + " constant = A=0x22, B=0x33 (deterministic)\n" " uniform = FP4 U(-3,3)\n" - " gaussian = N(0,1) [norm-dist / LLM-like]\n" - " trig = trig_float in [-2,2] [optimistic pattern]\n" - " random = pure random e2m1 codes [overly pessimistic]\n" - " constant = A=0x22, B=0x33 (deterministic)", + " norm = N(0,1) [norm-dist / LLM-like]", ) parser.add_argument( "--scale-init", dest="scale_init", - nargs="*", - choices=["auto", "pow2_binomial", "gaussian", "random", "constant"], + nargs="+", + choices=["auto", "pow2_binomial", "zero", "constant", "uniform", "norm"], default=None, help="SCALE init distribution(s) (by scale format)\n" "Default (unset): perf/profile = 'constant auto', func = 'auto'\n" " auto = format-recommended: mxfp4/E8M0 -> pow2_binomial,\n" " nvfp4/E4M3 -> gaussian(0.34375,0.08)\n" " pow2_binomial = 2^(Binomial(21,0.5)-11) [E8M0 only]\n" - " gaussian = N(0.34375,0.08) [E4M3 only]\n" - " random = random on-wire byte, modest range\n" - " constant = neutral scale (2^0 = 1.0)", + " zero = all-zero scale bytes\n" + " constant = neutral scale (2^0 = 1.0)\n" + " uniform = U(0.5,2) -> nearest on-wire byte\n" + " norm = N(1,0.25) -> nearest on-wire byte", ) parser.add_argument( "--seed", @@ -621,6 +627,39 @@ def main(): default=0, help="RNG seed; same seed -> bit-identical data/scale buffers", ) + parser.add_argument( + "--warmup", + type=int, + default=2, + help="warmup iterations before timing (run_perftest num_warmup)", + ) + parser.add_argument( + "--iters", + type=int, + default=None, + help="timed iterations (run_perftest num_iters); unset -> mode default " + "(func=5, perf/profile=101)", + ) + parser.add_argument( + "--graph", + action="store_true", + help="also time via HIP graph replay (run_perftest testGraph), " + "minimizing inter-kernel gaps", + ) + parser.add_argument( + "--rotate", + type=int, + default=0, + help="rotating input-buffer copies to defeat the L2 hot-cache " + "(run_perftest num_rotate_args); 0 = auto-size from L2", + ) + parser.add_argument( + "--json", + dest="json_out", + default=None, + help="also write the full summary table (all columns) as JSON records " + "to this path, for CI/regression", + ) parser.add_argument( "--knl-name", dest="knl_name", @@ -695,14 +734,38 @@ def main(): seed=args.seed, mode=args.mode, knl_name=args.knl_name, + num_warmup=args.warmup, + num_iters=args.iters, + test_graph=args.graph, + num_rotate=args.rotate, ) for apre, (di, si), intype, outtype, (M, N, K) in itertools.product( apre_list, init_pairs, args.intype, args.outtype, shapes ) ] - df = pd.DataFrame(rows) - # Keep knl_name (the actual .co); drop the columns constant within a table. - df = df.drop(columns=["seed", "gfx", "mode"], errors="ignore") + df_full = pd.DataFrame(rows) + # JSON keeps every column (config + algo details + results) so each record is + # self-describing for CI/regression; the markdown table below drops columns + # that are constant within a run for readability. + if args.json_out: + df_full.to_json(args.json_out, orient="records", indent=2) + aiter.logger.info( + "wrote JSON summary (%d rows) to %s", len(df_full), args.json_out + ) + # Keep knl_name (the actual .co) + tile; drop columns constant within a table + # (cluster is always 4x4). + df = df_full.drop( + columns=[ + "seed", + "gfx", + "mode", + "num_warmup", + "num_iters", + "test_graph", + "num_rotate", + ], + errors="ignore", + ) aiter.logger.info( "gemm_a4w4 (F4GEMM) summary (markdown):\n%s", df.to_markdown(index=False), diff --git a/op_tests/test_flydsl_qk_norm_rope_quant.py b/op_tests/test_flydsl_qk_norm_rope_quant.py index ff354f1cfb..28785b3992 100755 --- a/op_tests/test_flydsl_qk_norm_rope_quant.py +++ b/op_tests/test_flydsl_qk_norm_rope_quant.py @@ -17,17 +17,20 @@ python op_tests/test_flydsl_qk_norm_rope_quant.py python op_tests/test_flydsl_qk_norm_rope_quant.py -T 64 256 1024 -q fp8_1x128_e8m0 python op_tests/test_flydsl_qk_norm_rope_quant.py --no-quant # bf16 only + python op_tests/test_flydsl_qk_norm_rope_quant.py --data-init zero + python op_tests/test_flydsl_qk_norm_rope_quant.py --seed 42 --data-init uniform """ import argparse import itertools import math -import pandas as pd import torch import aiter from aiter import dtypes +from aiter.benchmark_data_init import add_data_init_args, fill, make_generator +from aiter.benchmark_reporting import print_json_table from aiter.jit.utils.chip_info import get_gfx from aiter.ops.flydsl import flydsl_qk_norm_rope_quant from aiter.test_common import benchmark, checkAllclose, run_perftest @@ -45,7 +48,6 @@ _FP8_DTYPE = dtypes.fp8 _FP8_MAX = float(torch.finfo(_FP8_DTYPE).max) - # ============================================================================ # Reference (pure torch) # ============================================================================ @@ -219,9 +221,12 @@ def test_flydsl_qk_norm_rope_quant( scale_dtype, q_weighted, quant, + seed=0, + data_init="norm", ): - torch.manual_seed(0) + torch.manual_seed(seed) device = torch.device("cuda") + generator = make_generator(seed, device=device) # Build cos/sin via a YaRN-style table covering all positions in T. max_pos = max(T, 64) @@ -231,10 +236,28 @@ def test_flydsl_qk_norm_rope_quant( cos = freqs.cos().to(torch.bfloat16).contiguous() sin = freqs.sin().to(torch.bfloat16).contiguous() - q = torch.randn(T, H * D, dtype=torch.bfloat16, device=device) * 0.1 + q = fill( + (T, H * D), + data_init, + generator, + dtype=torch.bfloat16, + device=device, + uniform=(-0.1, 0.1), + ) + if data_init == "norm": + q.mul_(0.1) # Mimic V4 KV split: kv = strided view into a wider tensor Q_LORA = 1536 - qkv_a = torch.randn(T, Q_LORA + D, dtype=torch.bfloat16, device=device) * 0.1 + qkv_a = fill( + (T, Q_LORA + D), + data_init, + generator, + dtype=torch.bfloat16, + device=device, + uniform=(-0.1, 0.1), + ) + if data_init == "norm": + qkv_a.mul_(0.1) _, kv = torch.split(qkv_a, [Q_LORA, D], dim=-1) kv_w = torch.randn(D, dtype=torch.bfloat16, device=device).abs() + 0.5 q_w = ( @@ -525,9 +548,10 @@ def _build_swa_case(T, mode, *, device): @benchmark() -def test_flydsl_swa_write(T, H, D, RD, mode): - torch.manual_seed(0) +def test_flydsl_swa_write(T, H, D, RD, mode, *, seed=0, data_init="norm"): + torch.manual_seed(seed) device = torch.device("cuda") + generator = make_generator(seed, device=device) bid, pos, index_t, num_rows, dest = _build_swa_case(T, mode, device=device) max_pos = int(pos.max().item()) + 4 @@ -538,11 +562,29 @@ def test_flydsl_swa_write(T, H, D, RD, mode): cos = freqs.cos().to(torch.bfloat16).contiguous() sin = freqs.sin().to(torch.bfloat16).contiguous() - q = torch.randn(T, H * D, dtype=torch.bfloat16, device=device) * 0.1 + q = fill( + (T, H * D), + data_init, + generator, + dtype=torch.bfloat16, + device=device, + uniform=(-0.1, 0.1), + ) + if data_init == "norm": + q.mul_(0.1) # Mimic the V4 KV split: kv is a strided view into a wider tensor, exactly # as the model hands it over. Q_LORA = 1536 - qkv_a = torch.randn(T, Q_LORA + D, dtype=torch.bfloat16, device=device) * 0.1 + qkv_a = fill( + (T, Q_LORA + D), + data_init, + generator, + dtype=torch.bfloat16, + device=device, + uniform=(-0.1, 0.1), + ) + if data_init == "norm": + qkv_a.mul_(0.1) _, kv = torch.split(qkv_a, [Q_LORA, D], dim=-1) kv_w = torch.randn(D, dtype=torch.bfloat16, device=device).abs() + 0.5 @@ -550,7 +592,14 @@ def test_flydsl_swa_write(T, H, D, RD, mode): # a write; they change no in-pool byte, so only a dirtied guard row can show # that one regressed. G = _SWA_GUARD_ROWS - pool = torch.zeros(G + num_rows + G, D, dtype=torch.bfloat16, device=device) + # A zero-filled pool cannot distinguish "the kernel wrote a zero row" from + # "the kernel skipped this row" when --data-init zero. Seed every row with an + # unreachable sentinel so the write-mask check remains valid for all data + # distributions. + sentinel = torch.finfo(torch.bfloat16).max + pool = torch.full( + (G + num_rows + G, D), sentinel, dtype=torch.bfloat16, device=device + ) swa_kv = pool[G : G + num_rows] mode_kw = ( {"swa_dest_rows": index_t} @@ -598,7 +647,7 @@ def test_flydsl_swa_write(T, H, D, RD, mode): # The scatter is a verbatim copy of kv_out, so the reference IS kv_out # gathered onto the rows the addressing mode selects. - expected = torch.zeros_like(swa_kv) + expected = torch.full_like(swa_kv, sentinel) n_written = 0 for t in range(T): if dest[t] < 0: @@ -620,10 +669,12 @@ def test_flydsl_swa_write(T, H, D, RD, mode): ) # A skipped token must reach NO row, not merely the right one. assert ( - int((swa_kv != 0).any(dim=1).sum()) == n_written + int((swa_kv != sentinel).any(dim=1).sum()) == n_written ), f"{mode}: a skipped token still reached the pool" - assert not pool[:G].any(), f"{mode}: scatter wrote BEFORE the pool" - assert not pool[G + num_rows :].any(), f"{mode}: scatter wrote PAST the pool" + assert (pool[:G] == sentinel).all(), f"{mode}: scatter wrote BEFORE the pool" + assert ( + pool[G + num_rows :] == sentinel + ).all(), f"{mode}: scatter wrote PAST the pool" # The scatter must not perturb the primary outputs. ref_q, ref_kv, _, _ = flydsl_qk_norm_rope_quant( @@ -742,6 +793,7 @@ def main(): action="store_true", help="bf16 only (ignore -q).", ) + add_data_init_args(parser, default_dist="norm", include_scale=False) args = parser.parse_args() # Smoke-test the advertised 4D cos/sin layout once before sweeping. @@ -749,10 +801,9 @@ def main(): quant_keys = ["bf16"] if args.no_quant else args.quant qweight_modes = [False, True] if args.qweight else [False] - rows = [] - for key, qw_mode, H, D, T in itertools.product( - quant_keys, qweight_modes, args.H, args.D, args.T + for data_init, key, qw_mode, H, D, T in itertools.product( + args.data_init, quant_keys, qweight_modes, args.H, args.D, args.T ): quant, group_size, scale_dtype, _ = _QUANT_OPTIONS[key] rows.append( @@ -765,17 +816,17 @@ def main(): scale_dtype=scale_dtype, q_weighted=qw_mode, quant=quant, + seed=args.seed, + data_init=data_init, ) ) - aiter.logger.info( - "flydsl_qk_norm_rope_quant summary (markdown):\n%s", - pd.DataFrame(rows).to_markdown(index=False), - ) + print_json_table("flydsl_qk_norm_rope_quant summary", rows) # Separate arg signature -> its own table (merging would scatter NaNs). # The scatter is decode-only and bf16-only; keep T in the decode range. swa_rows = [] - for mode, H, D, T in itertools.product( + for data_init, mode, H, D, T in itertools.product( + args.data_init, args.swa_mode, args.H, args.D, @@ -784,11 +835,18 @@ def main(): # reserved for the out-of-window sentinel). [t for t in args.T if 8 <= t <= 96] or [16, 64], ): - swa_rows.append(test_flydsl_swa_write(T, H, D, args.RD, mode)) - aiter.logger.info( - "flydsl_qk_norm_rope_quant fused SWA write summary (markdown):\n%s", - pd.DataFrame(swa_rows).to_markdown(index=False), - ) + swa_rows.append( + test_flydsl_swa_write( + T, + H, + D, + args.RD, + mode, + seed=args.seed, + data_init=data_init, + ) + ) + print_json_table("flydsl_qk_norm_rope_quant fused SWA write summary", swa_rows) if __name__ == "__main__": diff --git a/op_tests/test_gemm_a8w8_blockscale.py b/op_tests/test_gemm_a8w8_blockscale.py index 9281279bba..3517bbd495 100644 --- a/op_tests/test_gemm_a8w8_blockscale.py +++ b/op_tests/test_gemm_a8w8_blockscale.py @@ -15,7 +15,9 @@ from einops import repeat as eirp import aiter +from aiter import benchmark_data_init as bench_init from aiter import dtypes +from aiter.benchmark_reporting import print_json_table from aiter.ops.gemm_op_a8w8 import gemm_a8w8_blockscale_ck, gemm_a8w8_blockscale_cktile from aiter.ops.shuffle import shuffle_weight from aiter.test_common import benchmark, checkAllclose, perftest @@ -87,25 +89,53 @@ def run_triton(x, weightshuffle, x_scale, w_scale, dtype=dtypes.bf16, backend=No @benchmark() -def test_gemm(dtype, m, n, k, ck_preshuffle=True, use_flydsl=False): +def test_gemm( + dtype, + m, + n, + k, + ck_preshuffle=True, + use_flydsl=False, + data_init="uniform", + scale_init="auto", + seed=0, +): ret = {} block_shape_n, block_shape_k = block_shape scale_m = m scale_n = (n + block_shape_n - 1) // block_shape_n scale_k = (k + block_shape_k - 1) // block_shape_k - x = (torch.rand((m, k), dtype=dtypes.fp32, device="cuda") / 10).to(dtypes.fp8) - weight = (torch.rand((n, k), dtype=dtypes.fp32, device="cuda") / 10).to(dtypes.fp8) - x_scale = torch.rand([scale_m, scale_k], dtype=dtypes.fp32, device="cuda") - w_scale = torch.rand([scale_n, scale_k], dtype=dtypes.fp32, device="cuda") - use_flydsl_fp8_scale = use_flydsl and ck_preshuffle - if use_flydsl_fp8_scale: - FP8_E4M3_MAX = 448.0 - x_scale = fp4_utils.f32_to_mx_e8m0_scale( - x_scale * FP8_E4M3_MAX, dtype=fp4_utils.MxDtypeInt.FP8_E4M3 + generator = bench_init.make_generator(seed) + if data_init == "constant": + x = torch.full((m, k), 0.5, dtype=dtypes.fp32, device="cuda").to(dtypes.fp8) + weight = torch.full((n, k), 0.5, dtype=dtypes.fp32, device="cuda").to( + dtypes.fp8 + ) + else: + x = bench_init.fill_fp8((m, k), data_init, generator, dtype=dtypes.fp8) + weight = bench_init.fill_fp8((n, k), data_init, generator, dtype=dtypes.fp8) + + if scale_init == "constant": + x_scale_raw = torch.full( + (scale_m, scale_k), 0x7F, dtype=torch.uint8, device="cuda" + ) + w_scale_raw = torch.full( + (scale_n, scale_k), 0x7F, dtype=torch.uint8, device="cuda" ) - w_scale = fp4_utils.f32_to_mx_e8m0_scale( - w_scale * FP8_E4M3_MAX, dtype=fp4_utils.MxDtypeInt.FP8_E4M3 + else: + x_scale_raw = bench_init.fill_scale_e8m0( + (scale_m, scale_k), scale_init, generator + ) + w_scale_raw = bench_init.fill_scale_e8m0( + (scale_n, scale_k), scale_init, generator ) + use_flydsl_fp8_scale = use_flydsl and ck_preshuffle + if use_flydsl_fp8_scale: + x_scale = x_scale_raw.view(dtypes.fp8_e8m0) + w_scale = w_scale_raw.view(dtypes.fp8_e8m0) + else: + x_scale = fp4_utils.e8m0_to_f32(x_scale_raw) + w_scale = fp4_utils.e8m0_to_f32(w_scale_raw) a, _ = run_torch(x, weight, x_scale, w_scale, dtype) @@ -304,6 +334,30 @@ def test_splitk_correctness(m=4, n=2112, k=7168, dtype=dtypes.bf16, splitK=1): help="""N&K of mnk. e.g.: -nk 24576,1536""", ) +parser.add_argument( + "--data-init", + dest="data_init", + nargs="+", + choices=bench_init.DATA_DISTS, + default=None, + help="DATA initialization distribution(s), paired position-wise with " + "--scale-init (length-1 broadcasts). Default: constant uniform", +) +parser.add_argument( + "--scale-init", + dest="scale_init", + nargs="+", + choices=bench_init.E8M0_SCALE_DISTS, + default=None, + help="E8M0 SCALE initialization distribution(s), paired position-wise " + "with --data-init (length-1 broadcasts). Default: constant auto", +) +parser.add_argument( + "--seed", + type=int, + default=0, + help="RNG seed for input, weight, and scales (default: 0)", +) parser.add_argument( "--ck_preshuffle", type=dtypes.str2bool, @@ -344,6 +398,19 @@ def test_splitk_correctness(m=4, n=2112, k=7168, dtype=dtypes.bf16, splitK=1): args = parser.parse_args() +data_init_list = args.data_init or ["constant", "uniform"] +scale_init_list = args.scale_init or ["constant", "auto"] +if len(data_init_list) == 1: + data_init_list *= len(scale_init_list) +if len(scale_init_list) == 1: + scale_init_list *= len(data_init_list) +if len(data_init_list) != len(scale_init_list): + parser.error( + "--data-init and --scale-init must have equal length " + "(or length 1 to broadcast)" + ) +init_pairs = list(zip(data_init_list, scale_init_list)) + l_preshuffle = ( args.ck_preshuffle if isinstance(args.ck_preshuffle, list) else [args.ck_preshuffle] ) @@ -356,42 +423,40 @@ def test_splitk_correctness(m=4, n=2112, k=7168, dtype=dtypes.bf16, splitK=1): print(f"Loaded {len(shapes_df)} shapes from {args.csv}", flush=True) for dtype in args.dtype: for preshuffle in l_preshuffle: - for _, row in shapes_df.iterrows(): - ret = test_gemm( - dtype, - int(row["M"]), - int(row["N"]), - int(row["K"]), - ck_preshuffle=preshuffle, - use_flydsl=args.flydsl, - ) - df.append(ret) + for data_init, scale_init in init_pairs: + for _, row in shapes_df.iterrows(): + ret = test_gemm( + dtype, + int(row["M"]), + int(row["N"]), + int(row["K"]), + ck_preshuffle=preshuffle, + use_flydsl=args.flydsl, + data_init=data_init, + scale_init=scale_init, + seed=args.seed, + ) + df.append(ret) else: for dtype in args.dtype: for m in args.m: for n, k in args.nk: for ck_p in l_preshuffle: - ret = test_gemm( - dtype, m, n, k, ck_preshuffle=ck_p, use_flydsl=args.flydsl - ) - df.append(ret) - -df = pd.DataFrame(df) - -# Configure pandas to show all columns without truncation -pd.set_option("display.max_columns", None) -pd.set_option("display.width", None) -pd.set_option("display.max_colwidth", None) -pd.set_option("display.expand_frame_repr", False) - -print("\n" + "=" * 150) -print("COMPLETE PERFORMANCE SUMMARY (All Columns)") -print("=" * 150) -print(df.to_string(index=False)) -print("=" * 150) - -df_md = df.to_markdown(index=False) -aiter.logger.info("gemm_a8w8_blockscale summary (markdown):\n%s", df_md) + for data_init, scale_init in init_pairs: + ret = test_gemm( + dtype, + m, + n, + k, + ck_preshuffle=ck_p, + use_flydsl=args.flydsl, + data_init=data_init, + scale_init=scale_init, + seed=args.seed, + ) + df.append(ret) + +print_json_table("gemm_a8w8_blockscale summary", df) # Correctness check: verify split-K produces matching results print("\nRunning split-K correctness checks ...") diff --git a/op_tests/test_inverse_rope_group_quant.py b/op_tests/test_inverse_rope_group_quant.py index f1272cbe79..91adb7f908 100644 --- a/op_tests/test_inverse_rope_group_quant.py +++ b/op_tests/test_inverse_rope_group_quant.py @@ -16,11 +16,12 @@ import sys from collections import namedtuple -import pandas as pd import torch import aiter from aiter import dtypes +from aiter.benchmark_data_init import DATA_DISTS, fill, make_generator +from aiter.benchmark_reporting import print_json_table from aiter.jit.utils.chip_info import get_gfx from aiter.ops.inverse_rope_group_quant import ( SCALE_LAYOUTS, @@ -295,7 +296,7 @@ def _check_scale_layout(scale, s, g, ks, scale_layout, group_size, name): ), f"{name}: {scale_layout} scale should be {expect}, got {tuple(scale.shape)}" -def _make_inputs(s, h, head_dim, rd, dtype, seed=0): +def _make_inputs(s, h, head_dim, rd, dtype, data_init="norm", seed=0): """Build (o, positions, cos, sin) for one config. cos/sin are the 2D [max_pos, rd//2] the op takes. A model holding the @@ -305,12 +306,16 @@ def _make_inputs(s, h, head_dim, rd, dtype, seed=0): own call site, the way run_inverse_rope_inplace does for the triton rope. Shared by the sweep and the graph check so the two cannot drift. """ - torch.manual_seed(seed) + gen = make_generator(seed) positions = torch.arange(s, dtype=dtypes.i64) % MAX_POS # /10 keeps a group's amax away from fp8 saturation, like a real # post-softmax attention output. - o = torch.randn((s, h, head_dim), dtype=dtype) / 10 - theta = torch.randn((MAX_POS, rd // 2), dtype=dtypes.fp32) + o = ( + fill((s * h, head_dim), data_init, gen, dtype=dtype) + .view(s, h, head_dim) + .div_(10) + ) + theta = fill((MAX_POS, rd // 2), data_init, gen, dtype=dtypes.fp32) cos = torch.cos(theta).to(dtype).contiguous() sin = torch.sin(theta).to(dtype).contiguous() return o, positions, cos, sin @@ -440,12 +445,23 @@ def run_unfused(x, positions, cos, sin, num_groups, quant_group_size, rd, out): @benchmark() def test_inverse_rope_group_quant( - s, h, g, head_dim, rd, group_size, dtype, scale_layout + s, + h, + g, + head_dim, + rd, + group_size, + dtype, + scale_layout, + data_init="norm", + seed=0, ): d = h * head_dim // g scale_n = d // group_size - o, positions, cos, sin = _make_inputs(s, h, head_dim, rd, dtype) + o, positions, cos, sin = _make_inputs( + s, h, head_dim, rd, dtype, data_init=data_init, seed=seed + ) ref = run_torch(o, positions, cos, sin, g, group_size, rd) ref_rt = run_torch(o, positions, cos, sin, g, group_size, rd, roundtrip=True) @@ -546,14 +562,27 @@ def unfused_once(): return ret -def check_graph(s, h, g, head_dim, rd, group_size, dtype, scale_layout): +def check_graph( + s, + h, + g, + head_dim, + rd, + group_size, + dtype, + scale_layout, + data_init="norm", + seed=0, +): """Capture the op in a HIP graph, replay on fresh data, compare against eager. Not part of the perf table: this is a pass/fail check that the host-side dispatch tier and the pre-allocated buffers survive capture/replay. """ d = h * head_dim // g - o, positions, cos, sin = _make_inputs(s, h, head_dim, rd, dtype) + o, positions, cos, sin = _make_inputs( + s, h, head_dim, rd, dtype, data_init=data_init, seed=seed + ) x_fp8, x_scale = _alloc_outputs(s, g, d, group_size, scale_layout=scale_layout) kwargs = { "num_groups": g, @@ -575,7 +604,9 @@ def check_graph(s, h, g, head_dim, rd, group_size, dtype, scale_layout): inverse_rope_group_quant_cpp(o, positions, cos, sin, **kwargs) # Replay on new data, then compare against an eager run on the same data. - o2, positions2, cos2, sin2 = _make_inputs(s, h, head_dim, rd, dtype, seed=7) + o2, positions2, cos2, sin2 = _make_inputs( + s, h, head_dim, rd, dtype, data_init=data_init, seed=seed + 7 + ) o.copy_(o2) positions.copy_(positions2) cos.copy_(cos2) @@ -740,6 +771,19 @@ def main(): help="""Also run the HIP-graph capture/replay check over the same sweep. e.g.: --graph -s 1 4 32 128 300 512 700 2048""", ) + parser.add_argument( + "--data-init", + nargs="+", + choices=list(DATA_DISTS), + default=["norm"], + help="DATA initialization distribution(s) (default: norm)", + ) + parser.add_argument( + "--seed", + type=int, + default=0, + help="RNG seed for o and the RoPE cache source (default: 0)", + ) parser.add_argument( "--opus-tree", default=os.environ.get("AITER_OPUS_TREE"), @@ -756,13 +800,22 @@ def main(): for dtype in args.dtype: df = [] - for (h, g), s, head_dim, rd, group_size, scale_layout in itertools.product( + for ( + (h, g), + s, + head_dim, + rd, + group_size, + scale_layout, + data_init, + ) in itertools.product( args.hg, args.tokens, args.head_dim, args.rope_dim, args.group_size, args.scale_layout, + args.data_init, ): # n32k4 only exists at group 32: its four packed k groups are one # WMMA-K=128 step, so 4 * group_size has to be 128. The op rejects @@ -770,16 +823,32 @@ def main(): if scale_layout == "n32k4" and group_size != 32: continue ret = test_inverse_rope_group_quant( - s, h, g, head_dim, rd, group_size, dtype, scale_layout + s, + h, + g, + head_dim, + rd, + group_size, + dtype, + scale_layout, + data_init=data_init, + seed=args.seed, ) df.append(ret) if args.graph: - check_graph(s, h, g, head_dim, rd, group_size, dtype, scale_layout) - df = pd.DataFrame(df) - aiter.logger.info( - "inverse_rope_group_quant summary (markdown):\n%s", - df.to_markdown(index=False), - ) + check_graph( + s, + h, + g, + head_dim, + rd, + group_size, + dtype, + scale_layout, + data_init=data_init, + seed=args.seed, + ) + print_json_table("inverse_rope_group_quant summary", df) if args.graph: aiter.logger.info("all graph capture/replay checks passed") diff --git a/op_tests/test_mhc.py b/op_tests/test_mhc.py index 86405c78c3..4f0aa3a8eb 100644 --- a/op_tests/test_mhc.py +++ b/op_tests/test_mhc.py @@ -4,11 +4,12 @@ import argparse -import pandas as pd import torch import aiter from aiter import dtypes +from aiter.benchmark_data_init import add_data_init_args, fill, make_generator +from aiter.benchmark_reporting import print_json_table from aiter.jit.utils.chip_info import get_gfx_runtime from aiter.test_common import ( benchmark, @@ -37,7 +38,6 @@ TRITON_MHC_POST_PRE_MAX_M = 4096 torch.set_default_device("cuda") -# torch.cuda.manual_seed_all(0) # torch.set_printoptions(precision=3, linewidth=200, sci_mode=False) @@ -460,7 +460,15 @@ def mhc_pre_norm_split_hip( @benchmark() def test_mhc_pre( - m, hidden_size, hc_mult, test_hc_head=False, fuse_rmsnorm=False, fn_pack_bf16=False + m, + hidden_size, + hc_mult, + test_hc_head=False, + fuse_rmsnorm=False, + fn_pack_bf16=False, + dtype=dtypes.bf16, + data_init="norm", + seed=0, ): if fuse_rmsnorm and test_hc_head: raise ValueError("fuse_rmsnorm and hc_head are mutually exclusive") @@ -468,13 +476,14 @@ def test_mhc_pre( hc_mult2 = hc_mult * hc_mult hc_mult3 = hc_mult * 2 + hc_mult2 if not test_hc_head else hc_mult hc_hidden_size = hc_mult * hidden_size - residual = torch.randn(m, hc_mult, hidden_size, dtype=dtypes.bf16) - fn = torch.randn(hc_mult3, hc_hidden_size, dtype=dtypes.fp32) - hc_scale = torch.randn((3,), dtype=dtypes.fp32) * 0.1 - hc_base = torch.randn((hc_mult3,), dtype=dtypes.fp32) * 0.1 + gen = make_generator(seed) + residual = fill((m, hc_mult, hidden_size), data_init, gen, dtype=dtype) + fn = fill((hc_mult3, hc_hidden_size), data_init, gen, dtype=dtypes.fp32) + hc_scale = fill((3,), data_init, gen, dtype=dtypes.fp32) * 0.1 + hc_base = fill((hc_mult3,), data_init, gen, dtype=dtypes.fp32) * 0.1 norm_weight = None if fuse_rmsnorm: - norm_weight = torch.randn(hidden_size, dtype=dtypes.bf16) + norm_weight = fill((hidden_size,), data_init, gen, dtype=dtype) extra_args = { "rms_eps": 1e-6, "hc_pre_eps": 1e-6, @@ -680,11 +689,12 @@ def mhc_post_ref( @benchmark() -def test_mhc_post(m, hidden_size, hc_mult): - x = torch.randn(m, hidden_size, dtype=dtypes.bf16) - residual = torch.randn(m, hc_mult, hidden_size, dtype=dtypes.bf16) - post_layer_mix = torch.randn(m, hc_mult, 1, dtype=dtypes.fp32) - comb_res_mix = torch.randn(m, hc_mult, hc_mult, dtype=dtypes.fp32) +def test_mhc_post(m, hidden_size, hc_mult, dtype=dtypes.bf16, data_init="norm", seed=0): + gen = make_generator(seed) + x = fill((m, hidden_size), data_init, gen, dtype=dtype) + residual = fill((m, hc_mult, hidden_size), data_init, gen, dtype=dtype) + post_layer_mix = fill((m, hc_mult, 1), data_init, gen, dtype=dtypes.fp32) + comb_res_mix = fill((m, hc_mult, hc_mult), data_init, gen, dtype=dtypes.fp32) out_ref = mhc_post_ref(x, residual, post_layer_mix, comb_res_mix) out_hip, hip_us = run_perftest( mhc_post_hip, @@ -784,7 +794,15 @@ def mhc_post_pre_unfused_hip( @benchmark() def test_mhc_post_pre( - m, hidden_size, hc_mult, fuse_rmsnorm=False, large_m=False, fn_pack_bf16=False + m, + hidden_size, + hc_mult, + fuse_rmsnorm=False, + large_m=False, + fn_pack_bf16=False, + dtype=dtypes.bf16, + data_init="norm", + seed=0, ): """Fused mhc_post + mhc_pre: HIP ``mhc_fused_post_pre`` vs ref / unfused HIP / Triton.""" if hidden_size < 512: @@ -800,16 +818,17 @@ def test_mhc_post_pre( hc_mult3 = hc_mult * 2 + hc_mult2 hc_hidden_size = hc_mult * hidden_size - layer_input = torch.randn(m, hidden_size, dtype=dtypes.bf16) - residual_in = torch.randn(m, hc_mult, hidden_size, dtype=dtypes.bf16) - post_layer_mix = torch.randn(m, hc_mult, 1, dtype=dtypes.fp32) - comb_res_mix = torch.randn(m, hc_mult, hc_mult, dtype=dtypes.fp32) - fn = torch.randn(hc_mult3, hc_hidden_size, dtype=dtypes.fp32) - hc_scale = torch.randn((3,), dtype=dtypes.fp32) * 0.1 - hc_base = torch.randn((hc_mult3,), dtype=dtypes.fp32) * 0.1 + gen = make_generator(seed) + layer_input = fill((m, hidden_size), data_init, gen, dtype=dtype) + residual_in = fill((m, hc_mult, hidden_size), data_init, gen, dtype=dtype) + post_layer_mix = fill((m, hc_mult, 1), data_init, gen, dtype=dtypes.fp32) + comb_res_mix = fill((m, hc_mult, hc_mult), data_init, gen, dtype=dtypes.fp32) + fn = fill((hc_mult3, hc_hidden_size), data_init, gen, dtype=dtypes.fp32) + hc_scale = fill((3,), data_init, gen, dtype=dtypes.fp32) * 0.1 + hc_base = fill((hc_mult3,), data_init, gen, dtype=dtypes.fp32) * 0.1 norm_weight = None if fuse_rmsnorm: - norm_weight = torch.randn(hidden_size, dtype=dtypes.bf16) + norm_weight = fill((hidden_size,), data_init, gen, dtype=dtype) extra_args = { "rms_eps": 1e-6, @@ -1018,7 +1037,7 @@ def test_mhc_post_pre( choices=[dtypes.d_dtypes["fp16"], dtypes.d_dtypes["bf16"]], nargs="*", metavar="{fp16, bf16}", - default=["bf16"], + default=[dtypes.bf16], help="""Data type. e.g.: -d bf16""", ) @@ -1065,58 +1084,69 @@ def test_mhc_post_pre( "(mhc_post_pre). gfx950 native bf16 MFMA; gfx1250 wave32 bf16 WMMA (UNVERIFIED); " "other arches fall back to fp32.", ) +add_data_init_args(parser, default_dist="norm") args = parser.parse_args() df = [] for dtype in args.dtype: - for hidden_size in args.hidden_size: - for m in args.m: - for hc_mult in [4]: - ret = test_mhc_pre( - m=m, - hidden_size=hidden_size, - hc_mult=hc_mult, - test_hc_head=args.hc_head, - fuse_rmsnorm=args.fuse_rmsnorm, - fn_pack_bf16=args.fn_pack_bf16, - ) - df.append(ret) -df = pd.DataFrame(df) -df_md = df.to_markdown(index=False) -aiter.logger.info("mhc_pre summary (markdown):\n%s", df_md) - -if not args.hc_head: - df = [] - for dtype in args.dtype: + for data_init in args.data_init: for hidden_size in args.hidden_size: for m in args.m: for hc_mult in [4]: - ret = test_mhc_post(m=m, hidden_size=hidden_size, hc_mult=hc_mult) - df.append(ret) - df = pd.DataFrame(df) - df_md = df.to_markdown(index=False) - aiter.logger.info("mhc_post summary (markdown):\n%s", df_md) - - df = [] - for dtype in args.dtype: - for hidden_size in args.hidden_size: - for m in args.m: - for hc_mult in [4]: - ret = test_mhc_post_pre( + ret = test_mhc_pre( m=m, hidden_size=hidden_size, hc_mult=hc_mult, + test_hc_head=args.hc_head, fuse_rmsnorm=args.fuse_rmsnorm, - large_m=args.largeM, fn_pack_bf16=args.fn_pack_bf16, + dtype=dtype, + data_init=data_init, + seed=args.seed, ) - if ret.get("skipped"): - continue df.append(ret) +print_json_table("mhc_pre summary", df) + +if not args.hc_head: + df = [] + for dtype in args.dtype: + for data_init in args.data_init: + for hidden_size in args.hidden_size: + for m in args.m: + for hc_mult in [4]: + ret = test_mhc_post( + m=m, + hidden_size=hidden_size, + hc_mult=hc_mult, + dtype=dtype, + data_init=data_init, + seed=args.seed, + ) + df.append(ret) + print_json_table("mhc_post summary", df) + + df = [] + for dtype in args.dtype: + for data_init in args.data_init: + for hidden_size in args.hidden_size: + for m in args.m: + for hc_mult in [4]: + ret = test_mhc_post_pre( + m=m, + hidden_size=hidden_size, + hc_mult=hc_mult, + fuse_rmsnorm=args.fuse_rmsnorm, + large_m=args.largeM, + fn_pack_bf16=args.fn_pack_bf16, + dtype=dtype, + data_init=data_init, + seed=args.seed, + ) + if ret.get("skipped"): + continue + df.append(ret) if df: - df = pd.DataFrame(df) - df_md = df.to_markdown(index=False) - aiter.logger.info("mhc_post_pre summary (markdown):\n%s", df_md) + print_json_table("mhc_post_pre summary", df) else: aiter.logger.info("mhc_post_pre: all cases skipped") diff --git a/op_tests/test_mla_v4_kargpreld.py b/op_tests/test_mla_v4_kargpreld.py index ed0324125b..063636ed90 100644 --- a/op_tests/test_mla_v4_kargpreld.py +++ b/op_tests/test_mla_v4_kargpreld.py @@ -38,8 +38,13 @@ import aiter import aiter.mla # main no longer auto-imports submodules; need explicit from aiter import dtypes +from aiter.benchmark_data_init import DATA_DISTS, fill, make_generator from aiter.jit.utils.chip_info import get_gfx -from aiter.test_common import benchmark, checkAllclose, run_perftest +from aiter.test_common import ( + benchmark, + checkAllclose, + run_perftest, +) torch.set_default_device("cuda") @@ -62,7 +67,6 @@ # Perf iteration counts (kept out of the @benchmark signature so they don't # become table columns). main() overrides these from --iters / --warmup. _PERF = {"num_iters": 2, "num_warmup": 1} -_SEED = 0 # --------------------------------------------------------------------------- @@ -297,6 +301,7 @@ def _build_bf16_inputs( kv_seq_lens=64, q_seq_logical=4, seed=0, + data_init="norm", device="cuda", gqa_ratio=GQA_RATIO, attn_sink=True, @@ -311,19 +316,24 @@ def _build_bf16_inputs( mismatch shows up as an err blowup, not a silent pass. False -> per-head -inf ("no sink" no-op: exp(-inf - max) = 0). """ - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) + gen = make_generator(seed, device=device) total_q = batch * q_seq_logical num_page = batch * (kv_seq_lens // PAGE_SIZE) - q_bf16 = torch.randn( - (total_q, gqa_ratio, _QUANT_D), dtype=dtypes.bf16, device=device - ) - kv_bf16 = torch.randn( - (num_page, PAGE_SIZE, NUM_KV_HEADS, _QUANT_D), + q_bf16 = fill( + (total_q * gqa_ratio, _QUANT_D), + data_init, + gen, dtype=dtypes.bf16, device=device, - ) + ).view(total_q, gqa_ratio, _QUANT_D) + kv_bf16 = fill( + (num_page * PAGE_SIZE * NUM_KV_HEADS, _QUANT_D), + data_init, + gen, + dtype=dtypes.bf16, + device=device, + ).view(num_page, PAGE_SIZE, NUM_KV_HEADS, _QUANT_D) qo_indptr = ( torch.arange(0, batch + 1, dtype=torch.int32, device=device) * q_seq_logical @@ -345,7 +355,16 @@ def _build_bf16_inputs( if attn_sink: # randn*10 so the sink contributes materially (~15%) to the softmax; # well above tolerance, so a dropped/mis-scaled sink is a hard mismatch. - sink = torch.randn(num_heads, dtype=torch.float32, device=device) * 10.0 + sink = ( + fill( + (num_heads,), + data_init, + gen, + dtype=torch.float32, + device=device, + ) + * 10.0 + ) else: sink = torch.full( (num_heads,), float("-inf"), dtype=torch.float32, device=device @@ -377,6 +396,8 @@ def test_mla_v4_nm( num_kv_splits=1, gqa_ratio=GQA_RATIO, attn_sink=True, + data_init="norm", + seed=0, ): """Time each v4 nm kernel candidate, check it against the torch fp8-dequant reference, and return per-candidate `us` / `TFLOPS` / `TB/s` / `err`. @@ -413,7 +434,8 @@ def test_mla_v4_nm( batch=batch, kv_seq_lens=kv_seq_lens, q_seq_logical=q_seq_logical, - seed=_SEED, + seed=seed, + data_init=data_init, gqa_ratio=gqa_ratio, attn_sink=attn_sink, ) @@ -733,13 +755,18 @@ def main(): default=[True], help="attn sink value(s) to sweep. e.g. --attn-sink True False", ) + parser.add_argument( + "--data-init", + nargs="+", + choices=list(DATA_DISTS), + default=["norm"], + help="DATA initialization distribution(s) for Q, KV and attention sink", + ) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--iters", type=int, default=50, help="Perf timed iterations") parser.add_argument("--warmup", type=int, default=2, help="Perf warmup iterations") args = parser.parse_args() - global _SEED - _SEED = args.seed _PERF["num_iters"] = args.iters _PERF["num_warmup"] = args.warmup @@ -752,8 +779,20 @@ def main(): ] df = [] - for (nhead, decode_qlen), batch, kv_seq_lens, split_kv, sink in itertools.product( - nhead_combos, args.batch, args.kv_seq_lens, args.split_kv, args.attn_sink + for ( + (nhead, decode_qlen), + batch, + kv_seq_lens, + split_kv, + sink, + data_init, + ) in itertools.product( + nhead_combos, + args.batch, + args.kv_seq_lens, + args.split_kv, + args.attn_sink, + args.data_init, ): try: df.append( @@ -764,6 +803,8 @@ def main(): num_kv_splits=split_kv, gqa_ratio=nhead, attn_sink=sink, + data_init=data_init, + seed=args.seed, ) ) except (RuntimeError, AssertionError) as exc: diff --git a/op_tests/test_moe_ep.py b/op_tests/test_moe_ep.py index dc4db8f790..0df16cc0f9 100644 --- a/op_tests/test_moe_ep.py +++ b/op_tests/test_moe_ep.py @@ -374,7 +374,7 @@ def _per_1x32_mxfp4_quant(w): def _randn_or_const(shape, *, const_init, scale=1.0, dtype=dtypes.bf16, device="cuda"): """randn (scaled) by default; a constant VALUE tensor when const_init is set. - Mirrors test_flydsl_grouped_gemm_gfx1250.py's --const-init: the const path + Mirrors flydsl_tests/test_flydsl_grouped_gemm.py's --const-init: the const path fills with VALUE exactly (the ``scale`` only applies to the random path).""" if const_init is not None: return torch.full(shape, float(const_init), dtype=dtype, device=device) @@ -439,7 +439,7 @@ def test_fmoe_ep_mxfp4( expert (routed or active shared), deduplicated to one buffer row per token. MORI returns that count as the device scalar `total_recv_t`, which ATOM forwards to fused_moe as `num_local_tokens` (mirrors - test_mega_moe_gfx1250.py's DeviceMoEPipeline._layer_step, where total_recv_t + bench_mega_moe.py's DeviceMoEPipeline._layer_step, where total_recv_t comes straight from op.dispatch and feeds moe_forward's num_local_tokens). The dispatch buffer has `trim_M` rows with the full `topk` routing dimension. @@ -623,7 +623,7 @@ def test_fmoe_ep_mxfp4( # total_recv_t: device scalar matching MORI's dispatch return; fused_moe # gets it as num_local_tokens and processes only the first total_recv rows, # skipping the padded tail (mirrors DeviceMoEPipeline._layer_step in - # test_mega_moe_gfx1250.py, where total_recv_t from op.dispatch feeds + # bench_mega_moe.py, where total_recv_t from op.dispatch feeds # moe_forward's num_local_tokens). total_recv_t = torch.tensor([total_recv], dtype=dtypes.i32, device="cuda") num_local_tokens = total_recv_t @@ -679,7 +679,7 @@ def _dequant(w_qt, w_scale, orig_shape): # gugu (INTERLEAVE) stage1 layout so the EP path is routed through the # TDM batched GEMM (_grouped_a8w4_tdm_moe, gugu-only). gate/up are # row-interleaved ([g0,u0,g1,u1,...]) inside moe_shuffle_weight/scale, - # matching test_flydsl_grouped_gemm_gfx1250.py. + # matching flydsl_tests/test_flydsl_grouped_gemm.py. w1_a = moe_shuffle_weight( w1_u8, experts_cnt=total_local, is_guinterleave=True, gate_up=True ) @@ -809,7 +809,7 @@ def _dequant(w_qt, w_scale, orig_shape): # The grouped a8w4 path quantizes activations to fp8, so the reference # (bf16 activations) differs elementwise by more than atol/rtol=5e-2 # even without EP. Use the grouped tests' cosine criterion instead - # (test_flydsl_grouped_gemm_gfx1250.py: logits_diff < 0.01). + # (flydsl_tests/test_flydsl_grouped_gemm.py: logits_diff < 0.01). _logits_diff_tol = 0.01 err = logits_diff _verdict = "PASSED" if logits_diff < _logits_diff_tol else "FAILED" @@ -971,7 +971,7 @@ def _dequant(w_qt, w_scale, orig_shape): help="""initialize activations (input) and weights (w1/w2) to the constant VALUE instead of random values (mxfp4 EP tests only). Bare --const-init uses 0.0 (zero-init). Routing scores stay random so expert selection is unchanged. - Mirrors test_flydsl_grouped_gemm_gfx1250.py --const-init.""", + Mirrors flydsl_tests/test_flydsl_grouped_gemm.py --const-init.""", ) args = parser.parse_args() diff --git a/op_tests/test_moe_mxfp8_passthrough.py b/op_tests/test_moe_mxfp8_passthrough.py index b0372a6045..5c7811b965 100644 --- a/op_tests/test_moe_mxfp8_passthrough.py +++ b/op_tests/test_moe_mxfp8_passthrough.py @@ -18,7 +18,7 @@ Scope, stated so it is not overclaimed: this pins the passthrough against the path it replaces. It is not an absolute-correctness test for the a8w4 MoE, which -op_tests/test_moe_2stage.py and op_tests/flydsl_tests/test_flydsl_moe_a8w4.py +op_tests/test_moe_2stage.py and op_tests/flydsl_tests/test_flydsl_moe.py already cover against a torch reference. The harness mirrors how SGLang's MoRI dispatch calls this: per_1x32 MXFP4 expert diff --git a/op_tests/test_mxfp8fp4gemm.py b/op_tests/test_mxfp8fp4gemm.py index 3fd94a3b24..81ca21e035 100644 --- a/op_tests/test_mxfp8fp4gemm.py +++ b/op_tests/test_mxfp8fp4gemm.py @@ -28,22 +28,25 @@ import aiter from aiter import dtypes +from aiter.benchmark_data_init import ( + fill_fp4, + fill_fp8, + fill_scale_e8m0, + make_generator, +) from aiter.jit.utils.chip_info import get_gfx_runtime as get_gfx from aiter.ops.shuffle import ( shuffle_mxfp8fp4_a, shuffle_mxfp8fp4_b, shuffle_mxfp8fp4_scale, ) -from aiter.test_common import benchmark, checkAllclose, run_perftest +from aiter.test_common import ( + benchmark, + checkAllclose, + run_perftest, +) from aiter.utility import fp4_utils -try: - import bench_init -except ImportError as e: - if e.name != "bench_init": - raise - from op_tests import bench_init - torch.set_default_device("cuda") torch.set_printoptions(sci_mode=False) pd.set_option("display.max_columns", 30) @@ -188,44 +191,31 @@ def _ref(intype, A, B, sA, sB, M, N): return (A_f32 * sA_f) @ (B_f32 * sB_f).T -def _const_mxfp8(rows: int, k: int, val: float) -> torch.Tensor: - # Constant mxfp8 (e4m3): a single representable value, deterministic for perf. - return torch.full((rows, k), val, dtype=torch.float32).to(torch.float8_e4m3fn) - - def _prep( intype: str, M: int, N: int, K: int, apre: int, data_init: str, scale_init: str, gen ): """Build raw + shuffled device tensors and the f32 golden reference. - DATA and SCALE are sampled *independently* (bench_init), selected by + DATA and SCALE are sampled *independently* (test_common), selected by ``data_init`` / ``scale_init``: - data_init : uniform (FP8 U(-6,6) / FP4 U(-3,3)) [default] | gaussian | - trig | random | constant (A/B = 0.5) + data_init : uniform (FP8 U(-6,6) / FP4 U(-3,3)) [default] | norm | + zero | constant (A/B = 0.5) scale_init : auto (E8M0 -> pow2_binomial) [default] | pow2_binomial | - random | constant (neutral 0x7F -> 2^0 = 1.0) + zero | uniform | norm | constant (neutral 0x7F -> 2^0 = 1.0) """ # DATA: A is mxfp8 (e4m3); B is mxfp4 (e2m1 packed) for a8w4, else mxfp8. - if data_init == "constant": - A = _const_mxfp8(M, K, 0.5) - if intype == "a8w4": - B = torch.full((N, K // 2), 0x11, dtype=torch.uint8) # e2m1 nibble 0.5 - else: - B = _const_mxfp8(N, K, 0.5) - else: # uniform / gaussian / trig / random - A = bench_init.fill_fp8((M, K), data_init, gen) - if intype == "a8w4": - B = bench_init.fill_fp4((N, K), data_init, gen) - else: - B = bench_init.fill_fp8((N, K), data_init, gen) - - # SCALE: e8m0 per-32 for both operands. auto -> pow2_binomial for E8M0. - if scale_init == "constant": - sA = torch.full((M, K // MX_SCALE_BLOCK), 0x7F, dtype=torch.uint8) - sB = torch.full((N, K // MX_SCALE_BLOCK), 0x7F, dtype=torch.uint8) - else: # auto / pow2_binomial / random - sA = bench_init.fill_scale_e8m0((M, K // MX_SCALE_BLOCK), scale_init, gen) - sB = bench_init.fill_scale_e8m0((N, K // MX_SCALE_BLOCK), scale_init, gen) + # constant is built inside fill_fp8/fill_fp4 via constant=: A/B = 0.5, and + # the a8w4 fp4 B uses the e2m1 nibble byte 0x11 (= 0.5). + A = fill_fp8((M, K), data_init, gen, constant=0.5) + if intype == "a8w4": + B = fill_fp4((N, K), data_init, gen, constant=0x11) + else: + B = fill_fp8((N, K), data_init, gen, constant=0.5) + + # SCALE: e8m0 per-32 for both operands. auto -> pow2_binomial for E8M0; + # constant default is the neutral e8m0 byte 0x7F (2^0 = 1.0). + sA = fill_scale_e8m0((M, K // MX_SCALE_BLOCK), scale_init, gen) + sB = fill_scale_e8m0((N, K // MX_SCALE_BLOCK), scale_init, gen) # fp32 golden; the caller casts/quantizes it to the requested outtype. ref_f32 = _ref(intype, A, B, sA, sB, M, N) @@ -252,6 +242,10 @@ def test_gemm( seed=0, mode="perf", knl_name=None, + num_warmup=2, + num_iters=None, + test_graph=False, + num_rotate=0, ): # Skip unfittable shapes up front (before prep/shuffle) so they show as # "not support" rather than crashing on a shape assert / missing kernel. @@ -267,9 +261,12 @@ def test_gemm( N, K, ) + _tm, _tn = _heuristic_tile(M) return { "gfx": get_gfx(), "knl_name": knl_name or "(heuristic)", + "tile": f"{_tm}x{_tn}", + "cluster": "4x4", "asm us": float("nan"), "asm TFLOPS": float("nan"), "asm TB/s": float("nan"), @@ -279,11 +276,12 @@ def test_gemm( assert K % MX_SCALE_BLOCK == 0, f"K must be a multiple of {MX_SCALE_BLOCK}" out_dtype = _OUT_DTYPE[outtype] - gen = bench_init.make_generator(seed) # fixed seed -> bit-identical buffers + gen = make_generator(seed) # fixed seed -> bit-identical buffers inp, ref_f32 = _prep(intype, M, N, K, apre, data_init, scale_init, gen) ref = ref_f32.to(out_dtype) needTrace = mode == "profile" - num_iters = 5 if mode == "func" else 101 + # --iters overrides; unset keeps the mode default (func=5, perf/profile=101). + num_iters = num_iters if num_iters is not None else (5 if mode == "func" else 101) # Single ASM kernel under test, dispatched by intype. Inputs passed as ARGS so # run_perftest can rotate them (defeats the L2 hot-cache). Dispatch is @@ -326,6 +324,10 @@ def run_asm(A, B, sA, sB): _tile_m, _tile_n = _heuristic_tile(M) _label = f"f8gemm_{outtype}_{_middle}_{_pre}_{_tile_m}x{_tile_n}_4x4_ps" _report_active_tg(M, N, _tile_m, _tile_n, _label) + # Structured algo details (mxfp8fp4gemm.csv columns): the cpp-dispatch tile + # (M<=64 -> 64x512, else 256x256) and the 4x4 cluster; no splitk/unroll axis. + ret["tile"] = f"{_tile_m}x{_tile_n}" + ret["cluster"] = "4x4" # Only a missing .co is reported as "not support"; any other failure (OOM, # memory fault, shape assert, ...) must propagate, not show as a green cell. # An explicit --knl-name that isn't in the cfg is a real error (typo / missing @@ -339,6 +341,9 @@ def run_asm(A, B, sA, sB): cand, *cand_args, num_iters=num_iters, + num_warmup=num_warmup, + testGraph=test_graph, + num_rotate_args=num_rotate, needTrace=needTrace, ) except Exception as e: @@ -431,30 +436,31 @@ def main(): parser.add_argument( "--data-init", dest="data_init", - nargs="*", - choices=["constant", "uniform", "gaussian", "trig", "random"], + nargs="+", + choices=["zero", "constant", "uniform", "norm"], default=None, - help="DATA init distribution(s) (mblas-style; sampled independently of scale).\n" + help="DATA init distribution(s) (sampled independently of scale).\n" "Paired position-wise with --scale-init (length-1 broadcasts).\n" "Default (unset): perf/profile = 'constant uniform', func = 'uniform'\n" + " zero = all-zero on-wire codes\n" + " constant = A/B = 0.5 (deterministic)\n" " uniform = FP8 U(-6,6) / FP4 U(-3,3) [default]\n" - " gaussian = N(0,1) [norm-dist / LLM-like]\n" - " trig = trig_float in [-2,2] [optimistic pattern]\n" - " random = pure random on-wire codes [overly pessimistic]\n" - " constant = A/B = 0.5 (deterministic)", + " norm = N(0,1) [norm-dist / LLM-like]", ) parser.add_argument( "--scale-init", dest="scale_init", - nargs="*", - choices=["auto", "pow2_binomial", "random", "constant"], + nargs="+", + choices=["auto", "pow2_binomial", "zero", "constant", "uniform", "norm"], default=None, help="SCALE init distribution(s) (e8m0 for both operands)\n" "Default (unset): perf/profile = 'constant auto', func = 'auto'\n" " auto = E8M0 -> pow2_binomial [default]\n" " pow2_binomial = 2^(Binomial(21,0.5)-11)\n" - " random = random e8m0 byte, exp in [-2,2]\n" - " constant = neutral scale 0x7F (2^0 = 1.0)", + " zero = all-zero e8m0 bytes\n" + " constant = neutral scale 0x7F (2^0 = 1.0)\n" + " uniform = U(0.5,2) -> nearest e8m0 byte\n" + " norm = N(1,0.25) -> nearest e8m0 byte", ) parser.add_argument( "--seed", @@ -462,6 +468,39 @@ def main(): default=0, help="RNG seed; same seed -> bit-identical data/scale buffers", ) + parser.add_argument( + "--warmup", + type=int, + default=2, + help="warmup iterations before timing (run_perftest num_warmup)", + ) + parser.add_argument( + "--iters", + type=int, + default=None, + help="timed iterations (run_perftest num_iters); unset -> mode default " + "(func=5, perf/profile=101)", + ) + parser.add_argument( + "--graph", + action="store_true", + help="also time via HIP graph replay (run_perftest testGraph), " + "minimizing inter-kernel gaps", + ) + parser.add_argument( + "--rotate", + type=int, + default=0, + help="rotating input-buffer copies to defeat the L2 hot-cache " + "(run_perftest num_rotate_args); 0 = auto-size from L2", + ) + parser.add_argument( + "--json", + dest="json_out", + default=None, + help="also write the full summary table (all columns) as JSON records " + "to this path, for CI/regression", + ) parser.add_argument( "--knl-name", dest="knl_name", @@ -534,15 +573,39 @@ def shapes_for(intype): seed=args.seed, mode=args.mode, knl_name=args.knl_name, + num_warmup=args.warmup, + num_iters=args.iters, + test_graph=args.graph, + num_rotate=args.rotate, ) for apre, (di, si), intype, outtype in itertools.product( apre_list, init_pairs, args.intype, args.outtype ) for (M, N, K) in shapes_for(intype) ] - df = pd.DataFrame(rows) - # Keep knl_name (the actual .co); drop the columns constant within a table. - df = df.drop(columns=["seed", "gfx", "mode"], errors="ignore") + df_full = pd.DataFrame(rows) + # JSON keeps every column (config + algo details + results) so each record is + # self-describing for CI/regression; the markdown table below drops columns + # that are constant within a run for readability. + if args.json_out: + df_full.to_json(args.json_out, orient="records", indent=2) + aiter.logger.info( + "wrote JSON summary (%d rows) to %s", len(df_full), args.json_out + ) + # Keep knl_name (the actual .co) + tile; drop columns constant within a table + # (cluster is always 4x4). + df = df_full.drop( + columns=[ + "seed", + "gfx", + "mode", + "num_warmup", + "num_iters", + "test_graph", + "num_rotate", + ], + errors="ignore", + ) aiter.logger.info( "mxfp8fp4gemm (F8GEMM) summary (markdown):\n%s", df.to_markdown(index=False), diff --git a/op_tests/test_opus_a16w16_gemm.py b/op_tests/test_opus_a16w16_gemm.py index 4f44494deb..cef48d9c30 100644 --- a/op_tests/test_opus_a16w16_gemm.py +++ b/op_tests/test_opus_a16w16_gemm.py @@ -26,8 +26,17 @@ ) sys.exit(0) +from aiter.benchmark_data_init import ( + DATA_DISTS, + add_data_init_args, + fill, + make_generator, +) from aiter.ops.opus import gemm_a16w16_opus -from aiter.test_common import checkAllclose, run_perftest +from aiter.test_common import ( + checkAllclose, + run_perftest, +) try: from aiter.ops.opus import opus_gemm_workspace_init @@ -79,7 +88,41 @@ def _torch_ref(A: torch.Tensor, B: torch.Tensor, out_dtype): return torch.bmm(A.float(), B.float().transpose(-1, -2)).to(out_dtype) -def _make_b(batch: int, N: int, K: int) -> torch.Tensor: +# --------------------------------------------------------------------------- # +# Data initialization (bf16 operands) + seed +# --------------------------------------------------------------------------- # +# The distributions and the seeded generator come from aiter.test_common (the +# shared data-init API); a16w16 only needs bf16 DATA operands, so it wraps +# ``fill`` and reuses ``make_generator`` / ``add_data_init_args`` verbatim. +DATA_INITS = DATA_DISTS + + +def _make_tensor(shape, dist="norm", gen=None, const_val=1.0): + """Build a bf16 operand under the requested distribution. + + zero : all zeros + constant : filled with ``const_val`` + uniform : U(-1, 1) + norm : N(0, 1) [default; matches the original torch.randn path] + + Delegates to ``benchmark_data_init.fill`` so the operand init matches every + other op test; ``gen`` seeds the sampled dists (uniform/norm), zero/constant + ignore it. + """ + return fill( + shape, + dist, + gen, + dtype=torch.bfloat16, + device="cuda", + uniform=(-1.0, 1.0), + constant=const_val, + ) + + +def _make_b( + batch: int, N: int, K: int, dist: str = "norm", gen=None, const_val: float = 1.0 +) -> torch.Tensor: """Build a B that gemm_a16w16_opus accepts for both batch=1 and batch>1. The wrapper rejects 2D B + batch>1 because the opus launcher hardcodes @@ -87,19 +130,61 @@ def _make_b(batch: int, N: int, K: int) -> torch.Tensor: common "shared weight across batch" case, materialize an explicit `[batch, N, K]` tensor via the contiguous broadcast pattern. """ - B2D = torch.randn(N, K, device="cuda", dtype=torch.bfloat16) + B2D = _make_tensor((N, K), dist, gen, const_val) if batch == 1: return B2D return B2D.unsqueeze(0).expand(batch, -1, -1).contiguous() +def _make_a( + batch: int, M: int, K: int, dist: str = "norm", gen=None, const_val: float = 1.0 +) -> torch.Tensor: + """Build the [batch, M, K] bf16 activation under the requested dist.""" + return _make_tensor((batch, M, K), dist, gen, const_val) + + +# --------------------------------------------------------------------------- # +# FLOPS + Bandwidth +# --------------------------------------------------------------------------- # +def _tflops(batch, M, N, K, us): + """GEMM TFLOPS (2*b*M*N*K FLOP) from microseconds (None-safe).""" + return (2.0 * batch * M * N * K / us / 1e6) if us else None + + +def _tbs(batch, M, N, K, us, out_bytes=2): + """Operand traffic in TB/s (None-safe). + + Bytes moved = A[b,M,K]@bf16 + B@bf16 + out[b,M,N]@out_bytes. B is read once + per batch when materialized (batch>1) and once as a shared weight (batch=1). + bytes / (us*1e-6) / 1e12 == bytes / us / 1e6. + """ + if not us: + return None + a = batch * M * K * 2 + b = (batch * N * K if batch > 1 else N * K) * 2 + o = batch * M * N * out_bytes + return (a + b + o) / us / 1e6 + + def test_a16w16( - batch: int, M: int, N: int, K: int, out_dtype=torch.bfloat16, use_graph=False + batch: int, + M: int, + N: int, + K: int, + out_dtype=torch.bfloat16, + use_graph=False, + *, + dist="norm", + gen=None, + const_val=1.0, + iters=101, + warmup=2, + rotate=0, ): # gemm_a16w16_opus accepts either 2D or 3D A; test 3D to exercise the # batched reshape path. B is 2D when batch==1, 3D contiguous otherwise. - A = torch.randn(batch, M, K, device="cuda", dtype=torch.bfloat16) - B = _make_b(batch, N, K) + A = _make_a(batch, M, K, dist, gen, const_val) + B = _make_b(batch, N, K, dist, gen, const_val) ref = _torch_ref(A, B, out_dtype) @@ -110,6 +195,9 @@ def test_a16w16( None, out_dtype, testGraph=use_graph, + num_iters=iters, + num_warmup=warmup, + num_rotate_args=rotate, ) err = checkAllclose( @@ -119,11 +207,11 @@ def test_a16w16( rtol=0.1, atol=0.5, ) - flops = 2.0 * batch * M * N * K - tflops = flops / us / 1e6 + tflops = _tflops(batch, M, N, K, us) + tbs = _tbs(batch, M, N, K, us, out_bytes=Y.element_size()) print( f"[a16w16] batch={batch} M={M} N={N} K={K} dtype={out_dtype} " - f"| {us:.1f}us | {tflops:.2f} TFLOPs | err={err}" + f"| {us:.1f}us | {tflops:.2f} TFLOPs | {tbs:.3f} TB/s | err={err}" ) return err @@ -171,11 +259,17 @@ def load_opus_shapes(csv_path, gfx, N=None, K=None): sub = sub.sort_values("M").drop_duplicates(subset="M", keep="first") rows = [] for _, r in sub.iterrows(): + # splitK may be blank/NaN in some rows; keep it as an int when present. + try: + splitk = int(r["splitK"]) if not pd.isna(r.get("splitK")) else None + except (KeyError, ValueError, TypeError): + splitk = None rows.append( { "M": int(r["M"]), "csv_us": float(r["us"]), "kernelName": str(r.get("kernelName", "")), + "splitK": splitk, } ) return rows @@ -188,6 +282,13 @@ def test_opus_shapes_graph( K=7168, batch=1, out_dtype=torch.bfloat16, + *, + dist="norm", + gen=None, + const_val=1.0, + iters=101, + warmup=2, + rotate=0, ): """CUDA-graph-mode opus_gemm sweep with golden check. @@ -209,13 +310,16 @@ def test_opus_shapes_graph( passed = failed = 0 perf_rows = [] + out_bytes = torch.empty((), dtype=out_dtype).element_size() for r in rows: M = r["M"] csv_us = r["csv_us"] + kid = r.get("kernelName") or "" + splitk = r.get("splitK") tag = f"a16w16-graph b={batch} M={M} N={N} K={K}" try: - A = torch.randn(batch, M, K, device="cuda", dtype=torch.bfloat16) - B = _make_b(batch, N, K) + A = _make_a(batch, M, K, dist, gen, const_val) + B = _make_b(batch, N, K, dist, gen, const_val) ref = _torch_ref(A, B, out_dtype) # opus split-K workspace must be grown on the capture stream before # run_perftest's graph mode captures this shape. @@ -227,57 +331,94 @@ def test_opus_shapes_graph( None, out_dtype, testGraph=True, + num_iters=iters, + num_warmup=warmup, + num_rotate_args=rotate, ) err = checkAllclose(Y, ref, msg=tag, rtol=0.1, atol=0.5) - tflops = 2.0 * batch * M * N * K / us / 1e6 + tflops = _tflops(batch, M, N, K, us) + tbs = _tbs(batch, M, N, K, us, out_bytes=out_bytes) # Ratio of measured graph latency to the tuned CSV reference. ratio = (us / csv_us) if csv_us else float("nan") + splitk_str = "" if splitk in (None, "") else str(splitk) print( f"[PASS] {tag} | {us:.1f}us (csv {csv_us:.1f}us, " - f"{ratio:.2f}x) | {tflops:.2f} TFLOPs | err={err}" + f"{ratio:.2f}x) | {tflops:.2f} TFLOPs | {tbs:.3f} TB/s | err={err} " + f"| splitK={splitk_str or '-'} | kid={kid or '-'}" ) - perf_rows.append((M, us, csv_us, ratio)) + perf_rows.append((M, us, csv_us, ratio, tflops, tbs, splitk, kid)) passed += 1 except Exception as e: # noqa: BLE001 print(f"[FAIL] {tag} | {type(e).__name__}: {e}") failed += 1 if perf_rows: - print(f"\n{'-' * 64}") + print(f"\n{'-' * 88}") print(f"latency vs tuned CSV [{gfx}] N={N} K={K} batch={batch}") - print(f"{'-' * 64}") - print(f"{'M':>6} | {'graph us':>10} | {'csv us':>10} | {'ratio':>7} | note") - for M, us, csv_us, ratio in perf_rows: + print(f"{'-' * 88}") + print( + f"{'M':>6} | {'graph us':>10} | {'csv us':>10} | {'ratio':>7} | " + f"{'TFLOPs':>9} | {'TB/s':>7} | {'splitK':>6} | note | kernel(kid)" + ) + for M, us, csv_us, ratio, tflops, tbs, splitk, kid in perf_rows: # >20% slower than the tuned reference is flagged for a closer look. note = "" if ratio <= 1.20 else "SLOW >1.20x" - print(f"{M:>6} | {us:>10.2f} | {csv_us:>10.2f} | {ratio:>6.2f}x | {note}") + splitk_str = "-" if splitk in (None, "") else str(splitk) + print( + f"{M:>6} | {us:>10.2f} | {csv_us:>10.2f} | {ratio:>6.2f}x | " + f"{tflops:>9.2f} | {tbs:>7.3f} | {splitk_str:>6} | " + f"{note or '':<10} | {kid or '-'}" + ) print(f"\nSummary: {passed} passed, {failed} failed out of {len(rows)}") return failed == 0 -def test_a16w16_csv_sweep(csv_path: str, batch: int = 1): +def test_a16w16_csv_sweep( + csv_path: str, + batch: int = 1, + *, + out_dtype=torch.bfloat16, + dist="norm", + gen=None, + const_val=1.0, + iters=101, + warmup=2, + rotate=0, + use_graph=False, +): shapes = load_shapes_from_csv(csv_path) print(f"\n{'=' * 80}") print(f"a16w16 sweep from {csv_path}: {len(shapes)} unique shapes, batch={batch}") print("=" * 80) passed = failed = 0 + out_bytes = torch.empty((), dtype=out_dtype).element_size() for M, N, K in shapes: tag = f"a16w16 b={batch} M={M} N={N} K={K}" try: - A = torch.randn(batch, M, K, device="cuda", dtype=torch.bfloat16) - B = _make_b(batch, N, K) - ref = _torch_ref(A, B, torch.bfloat16) + A = _make_a(batch, M, K, dist, gen, const_val) + B = _make_b(batch, N, K, dist, gen, const_val) + ref = _torch_ref(A, B, out_dtype) + if use_graph: + _prewarm_opus_graph_workspace(A, B, out_dtype) Y, us = run_perftest( gemm_a16w16_opus, A, B, None, - torch.bfloat16, + out_dtype, + testGraph=use_graph, + num_iters=iters, + num_warmup=warmup, + num_rotate_args=rotate, ) err = checkAllclose(Y, ref, msg=tag, rtol=0.1, atol=0.5) - tflops = 2.0 * batch * M * N * K / us / 1e6 - print(f"[PASS] {tag} | {us:.1f}us | {tflops:.2f} TFLOPs | err={err}") + tflops = _tflops(batch, M, N, K, us) + tbs = _tbs(batch, M, N, K, us, out_bytes=out_bytes) + print( + f"[PASS] {tag} | {us:.1f}us | {tflops:.2f} TFLOPs | " + f"{tbs:.3f} TB/s | err={err}" + ) passed += 1 except Exception as e: # noqa: BLE001 print(f"[FAIL] {tag} | {type(e).__name__}: {e}") @@ -358,9 +499,58 @@ def test_a16w16_csv_sweep(csv_path: str, batch: int = 1): action="store_true", help="Use CUDA-graph mode for the single-shape / --csv_file paths too.", ) + # --- warmup / iteration controls --- + parser.add_argument( + "--iters", + type=int, + default=101, + help="Timed iterations passed to run_perftest (default: 101).", + ) + parser.add_argument( + "--warmup", + type=int, + default=2, + help="Warmup iterations passed to run_perftest (default: 2).", + ) + # --- rotating tensors --- + parser.add_argument( + "--rotate", + type=int, + default=0, + help=( + "num_rotate_args for run_perftest: number of rotated input copies " + "used to defeat L2 caching. 0 (default) lets the framework auto-size " + "the rotation from the L2 cache; 1 disables rotation." + ), + ) + # --- data initialization + seed (shared aiter.test_common API) --- + # add_data_init_args attaches --data-init / --scale-init / --seed; a16w16 + # has no scale operand, so --scale-init is accepted but unused. Default the + # DATA dist to norm to preserve the original torch.randn init. + add_data_init_args(parser, default_dist="norm") + parser.add_argument( + "--const-val", + type=float, + default=1.0, + help="Fill value used by --data-init constant (default: 1.0).", + ) args = parser.parse_args() out_dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32 + gen = make_generator(args.seed) + if len(args.data_init) != 1: + parser.error( + "--data-init accepts exactly one distribution for the a16w16 benchmark" + ) + data_init = args.data_init[0] + init_kwargs = { + "dist": data_init, + "gen": gen, + "const_val": args.const_val, + "iters": args.iters, + "warmup": args.warmup, + "rotate": args.rotate, + } # Default action (no -m and no --csv_file): auto-sweep the opus shapes in # CUDA-graph mode and print the vs-CSV latency table. So a bare @@ -377,10 +567,17 @@ def test_a16w16_csv_sweep(csv_path: str, batch: int = 1): K=args.k if args.k is not None else 7168, batch=batch, out_dtype=out_dtype, + **init_kwargs, ) sys.exit(0 if ok else 1) elif args.csv_file is not None: - test_a16w16_csv_sweep(args.csv_file, batch=(args.batch or 8)) + test_a16w16_csv_sweep( + args.csv_file, + batch=(args.batch or 8), + out_dtype=out_dtype, + use_graph=args.graph, + **init_kwargs, + ) else: # Clamp K>=128 so every kid the heuristic picks has K>=B_K (smallest is 128). k_eff = max(args.k if args.k is not None else 256, 128) @@ -391,4 +588,5 @@ def test_a16w16_csv_sweep(csv_path: str, batch: int = 1): k_eff, out_dtype=out_dtype, use_graph=args.graph, + **init_kwargs, ) diff --git a/op_tests/test_pa_sparse_prefill.py b/op_tests/test_pa_sparse_prefill.py index afce1a712d..99e9869df7 100644 --- a/op_tests/test_pa_sparse_prefill.py +++ b/op_tests/test_pa_sparse_prefill.py @@ -39,12 +39,18 @@ import torch import aiter # noqa: F401 (registers the top-level export) +from aiter.benchmark_data_init import DATA_DISTS, fill, make_generator +from aiter.benchmark_reporting import print_json_table from aiter.ops.mla_sparse_prefill import mla_sparse_prefill_fp8_asm from aiter.ops.pa_sparse_prefill_opus import ( pa_sparse_prefill_fp8_opus, pa_sparse_prefill_opus, ) -from aiter.test_common import benchmark, checkAllclose, perftest +from aiter.test_common import ( + benchmark, + checkAllclose, + perftest, +) try: from aiter.ops.triton.attention.pa_prefill_sparse import pa_prefill_sparse @@ -307,19 +313,25 @@ def _make_inputs( mode: str = "sparse", device: torch.device | str = "cuda", seed: int = 0, + data_init: str = "norm", ) -> dict: assert mode in _MODES - torch.manual_seed(seed) device = torch.device(device) + gen = make_generator(seed, device=device) - q = (torch.randn(n, h, d, device=device, dtype=torch.float32) * 0.5).to(dtype) + q = ( + fill((n * h, d), data_init, gen, dtype=torch.float32, device=device) + .view(n, h, d) + .mul_(0.5) + ).to(dtype) unified_kv = ( - torch.randn(total_pages, d, device=device, dtype=torch.float32) * 0.5 + fill((total_pages, d), data_init, gen, dtype=torch.float32, device=device) * 0.5 ).to(dtype) - kv = (torch.randn(total_tokens, d, device=device, dtype=torch.float32) * 0.5).to( - dtype - ) - attn_sink = torch.randn(h, device=device, dtype=torch.float32) * 0.25 + kv = ( + fill((total_tokens, d), data_init, gen, dtype=torch.float32, device=device) + * 0.5 + ).to(dtype) + attn_sink = fill((h,), data_init, gen, dtype=torch.float32, device=device) * 0.25 def _csr(total_rows: int, seed_offset: int): if mode == "sparse": @@ -357,19 +369,37 @@ def _make_inputs_fp8( mode: str = "sparse", device: torch.device | str = "cuda", seed: int = 0, + data_init: str = "norm", ) -> dict: """Returns ``{"kernel": ..., "ref": ...}``: the split fp8/bf16 tensors the kernels take, and the dequantized fp32 rows the reference takes. """ assert mode in _MODES - torch.manual_seed(seed) device = torch.device(device) + gen = make_generator(seed, device=device) def _streams(rows: int): nope_fp8, deq = _quantize_nope( - torch.randn(rows, _FP8_D_NOPE, device=device) * 0.5 + fill( + (rows, _FP8_D_NOPE), + data_init, + gen, + dtype=torch.float32, + device=device, + ) + * 0.5 + ) + rope = ( + fill( + (rows, _FP8_D_ROPE), + data_init, + gen, + dtype=torch.float32, + device=device, + ) + * 0.5 ) - rope = (torch.randn(rows, _FP8_D_ROPE, device=device) * 0.5).to(torch.bfloat16) + rope = rope.to(torch.bfloat16) row_fp32 = torch.cat([deq, rope.to(torch.float32)], dim=1) # [rows, 512] return nope_fp8, rope, row_fp32 @@ -380,7 +410,7 @@ def _streams(rows: int): ukn, ukr, ukv_fp32 = _streams(total_pages) kn, kr, kv_fp32 = _streams(total_tokens) - attn_sink = torch.randn(h, device=device, dtype=torch.float32) * 0.25 + attn_sink = fill((h,), data_init, gen, dtype=torch.float32, device=device) * 0.25 def _csr(total_rows: int, seed_offset: int): if mode == "sparse": @@ -450,8 +480,8 @@ def _csr(total_rows: int, seed_offset: int): @perftest() -def _profile_func(target_func, *args, **kwargs): - return target_func(*args, **kwargs) +def _profile_func(target_func, *, backend: str): + return target_func() # --------------------------------------------------------------------------- @@ -488,6 +518,7 @@ def run_pa_sparse_prefill( mode: str = "sparse", backends: tuple = _BACKENDS, seed: int = 0, + data_init: str = "norm", verify: bool = True, bench: bool = True, ) -> dict | None: @@ -498,7 +529,7 @@ def run_pa_sparse_prefill( softmax_scale = 1.0 / math.sqrt(d) msg = ( f"[N={n} H={h} D={d} total_pages={total_pages} total_tokens={total_tokens} " - f"prec={prec} mode={mode}]" + f"prec={prec} mode={mode} data_init={data_init} seed={seed}]" ) wanted = [b for b in _PREC_BACKENDS[prec] if b in backends] @@ -506,7 +537,15 @@ def run_pa_sparse_prefill( candidates: list = [] if prec == "fp8": - data = _make_inputs_fp8(n, h, total_pages, total_tokens, mode=mode, seed=seed) + data = _make_inputs_fp8( + n, + h, + total_pages, + total_tokens, + mode=mode, + seed=seed, + data_init=data_init, + ) kernel_inputs = data["kernel"] ref_fn, ref_inputs = _ref_pa_sparse_prefill_fp8, data["ref"] if "opus" in wanted: @@ -537,6 +576,7 @@ def run_pa_sparse_prefill( _PREC_TO_DTYPE[prec], mode=mode, seed=seed, + data_init=data_init, ) ref_fn, ref_inputs = _ref_pa_sparse_prefill_opus, kernel_inputs if "opus" in wanted: @@ -590,7 +630,7 @@ def run_pa_sparse_prefill( ) if bench: - _, lat_us = _profile_func(invoke) # (data, avg_us_per_iter) + _, lat_us = _profile_func(invoke, backend=name) flops = 4.0 * h * total_nnz * d tflops = flops / max(lat_us * 1e-6, 1e-12) / 1e12 row[f"{name} us"] = round(float(lat_us), 2) @@ -744,6 +784,13 @@ def test_pa_sparse_prefill(prec, n, h, total_pages, total_tokens, mode): default=0, help="RNG seed for input + CSR generation", ) +parser.add_argument( + "--data-init", + nargs="+", + choices=list(DATA_DISTS), + default=["norm"], + help="DATA initialization distribution(s) for Q, KV and attention sink", +) if __name__ == "__main__": @@ -751,12 +798,13 @@ def test_pa_sparse_prefill(prec, n, h, total_pages, total_tokens, mode): rows = [] # product varies its last argument fastest -> this is also the row order. - for prec, mode, h, n, pages_arg in itertools.product( + for prec, mode, h, n, pages_arg, data_init in itertools.product( args.prec, args.mode, args.h_q, args.n_tokens, args.total_pages, + args.data_init, ): total_pages = pages_arg if pages_arg > 0 else n # 0 is "mirror -n" total_tokens = args.total_tokens if args.total_tokens is not None else n @@ -770,6 +818,7 @@ def test_pa_sparse_prefill(prec, n, h, total_pages, total_tokens, mode): mode=mode, backends=tuple(args.backend), seed=args.seed, + data_init=data_init, verify=not args.no_verify, bench=not args.no_bench, ) @@ -784,11 +833,10 @@ def test_pa_sparse_prefill(prec, n, h, total_pages, total_tokens, mode): if drop_cols: df = df.drop(columns=drop_cols) # Column order otherwise follows whichever row first ran a backend. - lead = [c for c in ("prec", "mode", "h", "n") if c in df.columns] + lead = [c for c in ("prec", "mode", "data_init", "h", "n") if c in df.columns] rest = [c for c in df.columns if c not in lead] metrics = [c for b in _BACKENDS for c in rest if c.startswith(f"{b} ")] df = df[lead + [c for c in rest if c not in metrics] + metrics] - print() - print(df.to_string(index=False, na_rep="-")) # na_rep: backend not run + print_json_table("pa_sparse_prefill summary", df) sys.exit(0) sys.exit(0) diff --git a/op_tests/triton_tests/attention/test_mla_v4_triton.py b/op_tests/triton_tests/attention/test_mla_v4_triton.py new file mode 100644 index 0000000000..be966fa696 --- /dev/null +++ b/op_tests/triton_tests/attention/test_mla_v4_triton.py @@ -0,0 +1,709 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +# from __future__ import annotations + +import pytest +import torch +import triton + +from aiter.benchmark_data_init import fill, make_generator +from aiter.ops.triton.attention.pa_decode_sparse import pa_decode_sparse +from aiter.ops.triton.utils._triton import arch_info +from aiter.test_common import ( + benchmark, + checkAllclose, + run_perftest, +) + +# MLA v4 sparse-decode parity: D=512 heads, page_size=1 unified pool. +_PA_DECODE_SPARSE_D = 512 +_PERF = {"num_iters": 50, "num_warmup": 2} + + +def _sparse_attn_torch(q, kv, attn_sink, topk_idxs, softmax_scale): + """Per-batch sparse multi-head attention with sink in the denominator only. + + Shapes: + q: [B, M, H, D] + kv: [B, N, D] + attn_sink: [H] + topk_idxs: [B, M, K] int32, -1 means skip + Returns: + [B, M, H, D] same dtype as q. + """ + B, M, H, _D = q.shape + K = topk_idxs.shape[-1] + device = q.device + out_dtype = q.dtype + + valid = topk_idxs != -1 + safe_idxs = topk_idxs.clamp(min=0).long() + batch_idx = torch.arange(B, device=device).view(B, 1, 1).expand(B, M, K) + kv_gathered = kv[batch_idx, safe_idxs] # [B, M, K, D] + kv_f32 = kv_gathered.float() + kv_f32 = torch.where( + valid.unsqueeze(-1), kv_f32, torch.zeros((), dtype=kv_f32.dtype, device=device) + ) + + q_f32 = q.float() + scores = torch.einsum("bmhd,bmkd->bmhk", q_f32, kv_f32) * float(softmax_scale) + scores = scores.masked_fill(~valid.unsqueeze(2), float("-inf")) + + sink = attn_sink.float().view(1, 1, H, 1).expand(B, M, H, 1) + combined = torch.cat([scores, sink], dim=-1) + cmax = combined.amax(dim=-1, keepdim=True) + cmax = torch.where( + cmax == float("-inf"), + torch.zeros((), dtype=cmax.dtype, device=device), + cmax, + ) + weights = (combined - cmax).exp() + denom = weights.sum(dim=-1, keepdim=True) + weights = weights / denom.clamp(min=1e-30) + weights_kv = weights[..., :K] + out = torch.einsum("bmhk,bmkd->bmhd", weights_kv, kv_f32) + return out.to(out_dtype) + + +def pa_decode_sparse_reference( + q, unified_kv, kv_indices, kv_indptr, attn_sink, softmax_scale +): + """Pure-torch reference that materialises per-token KV via gather.""" + T = q.size(0) + indptr = kv_indptr.to(torch.int64) + spans = (indptr[1:] - indptr[:T]).clamp(min=0) + k_dim = int(spans.max().item()) if T > 0 else 1 + if k_dim == 0: + k_dim = 1 + topk_idxs = torch.full((T, k_dim), -1, device=q.device, dtype=torch.int32) + for t in range(T): + s = int(indptr[t].item()) + n = int(spans[t].item()) + if n > 0: + topk_idxs[t, :n] = kv_indices[s : s + n].to(torch.int32) + return _sparse_attn_torch( + q.unsqueeze(0), + unified_kv.unsqueeze(0), + attn_sink, + topk_idxs.unsqueeze(0), + softmax_scale, + ).squeeze(0) + + +# --------------------------------------------------------------------------- +# Input builder +# --------------------------------------------------------------------------- + + +def _make_inputs( + T: int, + H: int, + D: int, + kv_len_per_token: int, + total_pages: int, + dtype=torch.bfloat16, + seed: int = 0, + data_init: str = "norm", + include_sentinels: bool = False, + variable_len: bool = False, +): + torch.manual_seed(seed) + device = torch.device("cuda") + gen = make_generator(seed) + + q = ( + fill((T * H, D), data_init, gen, dtype=dtype, device=device) + .view(T, H, D) + .mul_(0.5) + ) + unified_kv = ( + fill((total_pages, D), data_init, gen, dtype=dtype, device=device) * 0.5 + ) + attn_sink = fill((H,), data_init, gen, dtype=torch.float32, device=device) * 0.1 + + # Per-token kv_len: fixed or random in [1, kv_len_per_token]. + if variable_len: + kv_lens = torch.randint( + low=1, + high=kv_len_per_token + 1, + size=(T,), + device=device, + dtype=torch.int64, + generator=gen, + ) + else: + kv_lens = torch.full((T,), kv_len_per_token, device=device, dtype=torch.int64) + + indptr = torch.zeros(T + 1, device=device, dtype=torch.int64) + indptr[1:] = kv_lens.cumsum(0) + total_indices = int(indptr[-1].item()) + + indices = torch.randint( + low=0, + high=total_pages, + size=(total_indices,), + device=device, + dtype=torch.int32, + generator=gen, + ) + if include_sentinels and total_indices > 0: + # Sprinkle a few -1 sentinels. + n_sentinel = max(1, total_indices // 16) + sentinel_pos = torch.randperm(total_indices, device=device, generator=gen)[ + :n_sentinel + ] + indices[sentinel_pos] = -1 + + indptr = indptr.to(torch.int32) + softmax_scale = float(D) ** -0.5 + return q, unified_kv, indices, indptr, attn_sink, softmax_scale + + +@benchmark() +def bench_mla_v4_triton_staged( + gqa_ratio, + batch, + kv_seq_lens, + num_kv_splits, + data_init="norm", + seed=0, +): + """Perf-only stage split: main kernel (s1) + reduce (s2) + total.""" + T = batch + H = gqa_ratio + D = _PA_DECODE_SPARSE_D + pages = T * kv_seq_lens + q, unified_kv, indices, indptr, sink, scale = _make_inputs( + T, + H, + D, + kv_seq_lens, + pages, + variable_len=False, + data_init=data_init, + seed=seed, + ) + pa_kwargs = { + "has_invalid": False, + "kv_splits": num_kv_splits, + "num_iters": _PERF["num_iters"], + "num_warmup": _PERF["num_warmup"], + "num_rotate_args": 1, + } + _, us_tot = run_perftest( + pa_decode_sparse, + q, + unified_kv, + indices, + indptr, + sink, + scale, + skip_reduce=False, + **pa_kwargs, + ) + if num_kv_splits > 1: + _, us_s1 = run_perftest( + pa_decode_sparse, + q, + unified_kv, + indices, + indptr, + sink, + scale, + skip_reduce=True, + **pa_kwargs, + ) + triton_s2 = max(0.0, us_tot - us_s1) + else: + us_s1 = us_tot + triton_s2 = 0.0 + return { + "triton_s1": round(us_s1, 2), + "triton_s2": round(triton_s2, 2), + "triton_tot": round(us_tot, 2), + } + + +@benchmark() +def bench_mla_v4_triton_perf( + gqa_ratio, + batch, + kv_seq_lens, + num_kv_splits, + data_init="norm", + seed=0, +): + """Perf sweep row for combo bench / gfx1250 Triton sparse MLA v4 decode. + + Shape ids mirror ``test_mla_v4_kargpreld.test_mla_v4_nm``: + T=batch (q_seq=1), H=gqa_ratio, ctx=kv_seq_lens, kv_splits=num_kv_splits, + D=512. + """ + T = batch + H = gqa_ratio + D = _PA_DECODE_SPARSE_D + pages = T * kv_seq_lens + q, unified_kv, indices, indptr, sink, scale = _make_inputs( + T, + H, + D, + kv_seq_lens, + pages, + variable_len=False, + data_init=data_init, + seed=seed, + ) + _, us = run_perftest( + pa_decode_sparse, + q, + unified_kv, + indices, + indptr, + sink, + scale, + has_invalid=False, + kv_splits=num_kv_splits, + num_iters=_PERF["num_iters"], + num_warmup=_PERF["num_warmup"], + num_rotate_args=1, + ) + flops = 4 * T * H * kv_seq_lens * D # QK^T + P@V + bpe = q.element_size() + nbytes = (T * H * D + T * kv_seq_lens * D + T * H * D) * bpe + return { + "us": round(us, 2), + "TFLOPS": round(flops / us / 1e6, 2), + "TB/s": round(nbytes / us / 1e6, 3), + } + + +# --------------------------------------------------------------------------- +# skip_reduce: the wrapper hands back the pre-reduce split-K partials and the +# caller is responsible for the log-sum-exp combine + sink fold. This mirrors +# the _pa_decode_sparse_reduce kernel in pure torch so we can validate the +# partials against the dense reference. +# --------------------------------------------------------------------------- + + +def _wrapper_main_kernel_params(T: int, H: int, D: int): + """Reproduce the (use_exp2, block_k) the wrapper picks for the main kernel. + + Must stay in sync with ``pa_decode_sparse``'s USE_EXP2 and block_k logic. + """ + use_gluon = arch_info.get_arch() == "gfx1250" + use_exp2 = True + if use_gluon: + if H >= 128: + block_h = 128 + elif H >= 64: + if T >= 2048: + block_h = 64 + elif T >= 32: + block_h = 32 + else: + block_h = 16 + elif H >= 32: + if T >= 256: + block_h = 32 + else: + block_h = 16 + else: + block_h = triton.next_power_of_2(H) + else: + block_h = triton.next_power_of_2(min(H, 16)) + if use_gluon: + block_k = 16 + if block_h == 128: + block_k = 32 + else: + block_k = 16 if D >= 256 else 32 + return use_exp2, block_k + + +def _reduce_partials_torch( + acc_partial, m_partial, l_partial, attn_sink, kv_indptr, block_k, use_exp2 +): + """Pure-torch port of _pa_decode_sparse_reduce. + + Shapes: + acc_partial: [T, KV_SPLITS, H_padded, D] fp32 + m_partial: [T, KV_SPLITS, H_padded] fp32 + l_partial: [T, KV_SPLITS, H_padded] fp32 + Returns [T, H, D] in attn_sink-implied output dtype (bf16/fp16 caller casts). + """ + T, kv_splits, _, D = acc_partial.shape + H = attn_sink.shape[0] + device = acc_partial.device + + expfn = torch.exp2 if use_exp2 else torch.exp + LOG2E = 1.4426950408889634 + sink_scale = LOG2E if use_exp2 else 1.0 + + indptr = kv_indptr.to(torch.int64) + kv_lens = (indptr[1 : T + 1] - indptr[:T]).clamp(min=0) + seg_ids = torch.arange(kv_splits, device=device) + sink = attn_sink.float() * sink_scale # [H] + + out = torch.empty(T, H, D, dtype=torch.float32, device=device) + for t in range(T): + n = int(kv_lens[t].item()) + # Match the kernel's tiles_per_segment / act_num_segments masking so we + # ignore the stale (uninitialised) partial-buffer slots that the split + # kernel early-returned on. + if n <= 0: + act_num_segments = 0 + else: + tiles_per_segment = triton.cdiv(n, kv_splits * block_k) + act_num_segments = triton.cdiv(n, tiles_per_segment * block_k) + seg_mask = seg_ids < act_num_segments # [KV_SPLITS] + + m_p = m_partial[t, :, :H].clone() # [KV_SPLITS, H] + l_p = l_partial[t, :, :H] + a_p = acc_partial[t, :, :H, :] # [KV_SPLITS, H, D] + m_p = torch.where(seg_mask[:, None], m_p, torch.full_like(m_p, float("-inf"))) + + m_max = m_p.max(dim=0).values # [H] + is_dead = m_p == float("-inf") # [KV_SPLITS, H] + alpha = torch.where(is_dead, torch.zeros_like(m_p), expfn(m_p - m_max[None, :])) + l_comb = torch.where(is_dead, torch.zeros_like(l_p), l_p * alpha).sum(0) # [H] + acc_comb = torch.where( + is_dead[:, :, None], torch.zeros_like(a_p), a_p * alpha[:, :, None] + ).sum( + 0 + ) # [H, D] + + m_final = torch.maximum(m_max, sink) + alpha_kv = expfn(m_max - m_final) + alpha_sink = expfn(sink - m_final) + l_final = l_comb * alpha_kv + alpha_sink + acc_final = acc_comb * alpha_kv[:, None] + denom = l_final.clamp(min=1e-30) + out[t] = torch.where( + l_final[:, None] > 0.0, + acc_final / denom[:, None], + torch.zeros_like(acc_final), + ) + return out + + +@pytest.mark.parametrize("T", [1, 64, 256, 2048]) +@pytest.mark.parametrize("H", [16, 32, 64, 128]) +@pytest.mark.parametrize("D", [512]) +@pytest.mark.parametrize("kv_len", [136, 388, 1024]) +@pytest.mark.parametrize("var_len", [True, False]) +@pytest.mark.parametrize("sentinels", [False]) +@pytest.mark.parametrize("skip_reduce", [False]) +def test_pa_decode_sparse_vs_reference( + T, H, D, kv_len, var_len, sentinels, skip_reduce +): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + + pages = T * kv_len + q, ukv, indices, indptr, sink, scale = _make_inputs( + T, + H, + D, + kv_len, + pages, + include_sentinels=sentinels, + variable_len=var_len, + ) + + ref = pa_decode_sparse_reference(q, ukv, indices, indptr, sink, scale) + result = pa_decode_sparse( + q, + ukv, + indices, + indptr, + sink, + scale, + has_invalid=sentinels, + skip_reduce=skip_reduce, + ) + + if isinstance(result, tuple): + # skip_reduce with the split-K path active (kv_splits > 1): the wrapper + # returns raw partials, so do the log-sum-exp combine + sink fold here. + acc_partial, m_partial, l_partial = result + use_exp2, block_k = _wrapper_main_kernel_params(T, H, D) + out = _reduce_partials_torch( + acc_partial, m_partial, l_partial, sink, indptr, block_k, use_exp2 + ).to(q.dtype) + else: + # kv_splits == 1 (skip_reduce is a no-op) or skip_reduce=False: the + # wrapper already returns the final output. + out = result + + tol_err_ratio = 0.01 + assert ( + checkAllclose( + out.to(torch.bfloat16), + ref.to(torch.bfloat16), + atol=5e-3, + rtol=5e-3, + tol_err_ratio=tol_err_ratio, + msg="pa_decode_sparse output", + ) + <= tol_err_ratio + ) + + +# --------------------------------------------------------------------------- +# FP8 KV cache quantization helpers +# --------------------------------------------------------------------------- + +_FP8_GROUP_SIZE = 64 +_FP8_DTYPE = torch.float8_e4m3fnuz + + +def _quantize_kv_fp8(unified_kv, group_size=_FP8_GROUP_SIZE): + """Quantize bf16/fp16 unified_kv to (fp8, scales) with 1xGROUP_SIZE block scaling. + + Returns (kv_fp8, kv_scales) where kv_fp8 is float8_e4m3fnuz and + kv_scales is [total_pages, D // group_size] fp32. + """ + total_pages, D = unified_kv.shape + assert D % group_size == 0 + num_groups = D // group_size + kv_f32 = unified_kv.float().view(total_pages, num_groups, group_size) + amax = kv_f32.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) + fp8_max = torch.finfo(_FP8_DTYPE).max + scales = (amax / fp8_max).squeeze(-1) # [total_pages, num_groups] + kv_scaled = kv_f32 / amax * fp8_max + kv_fp8 = kv_scaled.view(total_pages, D).to(_FP8_DTYPE) + return kv_fp8, scales.to(torch.float32) + + +def _dequant_kv_fp8(kv_fp8, kv_scales, group_size=_FP8_GROUP_SIZE): + """Dequantize for reference comparison.""" + total_pages, D = kv_fp8.shape + num_groups = D // group_size + kv_f32 = kv_fp8.float().view(total_pages, num_groups, group_size) + scales_expanded = kv_scales.unsqueeze(-1).expand( + total_pages, num_groups, group_size + ) + return (kv_f32 * scales_expanded).view(total_pages, D) + + +@pytest.mark.parametrize("T", [1, 32]) +@pytest.mark.parametrize("H", [16]) +@pytest.mark.parametrize("D", [512]) +@pytest.mark.parametrize("kv_len", [100]) +@pytest.mark.parametrize("var_len", [True, False]) +def test_pa_decode_sparse_fp8_vs_reference(T, H, D, kv_len, var_len): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + + pages = T * kv_len + q, ukv_bf16, indices, indptr, sink, scale = _make_inputs( + T, + H, + D, + kv_len, + pages, + variable_len=var_len, + ) + + # Quantize KV to fp8 + scales + kv_fp8, kv_scales = _quantize_kv_fp8(ukv_bf16) + + # Reference: dequant back to bf16, run the torch reference + ukv_deq = _dequant_kv_fp8(kv_fp8, kv_scales).to(q.dtype) + ref = pa_decode_sparse_reference(q, ukv_deq, indices, indptr, sink, scale) + + # Triton kernel with fp8 kv + kv_scales + out = pa_decode_sparse( + q, + kv_fp8, + indices, + indptr, + sink, + scale, + kv_scales=kv_scales, + has_invalid=False, + ) + + tol_err_ratio = 0.01 + assert ( + checkAllclose( + out.to(torch.bfloat16), + ref.to(torch.bfloat16), + atol=1e-2, + rtol=1e-2, + tol_err_ratio=tol_err_ratio, + msg="pa_decode_sparse output", + ) + <= tol_err_ratio + ) + + +def make_packed_cache(num_tokens, D, dtype): + device = "cuda" + rope = 64 # DSv4 RoPE dim, stored bf16 + block = 256 # packed cache page size + nope = D - rope # NoPE dim, stored fp8 e4m3 OCP + nb = triton.cdiv(num_tokens, block) + if dtype == "bf16": + cache = (torch.randn(nb, block, D, device=device) * 0.4).to(torch.bfloat16) + return cache, cache.reshape(nb * block, D).float() + + # per token: [nope fp8 (1B) | rope bf16 (2B) | 8 UE8M0 scale bytes] + data_bytes = nope + rope * 2 + scale_bytes = 8 + row_bytes = data_bytes + scale_bytes + cache = torch.zeros(nb, block, row_bytes, dtype=torch.uint8, device=device) + flat = cache.view(nb, block * row_bytes) + data = flat[:, : block * data_bytes].view(nb, block, data_bytes) + scales_region = flat[:, block * data_bytes :].view(nb, block, scale_bytes) + nope_fp8 = (torch.randn(nb, block, nope, device=device) * 0.4).to( + torch.float8_e4m3fn + ) + data[:, :, :nope] = nope_fp8.view(torch.uint8) + rope_bf16 = (torch.randn(nb, block, rope, device=device) * 0.4).to(torch.bfloat16) + data[:, :, nope:data_bytes] = rope_bf16.view(torch.uint8).view(nb, block, rope * 2) + num_groups = nope // 64 + exps = torch.randint( + 124, 130, (nb, block, num_groups), device=device, dtype=torch.uint8 + ) + scales_region[:, :, :num_groups] = exps + scales = torch.exp2(exps.float() - 127.0).repeat_interleave(64, dim=2) + kv_deq = torch.cat([nope_fp8.float() * scales, rope_bf16.float()], dim=2) + return cache, kv_deq.reshape(nb * block, D) + + +def widen_to_int32_overflow(cache, kv_deq): + """Re-lay ``cache`` as a strided view whose span exceeds a 32-bit offset. + + Same nelement() and same contents, but the dim-0 pitch is stretched so the + last block sits past 2**31 bytes. Only the blocks themselves are written; + the padding between them is left uninitialised, so the pool costs its + address space but not the time to fill it. + """ + nb, block, row = cache.shape + itemsize = cache.element_size() + pitch = triton.cdiv(2**31, max(1, nb - 1) * itemsize) + pitch = max(pitch, block * row) + # the packed fp8 cache is viewed as bfloat16, which needs an even stride + pitch += pitch % 2 + pool = torch.empty( + pitch * (nb - 1) + block * row, dtype=cache.dtype, device=cache.device + ) + view = pool.as_strided((nb, block, row), (pitch, row, 1)) + view.copy_(cache) + assert view.stride(0) * itemsize * (nb - 1) >= 2**31 + return view, kv_deq + + +def two_loop_reference( + q, + main_deq, + main_idx, + main_indptr, + extra_deq, + extra_idx, + extra_indptr, + attn_sink, + softmax_scale, +): + """Reference for the SWA(main) + top-k(extra) two-loop: concatenate the two + dequantized pools, merge the two ragged index sets (extra slots shifted past + the main pool), then reuse ``pa_decode_sparse_reference``. + """ + main_pages = main_deq.shape[0] + combined = torch.cat([main_deq, extra_deq], dim=0).to(q.dtype) + T = main_indptr.numel() - 1 + mi, mp = main_idx.long(), main_indptr.long() + ei, ep = extra_idx.long(), extra_indptr.long() + rows, lens = [], [] + for tok in range(T): + row = torch.cat( + [mi[mp[tok] : mp[tok + 1]], ei[ep[tok] : ep[tok + 1]] + main_pages] + ) + rows.append(row) + lens.append(row.numel()) + combined_idx = torch.cat(rows).to(torch.int32) + combined_indptr = torch.zeros(T + 1, dtype=torch.int32, device=q.device) + combined_indptr[1:] = torch.tensor(lens, device=q.device).cumsum(0) + return pa_decode_sparse_reference( + q, combined, combined_idx, combined_indptr, attn_sink, softmax_scale + ) + + +@pytest.mark.parametrize("T", [1, 32, 128]) +@pytest.mark.parametrize("H", [16]) +@pytest.mark.parametrize("D", [512]) +@pytest.mark.parametrize("main_len", [128]) +@pytest.mark.parametrize("extra_len", [8, 256]) +@pytest.mark.parametrize("dtype", ["bf16", "fp8"]) +@pytest.mark.parametrize("strided_cache", [False, True]) +def test_pa_decode_sparse_two_loop(T, H, D, main_len, extra_len, dtype, strided_cache): + """gfx950 vLLM DSv4 decode path: SWA (main) + top-k (extra) two-loop over + packed caches. fp8 (fp8_ds_mla) is the vLLM production format; bf16 is also + exercised. Skipped off gfx950 (extra_* is a packed-only gluon path).""" + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + if arch_info.get_arch() != "gfx950": + pytest.skip("two-loop (extra_*) is a gfx950 packed-cache-only path") + if strided_cache: + # The pool has to span >2 GiB for the offsets to overflow, so pin the + # regression to one shape -- the fp8 production format at the largest T + # -- rather than paying it on all 24 combinations. + if dtype != "fp8" or T != 128: + pytest.skip("strided-cache case is pinned to the fp8 T=128 shape") + if torch.cuda.mem_get_info()[0] < 4 * 1024**3: + pytest.skip("needs ~3 GiB free for the >2 GiB strided pool") + + device = "cuda" + torch.manual_seed(0) + q = torch.randn(T, H, D, dtype=torch.bfloat16, device=device) * 0.125 + attn_sink = torch.randn(H, dtype=torch.float32, device=device) * 0.1 + softmax_scale = float(D) ** -0.5 + + # main = contiguous SWA window per query + main_cache, main_deq = make_packed_cache(T * main_len, D, dtype) + query_base = (torch.arange(T, device=device) * main_len)[:, None] + main_idx = ( + (query_base + torch.arange(main_len, device=device)).to(torch.int32).reshape(-1) + ) + main_indptr = torch.arange( + 0, T * main_len + 1, main_len, dtype=torch.int32, device=device + ) + # extra = scattered top-k over a pool + extra_pool = T * extra_len + extra_cache, extra_deq = make_packed_cache(extra_pool, D, dtype) + if strided_cache: + extra_cache, extra_deq = widen_to_int32_overflow(extra_cache, extra_deq) + extra_idx = torch.randint( + 0, extra_pool, (T, extra_len), device=device, dtype=torch.int32 + ).reshape(-1) + extra_indptr = torch.arange( + 0, T * extra_len + 1, extra_len, dtype=torch.int32, device=device + ) + + ref = two_loop_reference( + q, + main_deq, + main_idx, + main_indptr, + extra_deq, + extra_idx, + extra_indptr, + attn_sink, + softmax_scale, + ) + out = pa_decode_sparse( + q, + main_cache, + main_idx, + main_indptr, + attn_sink, + softmax_scale, + extra_cache=extra_cache, + extra_indices=extra_idx, + extra_indptr=extra_indptr, + ) + + tol = 1e-2 if dtype == "fp8" else 5e-3 + torch.testing.assert_close(out, ref, atol=tol, rtol=tol)