# SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. from typing import Optional import torch import triton from ..._triton_kernels.gemm.basic.gemm_a16w16 import ( _gemm_a16_w16_kernel, _get_config as _get_triton_config, ) from ..._triton_kernels.common.splitk_reduce import ( _gemm_splitk_reduce_kernel, ) from ..._triton_kernels.activation import _get_activation_from_str from ...utils.gemm_config_utils import get_gemm_config from ...utils.logger import AiterTritonLogger from ...utils._triton.arch_info import get_arch from ...utils.common_utils import serialize_dict, deserialize_str from ..._aiter_compat.torch_guard import torch_compile_guard _LOGGER = AiterTritonLogger() _GLUON_SUPPORTED_ARCHS = ("gfx1250",) def _is_gluon_available(): """Check if the gluon backend is available for the current GPU architecture.""" try: return any(supported in get_arch() for supported in _GLUON_SUPPORTED_ARCHS) except Exception: return False def gemm_a16w16_fake_tensor( x: torch.Tensor, w: torch.Tensor, bias: Optional[torch.Tensor] = None, dtype: Optional[torch.dtype] = torch.bfloat16, y: Optional[torch.Tensor] = None, config: Optional[str] = None, activation: Optional[str] = None, skip_reduce: Optional[bool] = False, kernel_type: str = "bandwidth_bound", backend: Optional[str] = None, ) -> torch.Tensor: M, K = x.shape N, _ = w.shape # [triton only] split-K with skip_reduce returns the unreduced partials. if skip_reduce: cfg = deserialize_str(config) if config else _get_triton_config(M, N, K)[0] num_ksplit = cfg.get("NUM_KSPLIT", 1) if num_ksplit > 1: return torch.empty((num_ksplit, M, N), dtype=torch.float32, device=x.device) if y is not None: return y return torch.empty((M, N), dtype=dtype, device=x.device) @torch_compile_guard(gen_fake=gemm_a16w16_fake_tensor) def gemm_a16w16_( x: torch.Tensor, w: torch.Tensor, bias: Optional[torch.Tensor] = None, dtype: Optional[torch.dtype] = torch.bfloat16, y: Optional[torch.Tensor] = None, config: Optional[str] = None, activation: Optional[str] = None, skip_reduce: Optional[bool] = False, kernel_type: str = "bandwidth_bound", backend: Optional[str] = None, ) -> torch.Tensor: """ Computes 16 bit matrix multiplication Y = X @ W^T Uses the gluon backend automatically on supported architectures (gfx1250) and the triton backend everywhere else. Pass ``backend`` to force a choice. Args: x (torch.Tensor): Input matrix with shape (M, K). w (torch.Tensor): Weight matrix with shape (N, K), internally transposed. 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[str]): Serialized kernel tuning parameters. activation (Optional[str]): Activation function ("gelu", "gelu_tanh", "silu", "silu_exp2", "relu"). skip_reduce (Optional[bool]): [triton only] Skip reduction of split-K partial results. Returns shape (NUM_KSPLIT, M, N) instead of (M, N). kernel_type (str): [gluon only] Kernel variant ("bandwidth_bound", "compute_bound"). backend (Optional[str]): "triton", "gluon", or None (auto-detect). Returns: torch.Tensor: Output with shape (M, N) or (NUM_KSPLIT, M, N) if skip_reduce=True. """ config = deserialize_str(config) if config is not None else None 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'" if backend == "gluon": assert ( _is_gluon_available() ), f"Gluon backend requires one of {_GLUON_SUPPORTED_ARCHS}, got '{get_arch()}'" from ..._gluon_kernels.gfx1250.gemm.basic.gemm_a16w16 import ( _KERNEL_MAP, create_shared_layouts, create_wmma_layouts, ) assert ( kernel_type in _KERNEL_MAP ), f"Unknown kernel_type '{kernel_type}', must be one of {list(_KERNEL_MAP.keys())}" _LOGGER.info( f"GEMM_A16W16 [gluon/gfx1250]: x={tuple(x.shape)} w={tuple(w.shape)} " f"kernel={kernel_type}" ) assert x.dtype in ( torch.float16, torch.bfloat16, ), f"Activations (x) must be fp16 or bf16, got {x.dtype}" assert w.dtype in ( torch.float16, torch.bfloat16, ), f"Weights (w) must be fp16 or bf16, got {w.dtype}" assert x.shape[1] == w.shape[1], "Incompatible matrix shapes." M, _ = x.shape K = x.stride(0) N, _ = w.shape if config is None: config, _ = get_gemm_config("GEMM-A16W16", M, N, K) BLOCK_M = config["BLOCK_M"] BLOCK_N = config["BLOCK_N"] BLOCK_K = config["BLOCK_K"] NUM_BUFFERS = config.get("NUM_BUFFERS", 2) num_warps = config["num_warps"] # The kernels walk K with update_tensor_descriptor(add_offsets=...), # which advances the load position without shrinking the descriptor's # OOB bound, so a partial last K-tile would read past the end of K. # Require K to be a multiple of BLOCK_K rather than padding here. (M and # N may be unaligned: their descriptor bounds + store mask zero-fill the # partial tiles.) assert ( K % BLOCK_K == 0 ), f"K ({K}) must be a multiple of BLOCK_K ({BLOCK_K}) for the gluon a16w16 GEMM" # Clamp the software-pipeline depth to the number of K-tiles. # # The prologue/epilogue walk a fixed number of K-tiles determined by # NUM_BUFFERS, independent of how many real tiles exist. If NUM_BUFFERS # exceeds that count the pipeline loop counts go negative, so cap the # depth at the real tile count. Variants differ in reach and in the # minimum depth they require: # bandwidth_bound : reaches num_k_tiles -> cap = num_k_tiles # compute_bound : preloads one tile ahead (needs num_k_tiles >= NB + 1) # -> cap = num_k_tiles - 1 num_k_tiles = triton.cdiv(K, BLOCK_K) _MIN_BUFFERS = {"bandwidth_bound": 1, "compute_bound": 2} _DEPTH_SLACK = {"compute_bound": 1} # Fall back to the bandwidth_bound kernel when the requested variant cannot # satisfy its minimum pipeline depth for this K. The bandwidth_bound kernel # has no such floor (min depth 1) and is valid for every K, so we downgrade # rather than error. depth_cap = num_k_tiles - _DEPTH_SLACK.get(kernel_type, 0) if depth_cap < _MIN_BUFFERS[kernel_type]: needed = _MIN_BUFFERS[kernel_type] + _DEPTH_SLACK.get(kernel_type, 0) _LOGGER.info( f"GEMM_A16W16 [gluon/gfx1250]: kernel_type='{kernel_type}' needs " f"num_k_tiles>={needed} but num_k_tiles={num_k_tiles} " f"(K={K}, BLOCK_K={BLOCK_K}); falling back to kernel_type='bandwidth_bound'." ) kernel_type = "bandwidth_bound" depth_cap = num_k_tiles # bandwidth_bound: depth slack 0, min depth 1 NUM_BUFFERS = min(NUM_BUFFERS, depth_cap) w = w.T # Operand layout in BLAS TT/TN/NT/NN form: 'T' (row-major, trailing dim # contiguous) or 'N' (column-major, leading dim contiguous). First char # is x (A), second is w (B, after the internal transpose above). if x.stride(1) == 1: layout = "T" elif x.stride(0) == 1: layout = "N" else: raise ValueError( f"x must be contiguous in at least one dimension, got strides {x.stride()}" ) if w.stride(1) == 1: layout += "T" elif w.stride(0) == 1: layout += "N" else: raise ValueError( f"w must be contiguous in at least one dimension, got strides {w.stride()}" ) if y is None: y = torch.empty((M, N), dtype=dtype, device=x.device) wmma_layout, operand_a, operand_b = create_wmma_layouts(num_warps) shared_a, shared_b = create_shared_layouts(BLOCK_M, BLOCK_N, BLOCK_K, layout) grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N), 1) _KERNEL_MAP[kernel_type][grid]( x, w, y, bias, M, N, K, x.stride(0), x.stride(1), w.stride(0), w.stride(1), y.stride(0), y.stride(1), BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, NUM_BUFFERS=NUM_BUFFERS, LAYOUT=layout, SHARED_LAYOUT_A=shared_a, SHARED_LAYOUT_B=shared_b, WMMA_LAYOUT=wmma_layout, OPERAND_LAYOUT_A=operand_a, OPERAND_LAYOUT_B=operand_b, activation=_get_activation_from_str(activation) if activation else None, USE_ACTIVATION=activation is not None, ADD_BIAS=(bias is not None), num_warps=num_warps, ) return y _LOGGER.info(f"GEMM_A16W16 [triton]: x={tuple(x.shape)} w={tuple(w.shape)}") assert x.shape[1] == w.shape[1], "Incompatible matrix shapes." M, K = x.shape N, K = w.shape w = w.T if config is None: config, _ = _get_triton_config(M, N, K) if y is None and (config["NUM_KSPLIT"] == 1 or not skip_reduce): y = torch.empty((M, N), dtype=dtype, device=x.device) if config["NUM_KSPLIT"] > 1: y_pp = torch.empty( (config["NUM_KSPLIT"], M, N), dtype=torch.float32, device=y.device if y is not None else x.device, ) else: y_pp = None grid = lambda META: ( # noqa: E731 ( META["NUM_KSPLIT"] * triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]) ), ) _gemm_a16_w16_kernel[grid]( x, w, bias, y if config["NUM_KSPLIT"] == 1 else y_pp, 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), activation=_get_activation_from_str(activation) if activation else "", use_activation=activation is not None, ADD_BIAS=(bias is not None), SKIP_REDUCE=skip_reduce, **config, ) if config["NUM_KSPLIT"] > 1: if skip_reduce: return y_pp 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_splitk_reduce_kernel[grid_reduce]( y_pp, y, bias, 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"]), ADD_BIAS=(bias is not None), activation=_get_activation_from_str(activation) if activation else "", use_activation=activation is not None, KERNEL_NAME="_gemm_a16w16_reduce_kernel", ) return y def gemm_a16w16( x, w, bias: Optional[torch.Tensor] = None, dtype: Optional[torch.dtype] = torch.bfloat16, y: Optional[torch.Tensor] = None, config: Optional[dict] = None, activation: Optional[str] = None, skip_reduce: Optional[bool] = False, kernel_type: str = "bandwidth_bound", backend: Optional[str] = None, ): """ Computes 16 bit matrix multiplication Y = X @ W^T Uses the gluon backend automatically on supported architectures (gfx1250) and the triton backend everywhere else. Pass ``backend`` to force a choice. See ``gemm_a16w16_`` for the full argument description; ``config`` is a dict here and is serialized before dispatch so the op is torch.compile-traceable. """ # dtype must be a torch.dtype at the custom-op boundary (callers sometimes # pass a placeholder when a preallocated y already fixes the output dtype). if not isinstance(dtype, torch.dtype): dtype = torch.bfloat16 config_hashable = serialize_dict(config) if config else None return gemm_a16w16_( x, w, bias, dtype, y, config_hashable, activation, skip_reduce, kernel_type, backend, )