[ROCm][Perf] Tune wvSplitK on gfx1151 (#40784)
Signed-off-by: Matthias Gehre <matthias.gehre@amd.com>
Matthias Gehre committed
Jun 25, 2026 at 08:17 UTC
77c1d9fe9b2b4f49fe8d98d5da5d9efc1134b9cd
1 file changed
+53
-21
csrc/rocm/skinny_gemms.cu
+53
-21
@@ -70,6 +70,15 @@ bool on_gfx12() {
70
return result;
71
}
72
73
+bool on_gfx1151() {
74
+ static const bool result = [] {
75
+ const auto* dprops = at::cuda::getCurrentDeviceProperties();
76
+ const std::string device_arch = dprops->gcnArchName;
77
+ return device_arch.find("gfx1151") != std::string::npos;
78
+ }();
79
+ return result;
80
+}
81
+
82
#if defined(NDEBUG)
83
#undef NDEBUG
84
#include <assert.h>
@@ -1237,6 +1246,45 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
1246
WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 2, __N) \
1247
}
1248
1249
+// WVSPLITK_CFG arguments are: (THRDS, WVPRGRP, YTILE, UNRL, N).
1250
+// THRDS = wavefront width (32 on GFX11/GFX12, 64 on GFX9)
1251
+// WVPRGRP= waves per group (always 16)
1252
+// YTILE = output rows per thread tile
1253
+// UNRL = K-loop unroll factor
1254
+// N = batch size (passed through from the switch in wvSplitK)
1255
+#define WVSPLIT_TILE(_sYT, __N) \
1256
+ { \
1257
+ if (on_gfx1151()) { \
1258
+ bool fit_lds = (Kbp_in * N_in <= max_lds_len); \
1259
+ if (_sYT <= 1) \
1260
+ WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/4, \
1261
+ __N) \
1262
+ else if ((K_in % 1024 == 512) && K_in >= 1536 && \
1263
+ (_sYT >= 40 || K_in >= 4096)) \
1264
+ WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/4, /*UNRL=*/1, \
1265
+ __N) \
1266
+ else if (K_in < 1024) \
1267
+ WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/2, /*UNRL=*/4, \
1268
+ __N) \
1269
+ else if (K_in <= 2048 && (__N >= 2 || _sYT <= 26)) \
1270
+ WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/4, \
1271
+ __N) \
1272
+ else if (__N >= 2 && !fit_lds) \
1273
+ WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/4, \
1274
+ __N) \
1275
+ else if (__N == 1) \
1276
+ WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/2, \
1277
+ __N) \
1278
+ else \
1279
+ WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/1, \
1280
+ __N) \
1281
+ } else if (on_gfx1x()) { /* gfx1100/gfx1150/GFX12, wave32 */ \
1282
+ WVSPLIT_TILE_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, _sYT, __N) \
1283
+ } else { /* GFX9, wave64 */ \
1284
+ WVSPLIT_TILE_CFG(/*THRDS=*/64, /*WVPRGRP=*/16, _sYT, __N) \
1285
+ } \
1286
+ }
1287
+
1288
AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "wvSplitK", [&] {
1289
using fptype = typename scalar<scalar_t>::type;
1290
fptype* af4 = reinterpret_cast<fptype*>(in_a.data_ptr());
@@ -1251,37 +1299,21 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
1299
// then cut the active waves to balance their distribution...
1300
int sYT = (M_in + CuCount * 4 - 1) / (CuCount * 4);
1301
1254
- const bool use_wave32 = on_gfx1x();
1302
switch (N_in) {
1303
case 1:
1257
- if (use_wave32)
1258
- WVSPLIT_TILE_CFG(32, 16, sYT, 1)
1259
- else
1260
- WVSPLIT_TILE_CFG(64, 16, sYT, 1)
1304
+ WVSPLIT_TILE(sYT, 1)
1305
break;
1306
case 2:
1263
- if (use_wave32)
1264
- WVSPLIT_TILE_CFG(32, 16, sYT, 2)
1265
- else
1266
- WVSPLIT_TILE_CFG(64, 16, sYT, 2)
1307
+ WVSPLIT_TILE(sYT, 2)
1308
break;
1309
case 3:
1269
- if (use_wave32)
1270
- WVSPLIT_TILE_CFG(32, 16, sYT, 3)
1271
- else
1272
- WVSPLIT_TILE_CFG(64, 16, sYT, 3)
1310
+ WVSPLIT_TILE(sYT, 3)
1311
break;
1312
case 4:
1275
- if (use_wave32)
1276
- WVSPLIT_TILE_CFG(32, 16, sYT, 4)
1277
- else
1278
- WVSPLIT_TILE_CFG(64, 16, sYT, 4)
1313
+ WVSPLIT_TILE(sYT, 4)
1314
break;
1315
case 5:
1281
- if (use_wave32)
1282
- WVSPLIT_TILE_CFG(32, 16, sYT, 5)
1283
- else
1284
- WVSPLIT_TILE_CFG(64, 16, sYT, 5)
1316
+ WVSPLIT_TILE(sYT, 5)
1317
break;
1318
default:
1319
throw std::runtime_error(