# 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.common.splitk_reduce import ( _gemm_splitk_reduce_kernel, ) from ..._triton_kernels.gemm.basic.gemm_a16w8_blockscale import ( _gemm_a16w8_blockscale_kernel, _gemm_a16w8_blockscale_preshuffle_kernel, _get_config, ) from ...utils.logger import AiterTritonLogger from ...utils.gemm_config_utils import compute_splitk_params _LOGGER = AiterTritonLogger() def gemm_a16w8_blockscale( x: torch.Tensor, w: torch.Tensor, w_scale: torch.Tensor, dtype: Optional[float] = torch.bfloat16, y: Optional[torch.Tensor] = None, prequant: Optional[bool] = False, config: Optional[dict] = None, skip_reduce: Optional[bool] = False, ): """ 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). - 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)} 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 config is None: config, _ = _get_config(M, N, K) return_y_pp = config["NUM_KSPLIT"] > 1 and skip_reduce if config["NUM_KSPLIT"] > 1: y_pp = torch.empty( (config["NUM_KSPLIT"], M, N), dtype=torch.float32, device=x.device ) else: y_pp = None if y is None and not return_y_pp: y = torch.empty((M, N), dtype=dtype, device=x.device) compute_splitk_params(config, 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])) DTYPE_MAX = ( torch.finfo(w.dtype).max if torch.is_floating_point(w) else torch.iinfo(w.dtype).max ) # grid = (config["NUM_KSPLIT"], triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv(N, config["BLOCK_SIZE_N"]),) grid = lambda META: ( # noqa: E731 ( META["NUM_KSPLIT"] * triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]) ), ) _gemm_a16w8_blockscale_kernel[grid]( x, w, y if config["NUM_KSPLIT"] == 1 else y_pp, 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), w_scale.stride(0), w_scale.stride(1), PREQUANT=prequant, DTYPE_MAX=DTYPE_MAX, DTYPE_MIN=-DTYPE_MAX, **config, ) if return_y_pp: return y_pp elif 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_splitk_reduce_kernel[grid_reduce]( y_pp, y, None, 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=False, activation="", use_activation=False, KERNEL_NAME="_gemm_a8w8_blockscale_reduce_kernel", ) return y def gemm_a16w8_blockscale_preshuffle( x: torch.Tensor, w: torch.Tensor, w_scale: torch.Tensor, dtype: Optional[float] = torch.bfloat16, y: Optional[torch.Tensor] = None, prequant: Optional[bool] = False, config: Optional[dict] = None, skip_reduce: Optional[bool] = False, ): """ 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). - 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)} w_scale={tuple(w_scale.shape)}" ) M, K = x.shape N, K = w.shape N = N * 16 K = K // 16 # Check constraints. assert x.shape[1] == w.shape[1] // 16, "Incompatible dimensions!!!" if config is None: config, _ = _get_config(M, N, K, True) return_y_pp = config["NUM_KSPLIT"] > 1 and skip_reduce if config["NUM_KSPLIT"] > 1: y_pp = torch.empty( (config["NUM_KSPLIT"], M, N), dtype=torch.float32, device=x.device ) else: y_pp = None if y is None and not return_y_pp: y = torch.empty((M, N), dtype=dtype, device=x.device) compute_splitk_params(config, 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[1]) ) # scale_block_size_k config["GROUP_N"] = triton.next_power_of_2( triton.cdiv(N, w_scale.shape[0]) ) # scale_block_size_n assert ( config["GROUP_K"] == config["BLOCK_SIZE_K"] ), "GROUP_K must equal BLOCK_SIZE_K" DTYPE_MAX = ( torch.finfo(w.dtype).max if torch.is_floating_point(w) else torch.iinfo(w.dtype).max ) # grid = (config["NUM_KSPLIT"], triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv(N, config["BLOCK_SIZE_N"]),) grid = lambda META: ( # noqa: E731 ( META["NUM_KSPLIT"] * triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]) ), ) _gemm_a16w8_blockscale_preshuffle_kernel[grid]( x, w, y if config["NUM_KSPLIT"] == 1 else y_pp, 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), w_scale.stride(0), w_scale.stride(1), PREQUANT=prequant, DTYPE_MAX=DTYPE_MAX, DTYPE_MIN=-DTYPE_MAX, **config, ) if return_y_pp: return y_pp elif 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_splitk_reduce_kernel[grid_reduce]( y_pp, y, None, 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=False, activation="", use_activation=False, KERNEL_NAME="_gemm_a8w8_blockscale_reduce_kernel", ) return y