From 2e433c9254478f0de97a9ea10e800c2880125acf Mon Sep 17 00:00:00 2001 From: Sergei Kulik Date: Sat, 22 Aug 2026 10:48:10 +0930 Subject: [PATCH 1/2] Thread swizzling in kernel_mul_mm (Metal) for better cache locality. Avoids throughput collapse after hitting the SLC limit. --- ggml/src/ggml-metal/ggml-metal.metal | 57 ++++++++++++++++++++++++++-- 1 file changed, 53 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index 27f97b5e0798..d4446a1e22ce 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -10470,8 +10470,33 @@ kernel void kernel_mul_mm( constexpr int NL1 = NK/8; const int im = tgpig.z; - const int r0 = tgpig.y*NR0; - const int r1 = tgpig.x*NR1; + + int sy = tgpig.y; + int sx = tgpig.x; + + // Threadgroup swizzle. Walking the grid in groups of SWZ row tiles, cutting src1 + // DRAM traffic by ~SWZ. Only worth doing once src1 is large. + if ((size_t) args.ne1 * args.nb11 > (32u << 20)){ + const uint tile_bytes = NR0*(uint) args.nb01; + const uint want = clamp((10u << 20)/max(tile_bytes, 1u), 1u, 8u); + + const int SWZ = want >= 8 ? 8 : want >= 4 ? 4 : want >= 2 ? 2 : 1; + + const int nbx = (args.ne1 + NR1 - 1)/NR1; // src1 (batch) tiles + const int nby = (args.ne0 + NR0 - 1)/NR0; // src0 (row) tiles + + const int lin = (int) tgpig.y*nbx + (int) tgpig.x; + const int tpg = SWZ*nbx; + const int y0 = (lin/tpg)*SWZ; + const int lid = lin%tpg; + const int gh = (nby - y0) < SWZ ? (nby - y0) : SWZ; + + sy = y0 + lid%gh; + sx = lid/gh; + } + + const int r0 = sy*NR0; + const int r1 = sx*NR1; // if this block is of 64x32 shape or smaller const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0; @@ -10755,14 +10780,38 @@ kernel void kernel_mul_mm_id( constexpr int NL1 = NK/8; const int im = tgpig.z; // expert - const int r0 = tgpig.y*NR0; - const int r1 = tgpig.x*NR1; device const uint32_t * tpe_u32 = (device const uint32_t *) (htpe); device const int32_t * ids_i32 = (device const int32_t *) (hids); const int32_t neh1 = tpe_u32[im]; + int sy = tgpig.y; + int sx = tgpig.x; + + // Same threadgroup swizzle as kernel_mul_mm + if ((size_t) neh1 * args.nb11 > (32u << 20)) { + const uint tile_bytes = NR0*(uint) args.nb01; + const uint want = clamp((10u << 20)/max(tile_bytes, 1u), 1u, 8u); + + const int SWZ = want >= 8 ? 8 : want >= 4 ? 4 : want >= 2 ? 2 : 1; + + const int nbx = (args.ne21 + NR1 - 1)/NR1; + const int nby = (args.ne0 + NR0 - 1)/NR0; + + const int lin = (int) tgpig.y*nbx + (int) tgpig.x; + const int tpg = SWZ*nbx; + const int y0 = (lin/tpg)*SWZ; + const int lid = lin%tpg; + const int gh = (nby - y0) < SWZ ? (nby - y0) : SWZ; + + sy = y0 + lid%gh; + sx = lid/gh; + } + + const int r0 = sy*NR0; + const int r1 = sx*NR1; + if (r1 >= neh1) { return; } From 77c56f9a82a81164249d86d746228680076ff191 Mon Sep 17 00:00:00 2001 From: Sergei Kulik Date: Sat, 22 Aug 2026 13:20:42 +0930 Subject: [PATCH 2/2] Tests with large input to cover the new swizzling path. --- tests/test-backend-ops.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 8e3b273a1e46..ae4c319b2158 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9299,6 +9299,11 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 6, 4096, 5120, {1, 1}, {1, 1})); + // Threads swizzle path in Metal + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 640, 1056, 8192, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 320, 512, 17408, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q8_0, GGML_TYPE_F32, 2, 2, false, 128, 1056, 8192)); + // K not a multiple of 32 test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 64, 32, 65, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 64, 32, 80, {1, 1}, {1, 1}));