Hugging Face's logo

Kernels:
kernels-community
/
aiter-kernels
Trusted publisher

Kernel card Files
xet
 Community
aiter-kernels / build /torch-rocm /gemm /basic /gemm_a16w8_blockscale.py
kernels-bot's picture
kernels-bot
Uploaded using `kernel-builder`.
89181fc verified
Raw
History Blame
8.28 kB
# 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