[Triton/Gluon] Move gluon gemm_a8w8 kernel into _gluon_kernels/gfx950

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)

aiter/ops/triton/gluon/gemm_a8w8.py → 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