Commit e351231c4 for llama.cpp
commit e351231c4f4cdd89c88e696c46d0eb718c9e0ab5
Author: bri-prism <288398250+bri-prism@users.noreply.github.com>
Date: Fri Sep 25 02:15:33 2026 -0700
metal: FWHT kernels for block widths above 512 (#29095)
* metal: FWHT kernels for block widths above 512
The Metal FWHT covers widths 64 to 512, one row per simdgroup with N/32 values
per lane. Wider blocks need more registers per lane than that layout allows.
kernel_fwht_tg runs one row per threadgroup with 256 threads, so each thread
keeps N/256 values. Butterflies below the simdgroup width still shuffle, those
up to the threadgroup width go through threadgroup memory, and the rest stay in
registers. Same butterfly and sign convention as the simdgroup kernel.
Widths 64 to 512 keep the simdgroup kernel. 1024 through 8192 use the new one,
for both F32 and F16 sources.
The wide kernels allocate float[N] of threadgroup memory, 32 KB at 8192, so the
size check takes the device limit and reports those widths as unsupported where
they would not fit. Without that a device with less threadgroup memory would
accept the op and then abort on a nil pipeline.
test-backend-ops on M5 Pro: MUL_MAT_HADAMARD 26/26, MUL_MAT 1265/1265.
* cont : add TODOs
---------
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
diff --git a/ggml/src/ggml-metal/ggml-metal-common.cpp b/ggml/src/ggml-metal/ggml-metal-common.cpp
index 9c0b9474c..e43023cc6 100644
--- a/ggml/src/ggml-metal/ggml-metal-common.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-common.cpp
@@ -7,22 +7,29 @@
#include <vector>
-// must stay in sync with the kernel_fwht_<type>_<N> templates in misc.metal
-static bool ggml_metal_fwht_supported_size(int64_t n) {
- return n == 64 || n == 128 || n == 256 || n == 512;
+// must stay in sync with the kernel_fwht_<type>_<N> templates in misc.metal. Widths up to
+// 512 run on the simdgroup kernel and need no threadgroup memory. The wider ones allocate
+// float[N] per threadgroup, so they are only available where that fits.
+static bool ggml_metal_fwht_supported_size(int64_t n, size_t max_tg_mem) {
+ if (n == 64 || n == 128 || n == 256 || n == 512) {
+ return true;
+ }
+
+ if (n == 1024 || n == 2048 || n == 4096 || n == 8192) {
+ return (size_t) n * sizeof(float) <= max_tg_mem;
+ }
+
+ return false;
}
// the FWHT kernels handle a Hadamard-hinted MUL_MAT only under these conditions. supports_op
// and the dispatch must ask the same question: an F16 src1 that is admitted but then falls
// through reaches the generic path, which has no F32 src0 by F16 src1 kernel.
-bool ggml_metal_op_mul_mat_use_fwht(const struct ggml_tensor * op) {
- return ggml_get_op_params_i32(op, 1) == GGML_HINT_SRC0_IS_HADAMARD &&
- op->type == GGML_TYPE_F32 &&
- (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) &&
- ggml_is_contiguous(op->src[1]) &&
- ggml_is_contiguous(op) &&
- ggml_are_same_shape(op->src[1], op) &&
- ggml_metal_fwht_supported_size(op->src[1]->ne[0]);
+bool ggml_metal_op_mul_mat_use_fwht(const struct ggml_tensor * op, size_t max_tg_mem) {
+ return ggml_get_op_params_i32(op, 1) == GGML_HINT_SRC0_IS_HADAMARD && op->type == GGML_TYPE_F32 &&
+ (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) && ggml_is_contiguous(op->src[1]) &&
+ ggml_is_contiguous(op) && ggml_are_same_shape(op->src[1], op) &&
+ ggml_metal_fwht_supported_size(op->src[1]->ne[0], max_tg_mem);
}
bool ggml_metal_op_mul_mat_use_mm(const struct ggml_tensor * op, bool has_simdgroup_mm) {
diff --git a/ggml/src/ggml-metal/ggml-metal-common.h b/ggml/src/ggml-metal/ggml-metal-common.h
index e6a28d032..6b5a1883f 100644
--- a/ggml/src/ggml-metal/ggml-metal-common.h
+++ b/ggml/src/ggml-metal/ggml-metal-common.h
@@ -3,6 +3,7 @@
#pragma once
#include <stdbool.h>
+#include <stddef.h>
#ifdef __cplusplus
extern "C" {
@@ -48,7 +49,7 @@ bool ggml_mem_ranges_check(ggml_mem_ranges_t mrs, const struct ggml_tensor * ten
void ggml_graph_optimize(struct ggml_cgraph * gf);
// mat-mat vs mat-vec dispatch; used by both supports_op and ggml_metal_op_mul_mat*
-bool ggml_metal_op_mul_mat_use_fwht (const struct ggml_tensor * op);
+bool ggml_metal_op_mul_mat_use_fwht(const struct ggml_tensor * op, size_t max_tg_mem);
bool ggml_metal_op_mul_mat_use_mm (const struct ggml_tensor * op, bool has_simdgroup_mm);
bool ggml_metal_op_mul_mat_id_use_mm(const struct ggml_tensor * op, bool has_simdgroup_mm);
diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m
index c4b86e80d..fa58b8965 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.m
+++ b/ggml/src/ggml-metal/ggml-metal-device.m
@@ -1854,7 +1854,7 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
// the FWHT kernels read an F16 source directly; every other F16 src1 path
// still goes through ggml_metal_supports_mul_mat_op
if (op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F16 &&
- ggml_metal_op_mul_mat_use_fwht(op)) {
+ ggml_metal_op_mul_mat_use_fwht(op, dev->props.max_theadgroup_memory_size)) {
return has_simdgroup_reduction;
}
return ggml_metal_supports_mul_mat_op(
diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h
index 490dd83a1..eed85f283 100644
--- a/ggml/src/ggml-metal/ggml-metal-impl.h
+++ b/ggml/src/ggml-metal/ggml-metal-impl.h
@@ -1235,6 +1235,11 @@ typedef struct {
int32_t top_k; // k
} ggml_metal_kargs_top_k;
+// widths at or above this use the threadgroup FWHT kernel, one row per threadgroup
+// with GGML_METAL_FWHT_TG_NT threads, instead of one row per simdgroup
+#define GGML_METAL_FWHT_TG_MIN_N 1024
+#define GGML_METAL_FWHT_TG_NT 256
+
typedef struct {
int32_t ne01; // n_tokens
uint64_t nb01; // logits row stride
diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp
index 54d5b9808..8a46ec66a 100644
--- a/ggml/src/ggml-metal/ggml-metal-ops.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp
@@ -2343,6 +2343,13 @@ int ggml_metal_op_fwht(ggml_metal_op_t ctx, int idx) {
const int th_max = ggml_metal_pipeline_max_theads_per_threadgroup(pipeline);
const int simd_size = 32;
+ if (n >= GGML_METAL_FWHT_TG_MIN_N) {
+ GGML_ASSERT(th_max >= GGML_METAL_FWHT_TG_NT);
+ ggml_metal_encoder_dispatch_threadgroups(enc, nrows, 1, 1, GGML_METAL_FWHT_TG_NT, 1, 1);
+
+ return 1;
+ }
+
int sg_per_tg = 2;
sg_per_tg = std::min(sg_per_tg, th_max/simd_size);
sg_per_tg = std::max(sg_per_tg, 1);
@@ -2419,10 +2426,11 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) {
ggml_metal_library_t lib = ctx->lib;
ggml_metal_encoder_t enc = ctx->enc;
- if (ggml_metal_op_mul_mat_use_fwht(op)) {
+ const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev);
+
+ if (ggml_metal_op_mul_mat_use_fwht(op, props_dev->max_theadgroup_memory_size)) {
return ggml_metal_op_fwht(ctx, idx);
}
- const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev);
GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
diff --git a/ggml/src/ggml-metal/kernels/misc.metal b/ggml/src/ggml-metal/kernels/misc.metal
index 279d69f8f..d3b01978f 100644
--- a/ggml/src/ggml-metal/kernels/misc.metal
+++ b/ggml/src/ggml-metal/kernels/misc.metal
@@ -429,6 +429,81 @@ kernel void kernel_fwht(
}
}
+// Wide blocks: one row per threadgroup instead of per simdgroup, so each thread keeps
+// N/NT values rather than N/32. Butterflies below the simdgroup width still shuffle;
+// those up to NT go through threadgroup memory; the rest stay in registers.
+// TODO: try avoiding branch https://github.com/ggml-org/llama.cpp/pull/29094#discussion_r4049563223
+// TODO: try to unroll loops
+template<int N, int NT, typename src_t>
+kernel void kernel_fwht_tg(
+ constant ggml_metal_kargs_fwht & args,
+ device const src_t * src,
+ device float * dst,
+ uint3 tgpig[[threadgroup_position_in_grid]],
+ ushort sgitg[[simdgroup_index_in_threadgroup]],
+ ushort tiisg[[thread_index_in_simdgroup]],
+ ushort3 ntg[[threads_per_threadgroup]]) {
+
+ constexpr int NW = N_SIMDWIDTH;
+ constexpr int NE = N / NT;
+
+ threadgroup float shmem[N];
+
+ const float scale = 1.0f / sqrt((float) N);
+
+ const int64_t r = tgpig.x;
+ if (r >= args.nrows) {
+ return;
+ }
+
+ src += r * N;
+ dst += r * N;
+
+ const int tid = sgitg * NW + tiisg;
+
+ float reg[NE];
+ for (int i = 0; i < NE; i++) {
+ reg[i] = float(src[i*NT + tid])*scale;
+ }
+
+ for (int i = 1; i < NW; i *= 2) {
+ for (int j = 0; j < NE; j++) {
+ const float val = reg[j];
+ const float val2 = simd_shuffle_xor(val, i);
+ reg[j] = (tid & i) == 0 ? val2 + val : val2 - val;
+ }
+ }
+
+ for (int i = NW; i < NT; i *= 2) {
+ for (int j = 0; j < NE; j++) {
+ shmem[j*NT + tid] = reg[j];
+ }
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+ for (int j = 0; j < NE; j++) {
+ const float val = reg[j];
+ const float val2 = shmem[j*NT + (tid ^ i)];
+ reg[j] = (tid & i) == 0 ? val2 + val : val2 - val;
+ }
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+ }
+
+ for (int i = NT; i < N; i *= 2) {
+ const int step = i / NT;
+ for (int j = 0; j < NE; j += (2 * step)) {
+ for (int k = 0; k < step; k++) {
+ const float x = reg[j + k ];
+ const float y = reg[j + k + step];
+ reg[j + k] = x + y;
+ reg[j + k + step] = x - y;
+ }
+ }
+ }
+
+ for (int i = 0; i < NE; i++) {
+ dst[i*NT + tid] = reg[i];
+ }
+}
+
typedef decltype(kernel_fwht<64, float>) kernel_fwht_f32_t;
typedef decltype(kernel_fwht<64, half>) kernel_fwht_f16_t;
@@ -442,6 +517,16 @@ template [[host_name("kernel_fwht_f16_128")]] kernel kernel_fwht_f16_t kernel_fw
template [[host_name("kernel_fwht_f16_256")]] kernel kernel_fwht_f16_t kernel_fwht<256, half>;
template [[host_name("kernel_fwht_f16_512")]] kernel kernel_fwht_f16_t kernel_fwht<512, half>;
+template [[host_name("kernel_fwht_f32_1024")]] kernel kernel_fwht_f32_t kernel_fwht_tg<1024, GGML_METAL_FWHT_TG_NT, float>;
+template [[host_name("kernel_fwht_f32_2048")]] kernel kernel_fwht_f32_t kernel_fwht_tg<2048, GGML_METAL_FWHT_TG_NT, float>;
+template [[host_name("kernel_fwht_f32_4096")]] kernel kernel_fwht_f32_t kernel_fwht_tg<4096, GGML_METAL_FWHT_TG_NT, float>;
+template [[host_name("kernel_fwht_f32_8192")]] kernel kernel_fwht_f32_t kernel_fwht_tg<8192, GGML_METAL_FWHT_TG_NT, float>;
+
+template [[host_name("kernel_fwht_f16_1024")]] kernel kernel_fwht_f16_t kernel_fwht_tg<1024, GGML_METAL_FWHT_TG_NT, half>;
+template [[host_name("kernel_fwht_f16_2048")]] kernel kernel_fwht_f16_t kernel_fwht_tg<2048, GGML_METAL_FWHT_TG_NT, half>;
+template [[host_name("kernel_fwht_f16_4096")]] kernel kernel_fwht_f16_t kernel_fwht_tg<4096, GGML_METAL_FWHT_TG_NT, half>;
+template [[host_name("kernel_fwht_f16_8192")]] kernel kernel_fwht_f16_t kernel_fwht_tg<8192, GGML_METAL_FWHT_TG_NT, half>;
+
constant int FC_dsv4_hc_n_hc [[function_constant(FC_DSV4_HC + 0)]];
kernel void kernel_dsv4_hc_comb_f32(
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 89c18a2d5..6527049ff 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -10018,7 +10018,8 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 128, 4, 128, {2, 3}));
test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 256, 512, 256)); // many rows
test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 32, 1, 32)); // too small (N<64)
- test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 1024, 1, 1024)); // too big (N>512)
+ test_cases.emplace_back(
+ new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 16384, 1, 16384)); // too big (N>8192)
test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 64, 1, 64));
test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 128, 1, 128));
test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 256, 1, 256));
@@ -10026,6 +10027,16 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 128, 32, 128));
test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 128, 4, 128, {2, 3}));
test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 256, 512, 256)); // many rows
+ test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 1024, 1, 1024));
+ test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 2048, 1, 2048));
+ test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 4096, 1, 4096));
+ test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 8192, 1, 8192));
+ test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 1024, 7, 1024)); // many rows
+ test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 1024, 1, 1024));
+ test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 2048, 1, 2048));
+ test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 4096, 1, 4096));
+ test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 8192, 1, 8192));
+ test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 1024, 7, 1024)); // many rows
// FP4 activation precision (default = native W4A4, src1 GGML_PREC_Q8 = W4A8)
test_cases.emplace_back(new test_mul_mat_w4a8(GGML_TYPE_NVFP4, GGML_TYPE_F32, 32, 1, 256));