Skip to content

Commit 870b9e4

Browse files
lalala-shlalala-sh
authored
integrate deep gemm (ROCm#1265)
* integrate m grouped gemm * update ck * add limit for 950 * rename deepgeem
1 parent 02b6c7f commit 870b9e4

13 files changed

Lines changed: 1039 additions & 2 deletions

File tree

‎3rdparty/composable_kernel‎

Submodule composable_kernel updated 210 files

‎aiter/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ def getLogger():
5757
from .ops.gemm_op_a4w4 import *
5858
from .ops.batched_gemm_op_a8w8 import *
5959
from .ops.batched_gemm_op_bf16 import *
60+
from .ops.deepgemm import *
6061
from .ops.aiter_operator import *
6162
from .ops.activation import *
6263
from .ops.attention import *

‎aiter/jit/optCompilerConfig.json‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -269,6 +269,23 @@
269269
"is_standalone": "False",
270270
"blob_gen_cmd": "f'{AITER_CSRC_DIR}/ck_gemm_a8w8_bpreshuffle/gen_instances.py --working_path {{}} --tune_file {AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_FILE}'"
271271
},
272+
"module_deepgemm": {
273+
"srcs": [
274+
"f'{AITER_CSRC_DIR}/pybind/deepgemm_pybind.cu'",
275+
"f'{AITER_CSRC_DIR}/ck_deepgemm/deepgemm.cu'"
276+
277+
],
278+
"flags_extra_cc": [],
279+
"flags_extra_hip": [],
280+
"md_name": "'module_deepgemm'",
281+
"extra_ldflags": "None",
282+
"extra_include": ["f'{CK_DIR}/example/ck_tile/18_flatmm'", "f'{AITER_CSRC_DIR}/ck_deepgemm/include'"],
283+
"verbose": "False",
284+
"is_python_module": "True",
285+
"is_standalone": "False",
286+
"hip_clang_path": "os.environ.get('FLATMM_HIP_CLANG_PATH')",
287+
"blob_gen_cmd": "f'{AITER_CSRC_DIR}/ck_deepgemm/gen_instances.py --working_path {{}}'"
288+
},
272289
"module_gemm_a8w8_asm": {
273290
"srcs": [
274291
"f'{AITER_CSRC_DIR}/pybind/gemm_a8w8_asm_pybind.cu'",

‎aiter/ops/deepgemm.py‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
# SPDX-License-Identifier: MIT
2+
# Copyright (C) 2025, Advanced Micro Devices, Inc. All rights reserved.
3+
4+
from torch import Tensor
5+
from typing import Optional
6+
from ..jit.core import (
7+
compile_ops,
8+
)
9+
10+
11+
@compile_ops("module_deepgemm", fc_name="deepgemm")
12+
def deepgemm_ck(
13+
XQ: Tensor,
14+
WQ: Tensor,
15+
Y: Tensor,
16+
group_layout: Tensor,
17+
x_scale: Optional[Tensor] = None,
18+
w_scale: Optional[Tensor] = None,
19+
) -> Tensor: ...
20+
21+
22+
def deepgemm(
23+
XQ: Tensor,
24+
WQ: Tensor,
25+
Y: Tensor,
26+
group_layout: Tensor,
27+
x_scale: Optional[Tensor] = None,
28+
w_scale: Optional[Tensor] = None,
29+
):
30+
return deepgemm_ck(XQ, WQ, Y, group_layout, x_scale, w_scale)

‎csrc/ck_deepgemm/deepgemm.cu‎

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
// SPDX-License-Identifier: MIT
2+
// Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
3+
4+
#include "deepgemm_common.cuh"
5+
#include "deepgemm_lookup.h"
6+
#include "deepgemm_manifest.h"
7+
#include <cmath>
8+
#include "py_itfs_common.h"
9+
10+
using RowwiseKernel = std::function<
11+
torch::Tensor(torch::Tensor &, torch::Tensor &,
12+
torch::Tensor &, torch::Tensor &,
13+
std::optional<torch::Tensor>, std::optional<torch::Tensor>)>;
14+
15+
// Define a custom hash function for std::tuple<int, int, int>
16+
struct IntTupleHash
17+
{
18+
size_t operator()(const std::tuple<int, int, int> &t) const
19+
{
20+
auto hash1 = std::hash<int>{}(std::get<0>(t));
21+
auto hash2 = std::hash<int>{}(std::get<1>(t));
22+
auto hash3 = std::hash<int>{}(std::get<2>(t));
23+
return hash1 ^ hash2 ^ hash3;
24+
}
25+
};
26+
27+
// For certain high priority shapes, we directly use the best kernel rather
28+
// than use heuristics.
29+
using RowwiseKernelMap = std::unordered_map<
30+
std::tuple<int, int, int>,
31+
RowwiseKernel,
32+
IntTupleHash>;
33+
34+
template <typename ABDataType, typename AccDataType, typename CDataType>
35+
RowwiseKernel rowwise_heuristic_dispatch(int M, int N, int K)
36+
{
37+
// Apply shape heuristics to find a suitable kernel implementation.
38+
if (M < 128)
39+
{
40+
return deepgemm_256x32x64x256_16x16x64_1x4<ABDataType, AccDataType, CDataType>;
41+
}
42+
else
43+
{
44+
return deepgemm_256x128x128x128_16x16x64_1x4<ABDataType, AccDataType, CDataType>;
45+
}
46+
}
47+
48+
// Helper function to return the next largest power of 2
49+
static constexpr int nextPow2(unsigned int num)
50+
{
51+
if (num <= 1)
52+
return 1;
53+
return 1 << (CHAR_BIT * sizeof(num) - __builtin_clz(num - 1));
54+
}
55+
56+
template <typename ABDataType, typename AccDataType, typename CDataType>
57+
RowwiseKernel rowwise_dispatch(int M, int N, int K)
58+
{
59+
// TODO: add tuner @lalala-sh
60+
// For a given shape, either find the best kernel via lookup or heuristic.
61+
// For many small M shapes, we bucket them to the next largest kernel.
62+
// This is fine since kernels are padded anyway.
63+
64+
// static const auto lookup = [&]
65+
// {
66+
// return RowwiseKernelMap{GENERATE_LOOKUP_TABLE(ABDataType, AccDataType, CDataType)};
67+
// }();
68+
69+
// // First check if this shape(M,N,K) is available in the direct lookup.
70+
// auto it = lookup.find({M, N, K});
71+
// // If we found an optimal kernel, use it.
72+
// if (it != lookup.end())
73+
// {
74+
// return it->second;
75+
// }
76+
77+
// int padded_m = M;
78+
// if (M > 1 && M <= 16)
79+
// {
80+
// padded_m = 16;
81+
// }
82+
// else if (M <= 16384)
83+
// {
84+
// padded_m = nextPow2(M);
85+
// }
86+
// else if (M <= 20480)
87+
// {
88+
// padded_m = 20480;
89+
// }
90+
// // Second check if this shape(padded_m,N,K) is available in the direct lookup.
91+
// it = lookup.find({padded_m, N, K});
92+
// // If we found an optimal kernel, use it.
93+
// if (it != lookup.end())
94+
// {
95+
// return it->second;
96+
// }
97+
// Otherwise, use heuristics.
98+
return rowwise_heuristic_dispatch<ABDataType, AccDataType, CDataType>(M, N, K);
99+
}
100+
101+
torch::Tensor deepgemm(
102+
torch::Tensor &XQ,
103+
torch::Tensor &WQ,
104+
torch::Tensor &Y,
105+
torch::Tensor &grouped_layout,
106+
std::optional<torch::Tensor> x_scale,
107+
std::optional<torch::Tensor> w_scale)
108+
{
109+
TORCH_CHECK(XQ.dtype() == WQ.dtype(),
110+
"Weights and activations should both be int8/fp8!");
111+
if (x_scale != std::nullopt && w_scale != std::nullopt)
112+
TORCH_CHECK(x_scale.value().dtype() == w_scale.value().dtype(),
113+
"Scales should have the same dtype!");
114+
115+
int M = XQ.size(0);
116+
int N = WQ.size(0);
117+
int K = XQ.size(1);
118+
int KBatch = 1;
119+
120+
121+
122+
if (XQ.dtype() == at::ScalarType::BFloat16 || XQ.dtype() == at::ScalarType::Half)
123+
{
124+
if (XQ.dtype() == at::ScalarType::Half)
125+
{
126+
rowwise_dispatch<fp16, float, fp16>(M, N, K)(XQ, WQ, Y, grouped_layout, x_scale, w_scale);
127+
}
128+
else
129+
{
130+
rowwise_dispatch<bf16, float, bf16>(M, N, K)(XQ, WQ, Y, grouped_layout, x_scale, w_scale);
131+
}
132+
}
133+
else if (XQ.dtype() == torch_fp8)
134+
{
135+
if (Y.dtype() == at::ScalarType::Half)
136+
{
137+
rowwise_dispatch<fp8, float, fp16>(M, N, K)(XQ, WQ, Y, grouped_layout, x_scale, w_scale);
138+
}
139+
else if (Y.dtype() == at::ScalarType::BFloat16)
140+
{
141+
rowwise_dispatch<fp8, float, bf16>(M, N, K)(XQ, WQ, Y, grouped_layout, x_scale, w_scale);
142+
}
143+
}
144+
else
145+
{
146+
TORCH_CHECK(false, "Unsupported scales/output dtype!");
147+
}
148+
return Y;
149+
}

‎csrc/ck_deepgemm/deepgemm_common.py‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
# SPDX-License-Identifier: MIT
2+
# Copyright (C) 2025, Advanced Micro Devices, Inc. All rights reserved.
3+
from dataclasses import dataclass
4+
5+
6+
@dataclass
7+
class kernelInstance:
8+
BLOCK_SIZE: int
9+
# GroupCount: int
10+
MPerBLOCK: int
11+
NPerBLOCK: int
12+
KPerBLOCK: int
13+
WAVE_TILE_M: int
14+
WAVE_TILE_N: int
15+
WAVE_TILE_K: int
16+
WAVE_MAP_M: int
17+
WAVE_MAP_N: int
18+
19+
@property
20+
def name(self) -> str:
21+
return ("_").join(
22+
[
23+
"deepgemm",
24+
("x").join(
25+
map(
26+
lambda x: str(x),
27+
[
28+
self.BLOCK_SIZE,
29+
self.MPerBLOCK,
30+
self.NPerBLOCK,
31+
self.KPerBLOCK,
32+
],
33+
)
34+
),
35+
("x").join(
36+
map(
37+
lambda x: str(x),
38+
[self.WAVE_TILE_M, self.WAVE_TILE_N, self.WAVE_TILE_K],
39+
)
40+
),
41+
("x").join(map(lambda x: str(x), [self.WAVE_MAP_M, self.WAVE_MAP_N])),
42+
]
43+
)
44+
45+
46+
# fmt: off
47+
kernels_list = {
48+
# ( M, N, K): kernel: BLOCK_SIZE| MPerBLOCK| NPerBLOCK| KPerBLOCK| WAVE_TILE_M| WAVE_TILE_N| WAVE_TILE_K| WAVE_MAP_M| WAVE_MAP_N| LOOP_SCHED|PIPELINE_VERSION
49+
1: kernelInstance( 256, 128, 128, 128, 16, 16, 64, 1, 4),
50+
2: kernelInstance( 256, 128, 128, 128, 16, 16, 32, 1, 4),
51+
3: kernelInstance( 256, 32, 64, 256, 16, 16, 64, 1, 4),
52+
4: kernelInstance( 256, 32, 64, 256, 16, 16, 32, 1, 4),
53+
}
54+
55+
56+
default_kernels_dict = {
57+
# ( M, N, K): kernel: BLOCK_SIZE| MPerBLOCK| NPerBLOCK| KPerBLOCK| WAVE_TILE_M| WAVE_TILE_N| WAVE_MAP_M| WAVE_MAP_N| ABLOCK_TRANSFER| BBLOCK_TRANSFER| CBLOCK_TRANSFER| CBLOCK_SPV| CSHUFFLE_MX| CSHUFFLE_NX| LOOP_SCHED|PIPELINE_VERSION
58+
(-1): kernelInstance( 256, 128, 128, 128, 16, 16, 64, 1, 4),
59+
(-2): kernelInstance( 256, 128, 128, 128, 16, 16, 32, 1, 4),
60+
(-3): kernelInstance( 256, 32, 64, 256, 16, 16, 64, 1, 4),
61+
(-4): kernelInstance( 256, 32, 64, 256, 16, 16, 32, 1, 4),
62+
63+
}
64+
# fmt: on

0 commit comments

Comments
 (0)