Kernels:

kernels-community
/

aiter-kernels

Trusted publisher

Kernel card Files Files and versions

xet

Community

aiter-kernels / build /torch-rocm /gemm /basic /gemm_a16w16_atomic.py

kernels-bot

Uploaded using `kernel-builder`.

89181fc verified 25 days ago

Raw

History Blame

3.69 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.gemm.basic.gemm_a16w16_atomic import (
	_gemm_a16_w16_atomic_kernel,
	_get_config,
	)
	from ...utils.logger import AiterTritonLogger
	from ...utils.common_utils import serialize_dict, deserialize_str
	from ..._aiter_compat.torch_guard import torch_compile_guard

	_LOGGER = AiterTritonLogger()


	def gemm_a16w16_atomic_fake_tensor(
	x: torch.Tensor,
	w: torch.Tensor,
	dtype: Optional[torch.dtype] = torch.bfloat16,
	y: Optional[torch.Tensor] = None,
	config: Optional[str] = None,
	) -> torch.Tensor:
	if y is None:
	M, _ = x.shape
	_, N = w.shape
	return torch.zeros((M, N), dtype=dtype, device=x.device)
	return y


	@torch_compile_guard(gen_fake=gemm_a16w16_atomic_fake_tensor)
	def gemm_a16w16_atomic_(
	x: torch.Tensor,
	w: torch.Tensor,
	dtype: Optional[torch.dtype] = torch.bfloat16,
	y: Optional[torch.Tensor] = None,
	config: Optional[str] = None,
	) -> torch.Tensor:
	"""
	Computes 16 bit matrix multiplication Y = X @ W^T using atomic operations for split-K reduction.

	Args:
	x (torch.Tensor): BF16/FP16 input matrix matrix with shape (M, K).
	w (torch.Tensor): BF16/FP16 weight matrix with shape (N, K), internally transposed.
	dtype (Optional[torch.dtype]): Output datatype (BF16 or FP16).
	Note: BF16 atomic aggregation may have slight precision loss.
	y (Optional[torch.Tensor]): Pre-allocated output tensor with shape (M, N).
	Must be zero-initialized for split-K (NUM_KSPLIT > 1).
	config (Optional[str]): Kernel tuning parameters (BLOCK_SIZE_M, BLOCK_SIZE_N,
	BLOCK_SIZE_K, GROUP_SIZE_M, NUM_KSPLIT, cache_modifier).

	Returns:
	y (torch.Tensor): Output with shape (M, N).
	"""
	_LOGGER.info(
	f"GEMM_A16W16_ATOMIC: x.shape={tuple(x.shape)}, w.shape={tuple(w.shape)} "
	)

	w = w.T

	M, K = x.shape
	K, N = w.shape

	if config is None:
	config, _ = _get_config(M, N, K)
	else:
	config = deserialize_str(config)

	# For compatability reasons, these keys may not exist in the config
	# TODO: This needs to be embedded in the configs later
	if "NUM_KSPLIT" not in config:
	config["NUM_KSPLIT"] = 1
	if "cache_modifier" not in config:
	config["cache_modifier"] = ""

	if y is None:
	# atomic add requires 0 tensor
	if config["NUM_KSPLIT"] == 1:
	y = torch.empty((M, N), dtype=dtype, device=x.device)
	else:
	y = torch.zeros((M, N), dtype=dtype, device=x.device)

	grid = lambda META: ( # noqa: E731
	triton.cdiv(M, META["BLOCK_SIZE_M"])
	* triton.cdiv(N, META["BLOCK_SIZE_N"])
	* META["NUM_KSPLIT"],
	)
	# NOTE: if k split doesnt divide K evenly, this will waste compute
	SPLITK_BLOCK_SIZE = triton.cdiv(K, config["NUM_KSPLIT"])
	config["SPLITK_BLOCK_SIZE"] = SPLITK_BLOCK_SIZE
	_gemm_a16_w16_atomic_kernel[grid](
	x,
	w,
	y,
	M,
	N,
	K,
	x.stride(0),
	x.stride(1),
	w.stride(0),
	w.stride(1),
	y.stride(0),
	y.stride(1),
	**config,
	)

	return y


	def gemm_a16w16_atomic(
	x: torch.Tensor,
	w: torch.Tensor,
	dtype: Optional[torch.dtype] = torch.bfloat16,
	y: Optional[torch.Tensor] = None,
	config: Optional[dict] = None,
	) -> torch.Tensor:
	config_hashable = serialize_dict(config) if config else None
	return gemm_a16w16_atomic_(x, w, dtype, y, config_hashable)