File size: 3,688 Bytes
2976eec | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | # 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)
|