diff --git a/aiter/ops/triton/__init__.py b/aiter/ops/triton/__init__.py index bffa28311b..ca4f32f826 100644 --- a/aiter/ops/triton/__init__.py +++ b/aiter/ops/triton/__init__.py @@ -3,6 +3,7 @@ import importlib.util import sys +import warnings from types import SimpleNamespace # Try to import quant module @@ -52,6 +53,20 @@ for modules that were reorganized so that external repos (like sglang for example), which depend on the old module names, can still import it the old "way" of importing. """ +# Paths that only exist for backward compatibility and are on their way out. +_DEPRECATED_COMPAT_PATHS = ("gluon.gemm_a8w8", "gluon.gemm_a8w8_blockscale") + + +def _warn_if_deprecated(name, new_path): + if name in _DEPRECATED_COMPAT_PATHS: + warnings.warn( + f"aiter.ops.triton.{name} has moved to {new_path}; this path " + "will be removed in a future release.", + DeprecationWarning, + stacklevel=3, + ) + + # This is a mapping of the old module names to the new module names _BACKWARD_COMPAT_MAP = { # Batched GEMM modules (gemm/batched/) @@ -71,6 +86,8 @@ "gemm_a8w8_blockscale": "gemm.basic.gemm_a8w8_blockscale", "gemm_a8w8_per_token_scale": "gemm.basic.gemm_a8w8_per_token_scale", "gemm_a8w8": "gemm.basic.gemm_a8w8", + "gluon.gemm_a8w8": "gemm.basic.gemm_a8w8", + "gluon.gemm_a8w8_blockscale": "gemm.basic.gemm_a8w8_blockscale", "gemm_a8wfp4": "gemm.basic.gemm_a8wfp4", "gemm_afp4wfp4_pre_quant_atomic": "gemm.basic.gemm_afp4wfp4_pre_quant_atomic", "gemm_afp4wfp4": "gemm.basic.gemm_afp4wfp4", @@ -148,6 +165,7 @@ def __getattr__(name): """ if name in _BACKWARD_COMPAT_MAP: new_path = f"aiter.ops.triton.{_BACKWARD_COMPAT_MAP[name]}" + _warn_if_deprecated(name, new_path) module = importlib.import_module(new_path) sys.modules[f"aiter.ops.triton.{name}"] = module return module @@ -161,10 +179,11 @@ def _backward_compat_find_spec(fullname, path, target=None): from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4 import aiter.ops.triton.gemm_afp4wfp4 """ - if fullname.startswith("aiter.ops.triton.") and fullname.count(".") == 3: - name = fullname.split(".")[-1] + if fullname.startswith("aiter.ops.triton."): + name = fullname[len("aiter.ops.triton.") :] if name in _BACKWARD_COMPAT_MAP: new_path = f"aiter.ops.triton.{_BACKWARD_COMPAT_MAP[name]}" + _warn_if_deprecated(name, new_path) try: sys.modules[fullname] = importlib.import_module(new_path) return importlib.util.find_spec(new_path) diff --git a/aiter/ops/triton/gluon/gemm_a8w8.py b/aiter/ops/triton/_gluon_kernels/gfx950/gemm/basic/gemm_a8w8.py similarity index 74% rename from aiter/ops/triton/gluon/gemm_a8w8.py rename to aiter/ops/triton/_gluon_kernels/gfx950/gemm/basic/gemm_a8w8.py index da5d00f7e5..3861d9c044 100644 --- a/aiter/ops/triton/gluon/gemm_a8w8.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/gemm/basic/gemm_a8w8.py @@ -1,15 +1,13 @@ -import torch +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Gluon INT8/FP8 GEMM kernels for gfx950.""" + import triton from triton.experimental import gluon from triton.experimental.gluon import language as gl -from aiter.ops.triton.utils._triton import arch_info from aiter.ops.triton.utils._triton.pid_preprocessing import pid_grid, remap_xcd -from aiter.ops.triton.utils.device_info import get_num_xcds -from aiter.ops.triton.utils.gemm_config_utils import get_gemm_config -from aiter.ops.triton.utils.logger import AiterTritonLogger - -_LOGGER = AiterTritonLogger() @triton.heuristics( @@ -548,186 +546,3 @@ def _gemm_a8w8_preshuffled_kernel( c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) gl.amd.cdna4.buffer_store(stored_value=c, ptr=c_ptr, offsets=c_offs, mask=c_mask) - - -def _get_config( - M: int, - N: int, - K: int, -): - if arch_info.get_arch() != "gfx950": - raise ValueError( - "Gluon implementation is not supported on this device (requires CDNA4)." - ) - # get_gemm_config caches internally and returns a fresh deep copy. - config, _ = get_gemm_config("GEMM-A8W8", M, N, K, backend="gluon") - return config - - -def gemm_a8w8( - x: torch.Tensor, - w: torch.Tensor, - x_scale: torch.Tensor, - w_scale: torch.Tensor, - bias: torch.Tensor | None = None, - dtype: float | None = torch.bfloat16, - y: torch.Tensor | None = None, - config: dict | None = None, -): - """ - Computes 8 bit matrix multiplication Y = (X @ W^T) * (x_scale * w_scale) with optional bias. - INT8 inputs are scaled back to higher precision using per-tensor scale factors. - - Args: - x (torch.Tensor): INT8 input matrix with shape (M, K). - w (torch.Tensor): INT8 weight matrix with shape (N, K), internally transposed. - x_scale (torch.Tensor): Scale factor for x with shape (M, 1) or (M,). - w_scale (torch.Tensor): Scale factor for w with shape (1, N) or (N,). - bias (Optional[torch.Tensor]): Bias vector with shape (N,). - dtype (Optional[torch.dtype]): Output datatype (BF16 or FP16). - y (Optional[torch.Tensor]): Pre-allocated output tensor with shape (M, N). - config (Optional[dict]): Kernel tuning parameters (BLOCK_SIZE_M, BLOCK_SIZE_N, - BLOCK_SIZE_K, GROUP_SIZE_M). - - Returns: - torch.Tensor: Output with shape (M, N) in higher precision format. - """ - - _LOGGER.info( - f"GEMM_A8W8: x={tuple(x.shape)} w={tuple(w.shape)} x_scale={tuple(x_scale.shape)} w_scale={tuple(w_scale.shape)}" - ) - - # Check constraints. - assert x.shape[1] == w.shape[1], "Incompatible dimensions!!!" - assert x.dtype == w.dtype, "Input types must be the same" - - M, K = x.shape - N, K = w.shape - - # Transpose w (kernel expects (K, N)) - w = w.T - - if y is None: - y = torch.empty((M, N), dtype=dtype, device=x.device) - - if config is None: - config = _get_config(M, N, K) - - if x.dtype == torch.float8_e4m3fn: - fp8_format = "e4m3" - elif x.dtype == torch.float8_e5m2: - fp8_format = "e5m2" - else: - fp8_format = None # int8 case - - grid = ( - triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv(N, config["BLOCK_SIZE_N"]), - ) - _gemm_a8w8_kernel[grid]( - x, - w, - x_scale, - w_scale, - bias, - y, - M, - N, - K, - x.stride(0), - x.stride(1), - w.stride(0), - w.stride(1), - y.stride(0), - y.stride(1), - bias is not None, - NUM_XCDS=get_num_xcds(), - NUM_WARPS=config["num_warps"], - **config, - FP8_FORMAT=fp8_format, - ) - - return y - - -def gemm_a8w8_preshuffle( - x: torch.Tensor, - w: torch.Tensor, - x_scale: torch.Tensor, - w_scale: torch.Tensor, - bias: torch.Tensor | None = None, - dtype: float | None = torch.bfloat16, - y: torch.Tensor | None = None, - config: dict | None = None, -): - """ - Computes 8 bit matrix multiplication Y = (X @ W^T) * (x_scale * w_scale) with optional bias. - INT8 inputs are scaled back to higher precision using per-tensor scale factors. - - Args: - x (torch.Tensor): INT8 input matrix with shape (M, K). - w (torch.Tensor): INT8 weight matrix with shape (N*16, K//16), internally transposed. - x_scale (torch.Tensor): Scale factor for x with shape (M, 1) or (M,). - w_scale (torch.Tensor): Scale factor for w with shape (1, N) or (N,). - bias (Optional[torch.Tensor]): Bias vector with shape (N,). - dtype (Optional[torch.dtype]): Output datatype (BF16 or FP16). - y (Optional[torch.Tensor]): Pre-allocated output tensor with shape (M, N). - config (Optional[dict]): Kernel tuning parameters (BLOCK_SIZE_M, BLOCK_SIZE_N, - BLOCK_SIZE_K, GROUP_SIZE_M). - - Returns: - torch.Tensor: Output with shape (M, N) in higher precision format. - """ - - _LOGGER.info( - f"GEMM_A8W8: x={tuple(x.shape)} w={tuple(w.shape)} x_scale={tuple(x_scale.shape)} w_scale={tuple(w_scale.shape)}" - ) - - M, K = x.shape - N, K = w.shape - N = N * 16 - K = K // 16 - - if y is None: - y = torch.empty((M, N), dtype=dtype, device=x.device) - - if config is None: - config = _get_config(M, N, K) - - assert ( - K % config["BLOCK_SIZE_K"] == 0 - ), "K must be multiple of BLOCK_SIZE_K for preshuffling" - - if x.dtype == torch.float8_e4m3fn: - fp8_format = "e4m3" - elif x.dtype == torch.float8_e5m2: - fp8_format = "e5m2" - else: - fp8_format = None # int8 case - - grid = ( - triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv(N, config["BLOCK_SIZE_N"]), - ) - _gemm_a8w8_preshuffled_kernel[grid]( - x, - w, - x_scale, - w_scale, - bias, - y, - M, - N, - K, - x.stride(0), - x.stride(1), - w.stride(0), - w.stride(1), - y.stride(0), - y.stride(1), - bias is not None, - NUM_XCDS=get_num_xcds(), - NUM_WARPS=config["num_warps"], - **config, - FP8_FORMAT=fp8_format, - ) - - return y diff --git a/aiter/ops/triton/gluon/gemm_a8w8_blockscale.py b/aiter/ops/triton/_gluon_kernels/gfx950/gemm/basic/gemm_a8w8_blockscale.py similarity index 75% rename from aiter/ops/triton/gluon/gemm_a8w8_blockscale.py rename to aiter/ops/triton/_gluon_kernels/gfx950/gemm/basic/gemm_a8w8_blockscale.py index 3d3a9527e8..3516286e26 100644 --- a/aiter/ops/triton/gluon/gemm_a8w8_blockscale.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/gemm/basic/gemm_a8w8_blockscale.py @@ -1,25 +1,15 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -import functools +"""Gluon FP8 block-scale GEMM kernels for gfx950.""" -import torch import triton -from triton import language as tl from triton.experimental import gluon from triton.experimental.gluon import language as gl from triton.runtime.jit import constexpr_function -from aiter.ops.triton.utils._triton import arch_info from aiter.ops.triton.utils._triton.pid_preprocessing import pid_grid, remap_xcd -from aiter.ops.triton.utils.core import load_config_json -from aiter.ops.triton.utils.gemm_config_utils import resolve_config_dir -from aiter.ops.triton.utils.logger import AiterTritonLogger -_LOGGER = AiterTritonLogger() - - -# Supported (BLOCK_M, BLOCK_N) tiles; BLOCK_K=128 and NUM_WARPS=4 are baked in. _SUPPORTED_TILES = ((64, 128), (128, 128), (128, 256)) @@ -894,279 +884,3 @@ def _gemm_a8w8_blockscale_kernel( NEED_N_MASK=NEED_N_MASK, NUM_WARPS=NUM_WARPS, ) - - -@gluon.jit -def _gemm_a8w8_blockscale_reduce_kernel( - c_in_ptr, - c_out_ptr, - M, - N, - stride_c_in_k, - stride_c_in_m, - stride_c_in_n, - stride_c_out_m, - stride_c_out_n, - BLOCK_SIZE_M: gl.constexpr, # Note: Can be distinct from GEMM block size - BLOCK_SIZE_N: gl.constexpr, - ACTUAL_KSPLIT: gl.constexpr, - MAX_KSPLIT: gl.constexpr, -): - - pid_m = gl.program_id(axis=0) - pid_n = gl.program_id(axis=1) - - blocked_read: gl.constexpr = gl.BlockedLayout( # (MAX_KSPLIT, BLOCK_M, BLOCK_N) - size_per_thread=[1, 1, 4], - threads_per_warp=[1, 8, 8], - warps_per_cta=[1, 4, 1], - order=[2, 1, 0], - ) - - # blocked_write: gl.constexpr = gl.BlockedLayout( - # size_per_thread=[1, 4], # (BLOCK_M, BLOCK_N) - # threads_per_warp=[8, 8], - # warps_per_cta=[4, 1], - # order=[1, 0], - # ) - - offs_m = pid_m * BLOCK_SIZE_M + gl.arange( - 0, - BLOCK_SIZE_M, # keep dim 1 - gl.SliceLayout(0, gl.SliceLayout(2, blocked_read)), - ) - offs_n = pid_n * BLOCK_SIZE_N + gl.arange( - 0, - BLOCK_SIZE_N, # keep dim 2 - gl.SliceLayout(0, gl.SliceLayout(1, blocked_read)), - ) - offs_k = gl.arange( - 0, MAX_KSPLIT, gl.SliceLayout(1, gl.SliceLayout(2, blocked_read)) # keep dim 0 - ) - c_in_offs = ( - (offs_k[:, None, None] * stride_c_in_k) - + (offs_m[None, :, None] * stride_c_in_m) - + (offs_n[None, None, :] * stride_c_in_n) - ) - if ACTUAL_KSPLIT == MAX_KSPLIT: - c_in_mask = (offs_m[None, :, None] < M) & (offs_n[None, None, :] < N) - c = gl.amd.cdna4.buffer_load(c_in_ptr, c_in_offs, mask=c_in_mask, cache=".ca") - else: - c_in_mask = ( - (offs_m[None, :, None] < M) - & (offs_n[None, None, :] < N) - & (offs_k[:, None, None] < ACTUAL_KSPLIT) - ) - c = gl.amd.cdna4.buffer_load( - c_in_ptr, c_in_offs, mask=c_in_mask, cache=".ca" - ) # , other=0.0) - c = tl.sum(c, 0) - - c = c.to(c_out_ptr.type.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + gl.arange( - 0, BLOCK_SIZE_M, gl.SliceLayout(1, gl.SliceLayout(0, blocked_read)) - ) - offs_cn = pid_n * BLOCK_SIZE_N + gl.arange( - 0, BLOCK_SIZE_N, gl.SliceLayout(0, gl.SliceLayout(0, blocked_read)) - ) - c_out_offs = (offs_cm[:, None] * stride_c_out_m) + ( - offs_cn[None, :] * stride_c_out_n - ) - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) - - gl.amd.cdna4.buffer_store( - stored_value=c, ptr=c_out_ptr, offsets=c_out_offs, mask=c_mask - ) - - -@functools.lru_cache(maxsize=1024) -def _get_config_cached( - M: int, - N: int, - K: int, -): - if not arch_info.is_gluon_avail(): - raise ValueError( - "Gluon implementation is not supported on this device (requires CDNA4)." - ) - - # This family's configs live in the nested layout, so its directory is the - # only candidate and the returned name prefix is always empty: nested files - # carry no arch prefix and name their default DEFAULT.json. - cfg_dir, _ = resolve_config_dir("gemm", "GEMM-A8W8_BLOCKSCALE", backend="gluon") - - # Try specialized config first. - config_dict = load_config_json( - f"{cfg_dir}/GEMM-A8W8_BLOCKSCALE-N={N}-K={K}.json", - required=False, - ) - # Fall back to the general config (must exist). - if config_dict is None: - config_dict = load_config_json(f"{cfg_dir}/DEFAULT.json") - - # Config keys should be named M_LEQ_ or "any" - bounds = [] - for setting in config_dict: - potential_block_m = setting.replace("M_LEQ_", "") - if potential_block_m.isnumeric(): - bounds.append(int(potential_block_m)) - - # Walk buckets in ascending-M order; pick the smallest one whose tile - # the kernel currently supports. Unsupported buckets are skipped (those - # configs become live again once the kernel grows the corresponding - # padded-LDS layouts), so we may fall through to "any". - config = config_dict["any"] - for bound in sorted(bounds): - if M > bound or f"M_LEQ_{bound}" not in config_dict: - continue - candidate = config_dict[f"M_LEQ_{bound}"] - if (candidate["BLOCK_SIZE_M"], candidate["BLOCK_SIZE_N"]) in _SUPPORTED_TILES: - config = candidate - break - - return config - - -def _get_config( - M: int, - N: int, - K: int, -): - # Fresh copy per call, outside the lru boundary — the caller writes - # derived fields (SPLITK_BLOCK_SIZE here, GROUP_K/GROUP_N at the call - # site) into the returned dict. - config = _get_config_cached(M, N, K).copy() - - block_size_k = config["BLOCK_SIZE_K"] - num_k_blocks = triton.cdiv(K, block_size_k) - num_k_blocks_per_split = triton.cdiv(num_k_blocks, config["NUM_KSPLIT"]) - config["SPLITK_BLOCK_SIZE"] = num_k_blocks_per_split * block_size_k - - return config - - -def gemm_a8w8_blockscale( - x: torch.Tensor, - w: torch.Tensor, - x_scale: torch.Tensor, - w_scale: torch.Tensor, - dtype: float | None = torch.bfloat16, - y: torch.Tensor | None = None, - config: dict | None = None, -): - """ - Computes the 8 bit matmul Y = X x WT using the block-scale quantization approach. - - Key parameters: - - X: Matrix X with shape (M, K). - - W: Matrix W with shape (N, K). - - X_scale: Scale tensor for X with shape (M, *scale_k). - - W_scale: Scale tensor for W with shape (**scale_n, *scale_k). - - Returns: - - Y: The output matrix with shape (M, N). - - *scale_k = (K + scale_block_size_k - 1) // scale_block_size_k - **scale_n = (N + scale_block_size_n - 1) // scale_block_size_n - """ - _LOGGER.info( - f"GEMM_A8W8_BLOCKSCALE: x={tuple(x.shape)} w={tuple(w.shape)} x_scale={tuple(x_scale.shape)} w_scale={tuple(w_scale.shape)}" - ) - - M, K = x.shape - N, K = w.shape - - # Check constraints. - assert x.shape[1] == w.shape[1], "Incompatible dimensions!!!" - - # Transpose w and w_scale - w = w.T - w_scale = w_scale.T - - if y is None: - y = torch.empty((M, N), dtype=dtype, device=x.device) - - if config is None: - config = _get_config(M, N, K) - - # Scale block sizes - # TODO: need a better way to pass scale block sizes around - config["GROUP_K"] = triton.next_power_of_2(triton.cdiv(K, w_scale.shape[0])) - config["GROUP_N"] = triton.next_power_of_2(triton.cdiv(N, w_scale.shape[1])) - - if config["NUM_KSPLIT"] == 1: - assert ( - config["GROUP_K"] == config["BLOCK_SIZE_K"] - ), f"GROUP_K: {config['GROUP_K']} must equal BLOCK_SIZE_K: {config['BLOCK_SIZE_K']} when not using KSPLIT" - - if config["NUM_KSPLIT"] > 1: - y_pp = torch.empty( - (config["NUM_KSPLIT"], M, N), dtype=torch.float32, device=y.device - ) - else: - y_pp = None - - num_stages = config.get("num_stages", 2) - num_stages = max(num_stages, 2) - - # grid = (config["NUM_KSPLIT"], triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv(N, config["BLOCK_SIZE_N"]),) - grid = lambda META: ( - ( - META["NUM_KSPLIT"] - * triton.cdiv(M, META["BLOCK_SIZE_M"]) - * triton.cdiv(N, META["BLOCK_SIZE_N"]) - ), - ) - _gemm_a8w8_blockscale_kernel[grid]( - x, - w, - y if config["NUM_KSPLIT"] == 1 else y_pp, - x_scale, - w_scale, - M, - N, - K, - x.stride(0), - x.stride(1), - w.stride(0), - w.stride(1), - 0 if config["NUM_KSPLIT"] == 1 else y_pp.stride(0), - y.stride(0) if config["NUM_KSPLIT"] == 1 else y_pp.stride(1), - y.stride(1) if config["NUM_KSPLIT"] == 1 else y_pp.stride(2), - x_scale.stride(0), - x_scale.stride(1), - w_scale.stride(0), - w_scale.stride(1), - NUM_WARPS=config["num_warps"], - NUM_STAGES=num_stages, - **config, - ) - - if config["NUM_KSPLIT"] > 1: - REDUCE_BLOCK_SIZE_M = 32 - REDUCE_BLOCK_SIZE_N = 32 - ACTUAL_KSPLIT = triton.cdiv(K, config["SPLITK_BLOCK_SIZE"]) - - grid_reduce = ( - triton.cdiv(M, REDUCE_BLOCK_SIZE_M), - triton.cdiv(N, REDUCE_BLOCK_SIZE_N), - ) - - _gemm_a8w8_blockscale_reduce_kernel[grid_reduce]( - y_pp, - y, - M, - N, - y_pp.stride(0), - y_pp.stride(1), - y_pp.stride(2), - y.stride(0), - y.stride(1), - REDUCE_BLOCK_SIZE_M, - REDUCE_BLOCK_SIZE_N, - ACTUAL_KSPLIT, - triton.next_power_of_2(config["NUM_KSPLIT"]), - ) - - return y diff --git a/aiter/ops/triton/configs/CLAUDE.md b/aiter/ops/triton/configs/CLAUDE.md index cf2e15f0db..af4f49cdf8 100644 --- a/aiter/ops/triton/configs/CLAUDE.md +++ b/aiter/ops/triton/configs/CLAUDE.md @@ -123,11 +123,9 @@ Consequences to keep in mind: instead). Direct-path loaders bypass the resolver's directory probe. Grep for -`f"{AITER_TRITON_CONFIGS_PATH}/..."` before moving anything — -`gluon/gemm_a8w8_blockscale.py` still builds legacy `gemm/gluon/` paths by -hand (via `load_config_json`) and must be edited when its configs move. -`gluon/gemm_a8w8.py` and `gluon/gemm_afp4wfp4.py` go through -`get_gemm_config(backend="gluon")` and need no changes. +`f"{AITER_TRITON_CONFIGS_PATH}/..."` before moving anything. +`gluon/gemm_afp4wfp4.py` goes through `get_gemm_config(backend="gluon")` and +needs no changes. --- diff --git a/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/DEFAULT.json b/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/DEFAULT.json index 8bac290f85..c09d900733 100644 --- a/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/DEFAULT.json +++ b/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/DEFAULT.json @@ -1,16 +1,4 @@ { - "M_LEQ_16": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".ca", - "NUM_KSPLIT": 4 - }, "M_LEQ_128": { "BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, diff --git a/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=2112-K=7168.json b/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=2112-K=7168.json index 8747331bb9..f94ac349e4 100644 --- a/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=2112-K=7168.json +++ b/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=2112-K=7168.json @@ -1,52 +1,4 @@ { - "M_LEQ_16": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 16, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 2, - "num_stages": 1, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 14 - }, - "M_LEQ_32": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 4, - "num_warps": 2, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": null, - "NUM_KSPLIT": 7 - }, - "M_LEQ_64": { - "BLOCK_SIZE_M": 32, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 4, - "num_warps": 2, - "num_stages": 1, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": null, - "NUM_KSPLIT": 7 - }, - "M_LEQ_128": { - "BLOCK_SIZE_M": 64, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 2, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".ca", - "NUM_KSPLIT": 7 - }, "M_LEQ_2048": { "BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, diff --git a/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=3072-K=1536.json b/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=3072-K=1536.json index f23eb68ff6..f94ac349e4 100644 --- a/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=3072-K=1536.json +++ b/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=3072-K=1536.json @@ -1,52 +1,4 @@ { - "M_LEQ_16": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 2, - "num_stages": 1, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 6 - }, - "M_LEQ_32": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 4, - "num_warps": 2, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".ca", - "NUM_KSPLIT": 6 - }, - "M_LEQ_64": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 2, - "num_warps": 2, - "num_stages": 1, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".ca", - "NUM_KSPLIT": 1 - }, - "M_LEQ_128": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 2, - "num_warps": 2, - "num_stages": 1, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".ca", - "NUM_KSPLIT": 1 - }, "M_LEQ_2048": { "BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, diff --git a/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=4608-K=7168.json b/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=4608-K=7168.json index b6319fb0e3..f94ac349e4 100644 --- a/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=4608-K=7168.json +++ b/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=4608-K=7168.json @@ -1,52 +1,4 @@ { - "M_LEQ_16": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 2, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 14 - }, - "M_LEQ_32": { - "BLOCK_SIZE_M": 32, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 2, - "num_warps": 2, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".ca", - "NUM_KSPLIT": 14 - }, - "M_LEQ_64": { - "BLOCK_SIZE_M": 32, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 2, - "num_warps": 2, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".ca", - "NUM_KSPLIT": 8 - }, - "M_LEQ_128": { - "BLOCK_SIZE_M": 64, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 2, - "num_warps": 2, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".ca", - "NUM_KSPLIT": 7 - }, "M_LEQ_2048": { "BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, diff --git a/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=512-K=7168.json b/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=512-K=7168.json index 21e6a61c1a..f94ac349e4 100644 --- a/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=512-K=7168.json +++ b/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=512-K=7168.json @@ -1,52 +1,4 @@ { - "M_LEQ_16": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 2, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".ca", - "NUM_KSPLIT": 28 - }, - "M_LEQ_32": { - "BLOCK_SIZE_M": 32, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 2, - "num_warps": 2, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".ca", - "NUM_KSPLIT": 14 - }, - "M_LEQ_64": { - "BLOCK_SIZE_M": 32, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 2, - "num_warps": 2, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".ca", - "NUM_KSPLIT": 14 - }, - "M_LEQ_128": { - "BLOCK_SIZE_M": 32, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 2, - "num_warps": 2, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".ca", - "NUM_KSPLIT": 14 - }, "M_LEQ_2048": { "BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, diff --git a/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=7168-K=2048.json b/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=7168-K=2048.json index 098cf56ea0..f94ac349e4 100644 --- a/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=7168-K=2048.json +++ b/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=7168-K=2048.json @@ -1,52 +1,4 @@ { - "M_LEQ_16": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 2, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 4 - }, - "M_LEQ_32": { - "BLOCK_SIZE_M": 32, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 4, - "num_warps": 2, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 4 - }, - "M_LEQ_64": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 2, - "num_warps": 2, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".ca", - "NUM_KSPLIT": 1 - }, - "M_LEQ_128": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 2, - "num_warps": 2, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".ca", - "NUM_KSPLIT": 1 - }, "M_LEQ_2048": { "BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, diff --git a/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=7168-K=256.json b/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=7168-K=256.json index 790cd214de..f94ac349e4 100644 --- a/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=7168-K=256.json +++ b/aiter/ops/triton/configs/gfx950/gluon/gemm/gemm_a8w8_blockscale/GEMM-A8W8_BLOCKSCALE-N=7168-K=256.json @@ -1,52 +1,4 @@ { - "M_LEQ_16": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 16, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 2, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".ca", - "NUM_KSPLIT": 1 - }, - "M_LEQ_32": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 2, - "num_warps": 2, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".ca", - "NUM_KSPLIT": 1 - }, - "M_LEQ_64": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 4, - "num_warps": 2, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".ca", - "NUM_KSPLIT": 1 - }, - "M_LEQ_128": { - "BLOCK_SIZE_M": 32, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 4, - "num_warps": 2, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".ca", - "NUM_KSPLIT": 1 - }, "M_LEQ_2048": { "BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, diff --git a/aiter/ops/triton/gemm/basic/gemm_a8w8.py b/aiter/ops/triton/gemm/basic/gemm_a8w8.py index 19517c5d7d..f9fdfe3535 100644 --- a/aiter/ops/triton/gemm/basic/gemm_a8w8.py +++ b/aiter/ops/triton/gemm/basic/gemm_a8w8.py @@ -11,10 +11,19 @@ _gemm_a8w8_kernel, _get_config, ) +from aiter.ops.triton.utils._triton.arch_info import get_arch +from aiter.ops.triton.utils.device_info import get_num_xcds +from aiter.ops.triton.utils.gemm_config_utils import get_gemm_config from aiter.ops.triton.utils.logger import AiterTritonLogger +from aiter.ops.triton.utils.types import ( + get_scaled_dot_format_string, + torch_to_triton_dtype, +) _LOGGER = AiterTritonLogger() +_GLUON_SUPPORTED_ARCHS = ("gfx950",) + def gemm_a8w8( x: torch.Tensor, @@ -26,6 +35,7 @@ def gemm_a8w8( y: torch.Tensor | None = None, config: dict | None = None, skip_reduce: bool | None = False, + backend: str = "triton", ): """ Computes 8 bit matrix multiplication Y = (X @ W^T) * (x_scale * w_scale) with optional bias. @@ -60,12 +70,68 @@ def gemm_a8w8( w = w.T + assert backend in ( + "triton", + "gluon", + ), f"Unknown backend '{backend}', must be 'triton' or 'gluon'" + + if backend == "gluon": + assert ( + get_arch() in _GLUON_SUPPORTED_ARCHS + ), f"Gluon backend requires one of {_GLUON_SUPPORTED_ARCHS}, got '{get_arch()}'" + assert x.dtype == w.dtype, "Input types must be the same" + if config is None: - config, _ = _get_config(M, N, K) + if backend == "gluon": + config, _ = get_gemm_config("GEMM-A8W8", M, N, K, backend="gluon") + else: + config, _ = _get_config(M, N, K) - if y is None and (config["NUM_KSPLIT"] == 1 or not skip_reduce): + if y is None and (config.get("NUM_KSPLIT", 1) == 1 or not skip_reduce): y = torch.empty((M, N), dtype=dtype, device=x.device) + if backend == "gluon": + from aiter.ops.triton._gluon_kernels.gfx950.gemm.basic.gemm_a8w8 import ( + _gemm_a8w8_kernel as _gluon_gemm_a8w8_kernel, + ) + + _LOGGER.info( + f"GEMM_A8W8 [gluon/{get_arch()}]: x={tuple(x.shape)} w={tuple(w.shape)}" + ) + + fp8_format = ( + None + if x.dtype == torch.int8 + else get_scaled_dot_format_string(torch_to_triton_dtype[x.dtype]) + ) + grid = ( + triton.cdiv(M, config["BLOCK_SIZE_M"]) + * triton.cdiv(N, config["BLOCK_SIZE_N"]), + ) + _gluon_gemm_a8w8_kernel[grid]( + x, + w, + x_scale, + w_scale, + bias, + y, + M, + N, + K, + x.stride(0), + x.stride(1), + w.stride(0), + w.stride(1), + y.stride(0), + y.stride(1), + bias is not None, + NUM_XCDS=get_num_xcds(), + NUM_WARPS=config["num_warps"], + **config, + FP8_FORMAT=fp8_format, + ) + return y + if config["NUM_KSPLIT"] > 1: y_pp = torch.empty( (config["NUM_KSPLIT"], M, N), @@ -137,3 +203,92 @@ def gemm_a8w8( ) return y + + +def gemm_a8w8_preshuffle( + x: torch.Tensor, + w: torch.Tensor, + x_scale: torch.Tensor, + w_scale: torch.Tensor, + bias: torch.Tensor | None = None, + dtype: float | None = torch.bfloat16, + y: torch.Tensor | None = None, + config: dict | None = None, +): + """ + Computes 8 bit matrix multiplication Y = (X @ W^T) * (x_scale * w_scale) with optional bias, + taking weights in a pre-shuffled layout for better memory access. + + Args: + x (torch.Tensor): INT8/FP8 input matrix with shape (M, K). + w (torch.Tensor): INT8/FP8 weight matrix pre-shuffled to (N*16, K//16), + internally transposed. + x_scale (torch.Tensor): Scale factor for x with shape (M, 1) or (M,). + w_scale (torch.Tensor): Scale factor for w with shape (1, N) or (N,). + bias (Optional[torch.Tensor]): Bias vector with shape (N,). + dtype (Optional[torch.dtype]): Output datatype (BF16 or FP16). + y (Optional[torch.Tensor]): Pre-allocated output tensor with shape (M, N). + config (Optional[dict]): Kernel tuning parameters (BLOCK_SIZE_M, BLOCK_SIZE_N, + BLOCK_SIZE_K, GROUP_SIZE_M). + + Returns: + torch.Tensor: Output with shape (M, N) in higher precision format. + """ + assert ( + get_arch() in _GLUON_SUPPORTED_ARCHS + ), f"gemm_a8w8_preshuffle requires one of {_GLUON_SUPPORTED_ARCHS}, got '{get_arch()}'" + from aiter.ops.triton._gluon_kernels.gfx950.gemm.basic.gemm_a8w8 import ( + _gemm_a8w8_preshuffled_kernel as _gluon_gemm_a8w8_preshuffled_kernel, + ) + + _LOGGER.info( + f"GEMM_A8W8 PRESHUFFLE [gluon/{get_arch()}]: x={tuple(x.shape)} w={tuple(w.shape)}" + ) + + M, K = x.shape + N, K = w.shape + N = N * 16 + K = K // 16 + + if config is None: + config, _ = get_gemm_config("GEMM-A8W8", M, N, K, backend="gluon") + + if y is None: + y = torch.empty((M, N), dtype=dtype, device=x.device) + + assert ( + K % config["BLOCK_SIZE_K"] == 0 + ), "K must be multiple of BLOCK_SIZE_K for preshuffling" + + fp8_format = ( + None + if x.dtype == torch.int8 + else get_scaled_dot_format_string(torch_to_triton_dtype[x.dtype]) + ) + grid = ( + triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv(N, config["BLOCK_SIZE_N"]), + ) + _gluon_gemm_a8w8_preshuffled_kernel[grid]( + x, + w, + x_scale, + w_scale, + bias, + y, + M, + N, + K, + x.stride(0), + x.stride(1), + w.stride(0), + w.stride(1), + y.stride(0), + y.stride(1), + bias is not None, + NUM_XCDS=get_num_xcds(), + NUM_WARPS=config["num_warps"], + **config, + FP8_FORMAT=fp8_format, + ) + + return y diff --git a/aiter/ops/triton/gemm/basic/gemm_a8w8_blockscale.py b/aiter/ops/triton/gemm/basic/gemm_a8w8_blockscale.py index 7870b4791e..6cd6e1eaac 100644 --- a/aiter/ops/triton/gemm/basic/gemm_a8w8_blockscale.py +++ b/aiter/ops/triton/gemm/basic/gemm_a8w8_blockscale.py @@ -28,16 +28,10 @@ _FORCE_GFX1250_EX = os.environ.get("AITER_FORCE_GFX1250_EX", "0") == "1" _TRITON_VERSION = Version(triton.__version__) -_GLUON_SUPPORTED_ARCHS = ("gfx1250",) - - -def _is_gluon_available(): - """Check if the gluon backend is available for the current GPU architecture.""" - try: - arch = get_arch() - return any(s in arch for s in _GLUON_SUPPORTED_ARCHS) - except Exception: # noqa: BLE001 - return False +_GLUON_SUPPORTED_ARCHS = ("gfx950", "gfx1250") +_GLUON_PRESHUFFLE_ARCHS = ("gfx1250",) +_GLUON_DEFAULT_ARCHS = ("gfx1250",) +_GLUON_KERNEL_TYPES = ("bandwidth_bound", "compute_bound") def gemm_a8w8_blockscale( @@ -85,11 +79,18 @@ def gemm_a8w8_blockscale( w = w.T # (K, N) w_scale = w_scale.T # (scale_k, scale_n) - # Resolve backend up-front so the config is loaded from the backend's - # config dir (gemm//), falling back to the shared gemm/ dir. if backend is None: - backend = "gluon" if _is_gluon_available() else "triton" + backend = "gluon" if get_arch() in _GLUON_DEFAULT_ARCHS else "triton" backend = backend.lower() + assert backend in ( + "triton", + "gluon", + ), f"Unknown backend '{backend}', must be 'triton' or 'gluon'" + + if backend == "gluon": + assert ( + get_arch() in _GLUON_SUPPORTED_ARCHS + ), f"Gluon backend requires one of {_GLUON_SUPPORTED_ARCHS}, got '{get_arch()}'" if config is None: config, _ = _get_config(M, N, K, backend=backend) @@ -133,86 +134,67 @@ def gemm_a8w8_blockscale( ), # Effective launch grid dims: [NUM_KSPLIT, NUM_M_BLOCKS, NUM_N_BLOCKS] ) - if backend is None: - backend = "gluon" if _is_gluon_available() else "triton" - backend = backend.lower() - assert backend in ( - "triton", - "gluon", - ), f"Unknown backend '{backend}', must be 'triton' or 'gluon'" - + extra_constexpr = {} if backend == "gluon": + arch = get_arch() assert ( - _is_gluon_available() - ), f"Gluon backend requires one of {_GLUON_SUPPORTED_ARCHS}, got '{get_arch()}'" - from aiter.ops.triton._gluon_kernels.gfx1250.gemm.basic.gemm_a8w8_blockscale import ( - _KERNEL_MAP, - ) + kernel_type in _GLUON_KERNEL_TYPES + ), f"Unknown kernel_type '{kernel_type}', must be one of {list(_GLUON_KERNEL_TYPES)}" + if arch == "gfx950": + from aiter.ops.triton._gluon_kernels.gfx950.gemm.basic.gemm_a8w8_blockscale import ( + _gemm_a8w8_blockscale_kernel as gluon_kernel, + ) + + # gfx950 has one blockscale kernel, which serves both kernel types. + impl = gluon_kernel + extra_constexpr["NUM_WARPS"] = config["num_warps"] + extra_constexpr["NUM_STAGES"] = max(config.get("num_stages", 2), 2) + elif arch == "gfx1250": + from aiter.ops.triton._gluon_kernels.gfx1250.gemm.basic.gemm_a8w8_blockscale import ( + _KERNEL_MAP, + ) + + impl = _KERNEL_MAP[kernel_type] + warp_bases = [(0, 1)] + for i in range(int(math.log2(config["num_warps"] // 2))): + warp_bases.append((1 << i, 0)) + extra_constexpr["warp_bases"] = tuple(warp_bases) + config["NUM_BUFFERS"] = config.pop("num_stages", 1) + else: + raise AssertionError( + f"Gluon backend requires one of {_GLUON_SUPPORTED_ARCHS}, got '{arch}'" + ) - assert ( - kernel_type in _KERNEL_MAP - ), f"Unknown kernel_type '{kernel_type}', must be one of {list(_KERNEL_MAP.keys())}" _LOGGER.info( - f"GEMM_A8W8 BLOCKSCALE [gluon/gfx1250]: x={tuple(x.shape)} w={tuple(w.shape)} " - f"kernel={kernel_type}" - ) - - impl = _KERNEL_MAP[kernel_type] - extra_constexpr = {} - warp_bases = [(0, 1)] - for i in range(int(math.log2(config["num_warps"] // 2))): - warp_bases.append((1 << i, 0)) - extra_constexpr["warp_bases"] = tuple(warp_bases) - config["NUM_BUFFERS"] = config.pop("num_stages", 1) - - impl[grid]( - x, - w, - y if config["NUM_KSPLIT"] == 1 else y_pp, - x_scale, - w_scale, - M, - N, - K, - x.stride(0), - x.stride(1), - w.stride(0), - w.stride(1), - 0 if config["NUM_KSPLIT"] == 1 else y_pp.stride(0), - y.stride(0) if config["NUM_KSPLIT"] == 1 else y_pp.stride(1), - y.stride(1) if config["NUM_KSPLIT"] == 1 else y_pp.stride(2), - x_scale.stride(0), - x_scale.stride(1), - w_scale.stride(0), - w_scale.stride(1), - **config, - **extra_constexpr, + f"GEMM_A8W8 BLOCKSCALE [gluon/{arch}]: x={tuple(x.shape)} " + f"w={tuple(w.shape)} kernel={kernel_type}" ) else: impl = triton_gemm_a8w8_blockscale_kernel - impl[grid]( - x, - w, - y if config["NUM_KSPLIT"] == 1 else y_pp, - x_scale, - w_scale, - M, - N, - K, - x.stride(0), - x.stride(1), - w.stride(0), - w.stride(1), - 0 if config["NUM_KSPLIT"] == 1 else y_pp.stride(0), - y.stride(0) if config["NUM_KSPLIT"] == 1 else y_pp.stride(1), - y.stride(1) if config["NUM_KSPLIT"] == 1 else y_pp.stride(2), - x_scale.stride(0), - x_scale.stride(1), - w_scale.stride(0), - w_scale.stride(1), - **config, - ) + impl[grid]( + x, + w, + y if config["NUM_KSPLIT"] == 1 else y_pp, + x_scale, + w_scale, + M, + N, + K, + x.stride(0), + x.stride(1), + w.stride(0), + w.stride(1), + 0 if config["NUM_KSPLIT"] == 1 else y_pp.stride(0), + y.stride(0) if config["NUM_KSPLIT"] == 1 else y_pp.stride(1), + y.stride(1) if config["NUM_KSPLIT"] == 1 else y_pp.stride(2), + x_scale.stride(0), + x_scale.stride(1), + w_scale.stride(0), + w_scale.stride(1), + **config, + **extra_constexpr, + ) if config["NUM_KSPLIT"] > 1: if skip_reduce: @@ -301,7 +283,7 @@ def gemm_a8w8_blockscale_preshuffle( # Resolve backend up-front so the config is loaded from the backend's # config dir (gemm//), falling back to the shared gemm/ dir. if backend is None: - backend = "gluon" if _is_gluon_available() else "triton" + backend = "gluon" if get_arch() in _GLUON_PRESHUFFLE_ARCHS else "triton" backend = backend.lower() if config is None: @@ -371,9 +353,6 @@ def gemm_a8w8_blockscale_preshuffle( ) extra_constexpr = {} - if backend is None: - backend = "gluon" if _is_gluon_available(preshuffle=True) else "triton" - backend = backend.lower() assert backend in ( "triton", "gluon", @@ -381,8 +360,8 @@ def gemm_a8w8_blockscale_preshuffle( if backend == "gluon": assert ( - _is_gluon_available() - ), f"Gluon preshuffle requires one of {_GLUON_SUPPORTED_ARCHS}, got '{get_arch()}'" + get_arch() in _GLUON_PRESHUFFLE_ARCHS + ), f"Gluon preshuffle requires one of {_GLUON_PRESHUFFLE_ARCHS}, got '{get_arch()}'" from aiter.ops.triton._gluon_kernels.gfx1250.gemm.basic.gemm_a8w8_blockscale import ( _PRESHUFFLE_KERNEL_MAP, ) diff --git a/aiter/ops/triton/utils/types.py b/aiter/ops/triton/utils/types.py index 3515be9114..d4627036f0 100644 --- a/aiter/ops/triton/utils/types.py +++ b/aiter/ops/triton/utils/types.py @@ -71,6 +71,17 @@ def get_fp8_e4m3_dtype(): } +def get_scaled_dot_format_string(dtype: tl.dtype): + mapping = { + tl.float16: "fp16", + tl.bfloat16: "bf16", + tl.uint8: "e2m1", + tl.float8e4nv: "e4m3", + tl.float8e5: "e5m2", + } + return mapping[dtype] + + def _is_fp8(x): if x.dtype in { torch.float8_e4m3fnuz, diff --git a/op_tests/op_benchmarks/triton/bench_gemm_a8w8.py b/op_tests/op_benchmarks/triton/bench_gemm_a8w8.py index 6d663b8600..ef89248751 100644 --- a/op_tests/op_benchmarks/triton/bench_gemm_a8w8.py +++ b/op_tests/op_benchmarks/triton/bench_gemm_a8w8.py @@ -1,3 +1,4 @@ +import functools import math import sys from collections.abc import Callable @@ -5,10 +6,7 @@ import triton from aiter.ops.triton.gemm.basic.gemm_a8w8 import gemm_a8w8 as triton_gemm_a8w8 -from aiter.ops.triton.gluon.gemm_a8w8 import ( - gemm_a8w8 as gluon_gemm_a8w8, -) -from aiter.ops.triton.gluon.gemm_a8w8 import ( +from aiter.ops.triton.gemm.basic.gemm_a8w8 import ( gemm_a8w8_preshuffle as gluon_gemm_a8w8_preshuffle, ) from aiter.ops.triton.utils.types import str_to_torch_dtype @@ -131,7 +129,7 @@ def run_benchmark(args, defaults): if args.shuffle: impl = gluon_gemm_a8w8_preshuffle else: - impl = gluon_gemm_a8w8 + impl = functools.partial(triton_gemm_a8w8, backend="gluon") else: if args.shuffle: raise RuntimeError( diff --git a/op_tests/op_benchmarks/triton/bench_gemm_a8w8_blockscale.py b/op_tests/op_benchmarks/triton/bench_gemm_a8w8_blockscale.py index b6775a8962..67ff639ba9 100644 --- a/op_tests/op_benchmarks/triton/bench_gemm_a8w8_blockscale.py +++ b/op_tests/op_benchmarks/triton/bench_gemm_a8w8_blockscale.py @@ -1,3 +1,4 @@ +import functools import math from collections.abc import Callable @@ -10,9 +11,6 @@ from aiter.ops.triton.gemm.basic.gemm_a8w8_blockscale import ( gemm_a8w8_blockscale_preshuffle as triton_gemm_a8w8_blockscale_preshuffle, ) -from aiter.ops.triton.gluon.gemm_a8w8_blockscale import ( - gemm_a8w8_blockscale as gluon_gemm_a8w8_blockscale, -) from aiter.test_common import checkAllclose from op_tests.op_benchmarks.triton.utils.argparse import ( add_argparse_ff, @@ -156,7 +154,7 @@ def run_benchmark(args, defaults): args.shape and args.M ), "User can specify --shape or --model MODEL -M VAL exclusively" if args.gluon: - impl = gluon_gemm_a8w8_blockscale + impl = functools.partial(triton_gemm_a8w8_blockscale, backend="gluon") elif args.preshuffle: impl = triton_gemm_a8w8_blockscale_preshuffle else: diff --git a/op_tests/triton_tests/gemm/basic/test_gemm_a8w8.py b/op_tests/triton_tests/gemm/basic/test_gemm_a8w8.py index 09ef87ad8b..e780e92a29 100644 --- a/op_tests/triton_tests/gemm/basic/test_gemm_a8w8.py +++ b/op_tests/triton_tests/gemm/basic/test_gemm_a8w8.py @@ -1,18 +1,17 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +import functools +import importlib +import sys + import pytest import torch import torch.nn.functional as F from aiter.ops.shuffle import shuffle_weight from aiter.ops.triton.gemm.basic.gemm_a8w8 import gemm_a8w8 as triton_gemm_a8w8 -from aiter.ops.triton.gluon.gemm_a8w8 import ( - gemm_a8w8 as gluon_gemm_a8w8, -) -from aiter.ops.triton.gluon.gemm_a8w8 import ( - gemm_a8w8_preshuffle as gluon_gemm_a8w8_preshuffle, -) +from aiter.ops.triton.gemm.basic.gemm_a8w8 import gemm_a8w8_preshuffle from aiter.ops.triton.utils._triton import arch_info from aiter.ops.triton.utils.gemm_config_utils import ( compute_splitk_params, @@ -168,7 +167,7 @@ def test_gemm_fp8(in_dtype, m, n, k, impl: str): if impl in ["gluon", "gluon_shuffle"] and DEVICE_ARCH != "gfx950": pytest.skip( - "Gluon implementation is not supported on this device (requires CDNA4)." + "Gluon implementation is not supported on this device (requires gfx950)." ) if impl == "gluon_shuffle" and (n % 16 != 0 or k % 32 != 0): @@ -193,9 +192,9 @@ def test_gemm_fp8(in_dtype, m, n, k, impl: str): if impl == "triton": impl = triton_gemm_a8w8 elif impl == "gluon": - impl = gluon_gemm_a8w8 + impl = functools.partial(triton_gemm_a8w8, backend="gluon") elif impl == "gluon_shuffle": - impl = gluon_gemm_a8w8_preshuffle + impl = gemm_a8w8_preshuffle else: raise ValueError(f"Unknown implementation: {impl}") b = run_triton(x, weight_triton, x_scale, w_scale, bias, out_dtype, y, impl) @@ -229,7 +228,7 @@ def test_gemm_int8(out_dtype, m, n, k, layout, output, impl: str): if impl in ["gluon", "gluon_shuffle"] and DEVICE_ARCH != "gfx950": pytest.skip( - "Gluon implementation is not supported on this device (requires CDNA4)." + "Gluon implementation is not supported on this device (requires gfx950)." ) if impl == "gluon_shuffle" and (n % 16 != 0 or k % 32 != 0): @@ -254,9 +253,9 @@ def test_gemm_int8(out_dtype, m, n, k, layout, output, impl: str): if impl == "triton": impl = triton_gemm_a8w8 elif impl == "gluon": - impl = gluon_gemm_a8w8 + impl = functools.partial(triton_gemm_a8w8, backend="gluon") elif impl == "gluon_shuffle": - impl = gluon_gemm_a8w8_preshuffle + impl = gemm_a8w8_preshuffle else: raise ValueError(f"Unknown implementation: {impl}") b = run_triton(x, weight_triton, x_scale, w_scale, bias, out_dtype, y, impl) @@ -377,3 +376,14 @@ def test_gemm_splitk_skip_reduce(in_dtype, out_dtype, m, n, k, num_ksplit): b = y_pp.sum(dim=0).to(out_dtype) torch.testing.assert_close(a, b, atol=0.02, rtol=1e-2) + + +def test_legacy_gluon_import_path_warns(): + """The pre-move path still resolves here, but tells callers to move on.""" + legacy = "aiter.ops.triton.gluon.gemm_a8w8" + sys.modules.pop(legacy, None) + + with pytest.warns(DeprecationWarning, match="has moved to"): + mod = importlib.import_module(legacy) + + assert mod.gemm_a8w8.__module__ == "aiter.ops.triton.gemm.basic.gemm_a8w8" diff --git a/op_tests/triton_tests/gemm/basic/test_gemm_a8w8_blockscale.py b/op_tests/triton_tests/gemm/basic/test_gemm_a8w8_blockscale.py index fe99de006b..c2fd8a2dba 100644 --- a/op_tests/triton_tests/gemm/basic/test_gemm_a8w8_blockscale.py +++ b/op_tests/triton_tests/gemm/basic/test_gemm_a8w8_blockscale.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +import importlib +import sys + import pytest import torch import torch.nn.functional as F @@ -10,9 +13,6 @@ gemm_a8w8_blockscale, gemm_a8w8_blockscale_preshuffle, ) -from aiter.ops.triton.gluon.gemm_a8w8_blockscale import ( - gemm_a8w8_blockscale as gluon_gfx950_gemm_a8w8_blockscale, -) from aiter.ops.triton.utils._triton import arch_info from aiter.ops.triton.utils.types import get_fp8_dtypes, str_to_torch_dtype @@ -185,21 +185,30 @@ def test_gemm(dtype, M, N, K, layout, output, backend, shuffle): a = run_torch(x, weight, x_scale, w_scale, dtype) - if not shuffle and backend == "gluon" and DEVICE_ARCH == "gfx950": - impl = gluon_gfx950_gemm_a8w8_blockscale - else: - if shuffle: + if shuffle: - def impl(x, w, xs, ws, dt, y): - return gemm_a8w8_blockscale_preshuffle( - x, w, xs, ws, dt, y, backend=backend - ) + def impl(x, w, xs, ws, dt, y): + return gemm_a8w8_blockscale_preshuffle(x, w, xs, ws, dt, y, backend=backend) - else: + else: - def impl(x, w, xs, ws, dt, y): - return gemm_a8w8_blockscale(x, w, xs, ws, dt, y, backend=backend) + def impl(x, w, xs, ws, dt, y): + return gemm_a8w8_blockscale(x, w, xs, ws, dt, y, backend=backend) b = run_triton(x, weight_triton, x_scale_shuffled, w_scale, dtype, y, impl) torch.testing.assert_close(a, b, atol=0.01, rtol=1e-2) + + +def test_legacy_gluon_import_path_warns(): + """The pre-move path still resolves here, but tells callers to move on.""" + legacy = "aiter.ops.triton.gluon.gemm_a8w8_blockscale" + sys.modules.pop(legacy, None) + + with pytest.warns(DeprecationWarning, match="has moved to"): + mod = importlib.import_module(legacy) + + assert ( + mod.gemm_a8w8_blockscale.__module__ + == "aiter.ops.triton.gemm.basic.gemm_a8w8_blockscale" + )