Skip to content

Commit 026366c

Browse files
cthifacebook-github-bot
cthi
authored and
facebook-github-bot
committed
Support tuning cache for Cutlass FP8 GEMM (#4301)
Summary: Pull Request resolved: #4301 X-link: facebookresearch/FBGEMM#1377 This diff adds support for the tuning cache to the kernel. There should be no performance changes to the existing heuristics. - I refactored the kernel dispatch logic to instead return the kernel function, as it removes some duplication of the kernel invoke. - The next diff in this stack will add the new kernels D75820688, to make the review easier - Note that we are having some issues with adding the new kernels, as I have found this kernel is actually compiling 12 variants for each configuration, see D75820688 for more context. So for now we won't add the new kernels in D75820688, but we can just onboard it to auto tuning incase someone wants to compile them locally. Will revisit D75820688 later. Reviewed By: q10, jiawenliu64 Differential Revision: D75541025
1 parent 0a56667 commit 026366c

File tree

2 files changed

+147
-92
lines changed

2 files changed

+147
-92
lines changed

‎fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8f8bf16_rowwise.cu‎

Lines changed: 124 additions & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -12,164 +12,196 @@
1212
// clang-format on
1313

1414
#include "f8f8bf16_rowwise/f8f8bf16_rowwise_manifest.cuh"
15+
#include "fbgemm_gpu/quantize/tuning_cache.hpp"
16+
#include "fbgemm_gpu/quantize/utils.h"
1517

1618
namespace fbgemm_gpu {
1719

1820
#if CUDART_VERSION >= 12000
1921

2022
// FP8 Rowwise Cutlass kernel dispatch.
21-
at::Tensor dispatch_fp8_rowwise_kernel(
22-
at::Tensor XQ,
23-
at::Tensor WQ,
24-
at::Tensor x_scale,
25-
at::Tensor w_scale,
26-
bool use_fast_accum,
27-
std::optional<at::Tensor> bias = std::nullopt,
28-
std::optional<at::Tensor> output = std::nullopt) {
29-
int M = size_to_dim_(XQ.dim() - 1, XQ.sizes());
30-
int N = size_to_dim_(WQ.dim() - 1, WQ.sizes());
31-
int K = XQ.size(-1);
32-
static int arch = -1;
33-
// Avoid expensive cudaGetDeviceProperties call.
34-
if (arch < 0) {
35-
cudaDeviceProp prop;
36-
cudaGetDeviceProperties(&prop, 0);
37-
if (prop.major >= 10) {
38-
arch = 10;
39-
int runtimeVersion;
40-
C10_CUDA_CHECK(cudaRuntimeGetVersion(&runtimeVersion));
41-
TORCH_CHECK(
42-
runtimeVersion >= 12080,
43-
"FP8 GEMM on sm100a or above requires cuda >= 12.8");
44-
} else {
45-
arch = 9;
46-
}
47-
}
48-
23+
Kernel_f8f8bf16_rowwise
24+
get_kernel_via_heuristic(int arch, int M, int N, int K, bool use_fast_accum) {
4925
// Use shape heuristics to dispatch to optimized kernel configuration.
5026
if (arch == 10) {
5127
if (M <= 128) {
5228
if (N <= 1024) {
53-
return f8f8bf16_rowwise_128_32_128_1_1_1_10_f_f(
54-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
29+
return f8f8bf16_rowwise_128_32_128_1_1_1_10_f_f;
5530
} else {
56-
return f8f8bf16_rowwise_128_64_128_1_1_1_10_f_f(
57-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
31+
return f8f8bf16_rowwise_128_64_128_1_1_1_10_f_f;
5832
}
5933
} else if (M <= 1024) {
6034
if (N <= 1024) {
61-
return f8f8bf16_rowwise_128_256_128_2_1_1_10_f_f(
62-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
35+
return f8f8bf16_rowwise_128_256_128_2_1_1_10_f_f;
6336
} else {
64-
return f8f8bf16_rowwise_128_128_128_2_2_1_10_f_f(
65-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
37+
return f8f8bf16_rowwise_128_128_128_2_2_1_10_f_f;
6638
}
6739
} else if (M <= 2048) {
68-
return f8f8bf16_rowwise_128_256_128_2_1_1_10_f_f(
69-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
40+
return f8f8bf16_rowwise_128_256_128_2_1_1_10_f_f;
7041
} else {
7142
if (N <= 1024) {
72-
return f8f8bf16_rowwise_128_256_128_1_2_1_10_f_f(
73-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
43+
return f8f8bf16_rowwise_128_256_128_1_2_1_10_f_f;
7444
} else {
75-
return f8f8bf16_rowwise_128_256_128_2_1_1_10_f_f(
76-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
45+
return f8f8bf16_rowwise_128_256_128_2_1_1_10_f_f;
7746
}
7847
}
7948
} else {
8049
if (M <= 16) {
81-
return f8f8bf16_rowwise_64_16_128_1_1_1_9_f_f(
82-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
50+
return f8f8bf16_rowwise_64_16_128_1_1_1_9_f_f;
8351
} else if (M <= 32) {
8452
if (N <= 4096) {
85-
return f8f8bf16_rowwise_64_16_128_1_1_1_9_f_f(
86-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
53+
return f8f8bf16_rowwise_64_16_128_1_1_1_9_f_f;
8754
} else {
88-
return f8f8bf16_rowwise_64_32_128_2_1_1_9_f_f(
89-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
55+
return f8f8bf16_rowwise_64_32_128_2_1_1_9_f_f;
9056
}
9157
} else if (M <= 64) {
9258
if (N <= 2048) {
93-
return f8f8bf16_rowwise_64_16_128_1_1_1_9_f_f(
94-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
59+
return f8f8bf16_rowwise_64_16_128_1_1_1_9_f_f;
9560
} else if (N <= 4096) {
96-
return f8f8bf16_rowwise_64_32_128_2_1_1_9_f_f(
97-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
61+
return f8f8bf16_rowwise_64_32_128_2_1_1_9_f_f;
9862
} else {
99-
return f8f8bf16_rowwise_64_64_128_2_1_1_9_f_f(
100-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
63+
return f8f8bf16_rowwise_64_64_128_2_1_1_9_f_f;
10164
}
10265
} else if (M <= 128) {
10366
if (N <= 1024) {
104-
return f8f8bf16_rowwise_64_16_128_1_1_1_9_f_f(
105-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
67+
return f8f8bf16_rowwise_64_16_128_1_1_1_9_f_f;
10668
} else if (N <= 2048) {
107-
return f8f8bf16_rowwise_64_32_128_2_1_1_9_f_f(
108-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
69+
return f8f8bf16_rowwise_64_32_128_2_1_1_9_f_f;
10970
} else if (N <= 4096) {
110-
return f8f8bf16_rowwise_64_64_128_2_1_1_9_f_f(
111-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
71+
return f8f8bf16_rowwise_64_64_128_2_1_1_9_f_f;
11272
} else {
113-
return f8f8bf16_rowwise_64_128_128_1_1_1_9_f_f(
114-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
73+
return f8f8bf16_rowwise_64_128_128_1_1_1_9_f_f;
11574
}
11675
} else if (M <= 256) {
11776
if (N <= 1024) {
118-
return f8f8bf16_rowwise_64_32_128_2_1_1_9_f_f(
119-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
77+
return f8f8bf16_rowwise_64_32_128_2_1_1_9_f_f;
12078
} else if (N <= 2048) {
121-
return f8f8bf16_rowwise_64_64_128_2_1_1_9_f_f(
122-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
79+
return f8f8bf16_rowwise_64_64_128_2_1_1_9_f_f;
12380
} else if (N <= 4096) {
124-
return f8f8bf16_rowwise_64_128_128_1_1_1_9_f_f(
125-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
81+
return f8f8bf16_rowwise_64_128_128_1_1_1_9_f_f;
12682
} else {
127-
return f8f8bf16_rowwise_64_256_128_1_1_1_9_f_f(
128-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
83+
return f8f8bf16_rowwise_64_256_128_1_1_1_9_f_f;
12984
}
13085
} else if (M <= 512) {
13186
if (N <= 1024) {
132-
return f8f8bf16_rowwise_64_64_128_2_1_1_9_f_f(
133-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
87+
return f8f8bf16_rowwise_64_64_128_2_1_1_9_f_f;
13488
} else if (N <= 2048) {
135-
return f8f8bf16_rowwise_64_128_128_1_1_1_9_f_f(
136-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
89+
return f8f8bf16_rowwise_64_128_128_1_1_1_9_f_f;
13790
} else if (N <= 4096 || use_fast_accum == false) {
138-
return f8f8bf16_rowwise_64_256_128_1_1_1_9_f_f(
139-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
91+
return f8f8bf16_rowwise_64_256_128_1_1_1_9_f_f;
14092
} else {
141-
return f8f8bf16_rowwise_128_256_128_2_1_1_9_f_t(
142-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
93+
return f8f8bf16_rowwise_128_256_128_2_1_1_9_f_t;
14394
}
14495
} else if (M <= 1024) {
14596
if (N <= 1024) {
146-
return f8f8bf16_rowwise_64_128_128_1_1_1_9_f_f(
147-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
97+
return f8f8bf16_rowwise_64_128_128_1_1_1_9_f_f;
14898
} else if (N <= 2048 || use_fast_accum == false) {
149-
return f8f8bf16_rowwise_64_256_128_1_1_1_9_f_f(
150-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
99+
return f8f8bf16_rowwise_64_256_128_1_1_1_9_f_f;
151100
} else {
152-
return f8f8bf16_rowwise_128_256_128_2_1_1_9_f_t(
153-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
101+
return f8f8bf16_rowwise_128_256_128_2_1_1_9_f_t;
154102
}
155103
} else {
156104
if (M <= 2048 && N <= 1024) {
157-
return f8f8bf16_rowwise_64_256_128_2_1_1_9_f_f(
158-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
105+
return f8f8bf16_rowwise_64_256_128_2_1_1_9_f_f;
159106
} else if (K <= 4096 || use_fast_accum == false) {
160-
return f8f8bf16_rowwise_128_128_128_2_1_1_9_t_f(
161-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
107+
return f8f8bf16_rowwise_128_128_128_2_1_1_9_t_f;
162108
} else if (M > 8192 && N > 8192) {
163-
return f8f8bf16_rowwise_128_256_128_4_4_1_9_f_t(
164-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
109+
return f8f8bf16_rowwise_128_256_128_4_4_1_9_f_t;
165110
} else {
166-
return f8f8bf16_rowwise_128_256_128_2_1_1_9_f_t(
167-
XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
111+
return f8f8bf16_rowwise_128_256_128_2_1_1_9_f_t;
168112
}
169113
}
170114
}
171115
}
172116

117+
Kernel_f8f8bf16_rowwise get_kernel_via_tuning(
118+
int arch,
119+
int M,
120+
int N,
121+
int K,
122+
at::Tensor XQ,
123+
at::Tensor WQ,
124+
at::Tensor x_scale,
125+
at::Tensor w_scale,
126+
bool use_fast_accum,
127+
std::optional<at::Tensor> bias = std::nullopt,
128+
std::optional<at::Tensor> output = std::nullopt) {
129+
// One cache per kernel type
130+
static TuningCache cache("f8f8bf16_rowwise");
131+
132+
// Reducing amount of auto tuning by rounding up M to next power of 2.
133+
M = nextPowerOf2(M);
134+
// Use (M, N, K) shape as the key.
135+
const std::string shape_key =
136+
std::to_string(M) + "_" + std::to_string(N) + "_" + std::to_string(K);
137+
const auto& kernels = get_f8f8bf16_rowwise_kernels(arch);
138+
auto kernel = cache.findBestKernelMaybeAutotune(
139+
shape_key,
140+
kernels,
141+
XQ,
142+
WQ,
143+
x_scale,
144+
w_scale,
145+
use_fast_accum,
146+
bias,
147+
output);
148+
149+
return kernel;
150+
}
151+
152+
// FP8 Rowwise Cutlass kernel dispatch.
153+
at::Tensor dispatch_fp8_rowwise_kernel(
154+
at::Tensor XQ,
155+
at::Tensor WQ,
156+
at::Tensor x_scale,
157+
at::Tensor w_scale,
158+
bool use_fast_accum,
159+
std::optional<at::Tensor> bias = std::nullopt,
160+
std::optional<at::Tensor> output = std::nullopt) {
161+
int M = size_to_dim_(XQ.dim() - 1, XQ.sizes());
162+
int N = size_to_dim_(WQ.dim() - 1, WQ.sizes());
163+
int K = XQ.size(-1);
164+
165+
static int arch = -1;
166+
// Avoid expensive cudaGetDeviceProperties call.
167+
if (arch < 0) {
168+
cudaDeviceProp prop;
169+
cudaGetDeviceProperties(&prop, 0);
170+
if (prop.major >= 10) {
171+
arch = 10;
172+
int runtimeVersion;
173+
C10_CUDA_CHECK(cudaRuntimeGetVersion(&runtimeVersion));
174+
TORCH_CHECK(
175+
runtimeVersion >= 12080,
176+
"FP8 GEMM on sm100a or above requires cuda >= 12.8");
177+
} else {
178+
arch = 9;
179+
}
180+
}
181+
182+
// Select kernel to run via heuristics or tuning.
183+
auto kernel = [&]() {
184+
if (std::getenv("FBGEMM_AUTOTUNE_ENABLE")) {
185+
return get_kernel_via_tuning(
186+
arch,
187+
M,
188+
N,
189+
K,
190+
XQ,
191+
WQ,
192+
x_scale,
193+
w_scale,
194+
use_fast_accum,
195+
bias,
196+
output);
197+
} else {
198+
return get_kernel_via_heuristic(arch, M, N, K, use_fast_accum);
199+
}
200+
}();
201+
// Invoke kernel
202+
return kernel(XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
203+
}
204+
173205
void f8f8bf16_rowwise_out(
174206
at::Tensor XQ, // FP8
175207
at::Tensor WQ, // FP8

‎fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8f8bf16_rowwise/f8f8bf16_rowwise_manifest.cuh‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,4 +135,27 @@ at::Tensor f8f8bf16_rowwise_128_256_128_2_1_1_10_f_f(
135135
bool use_fast_accum = true,
136136
std::optional<at::Tensor> bias = std::nullopt,
137137
std::optional<at::Tensor> output = std::nullopt);
138+
139+
using Kernel_f8f8bf16_rowwise = at::Tensor (*)(
140+
at::Tensor,
141+
at::Tensor,
142+
at::Tensor,
143+
at::Tensor,
144+
bool,
145+
std::optional<at::Tensor>,
146+
std::optional<at::Tensor>);
147+
148+
inline const std::unordered_map<std::string, Kernel_f8f8bf16_rowwise>&
149+
get_f8f8bf16_rowwise_kernels(int arch) {
150+
static const std::unordered_map<std::string, Kernel_f8f8bf16_rowwise>
151+
kernelsSM90 = {};
152+
static const std::unordered_map<std::string, Kernel_f8f8bf16_rowwise>
153+
kernelsSM100 = {};
154+
if (arch == 10) {
155+
return kernelsSM100;
156+
} else {
157+
return kernelsSM90;
158+
}
159+
}
160+
138161
} // namespace fbgemm_gpu

0 commit comments

Comments
 (0)